diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 884e0ffc..9da5b576 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -4,8 +4,10 @@ ### Added +- Added `--session-id` to let CLI callers use an exact project-local session ID, creating it if missing ([#4874](https://github.com/earendil-works/pi/issues/4874)). - Added `excludeFromContext` flag to the `bash` RPC command for parity with the internal `executeBash` API ([#5039](https://github.com/earendil-works/pi/issues/5039)). + ### Fixed - Fixed user message transcript rendering to preserve user-authored ordered-list markers ([#5013](https://github.com/earendil-works/pi/issues/5013)). diff --git a/packages/coding-agent/src/cli/args.ts b/packages/coding-agent/src/cli/args.ts index e4cb27dd..19053db5 100644 --- a/packages/coding-agent/src/cli/args.ts +++ b/packages/coding-agent/src/cli/args.ts @@ -23,6 +23,7 @@ export interface Args { mode?: Mode; noSession?: boolean; session?: string; + sessionId?: string; fork?: string; sessionDir?: string; models?: string[]; @@ -95,6 +96,8 @@ export function parseArgs(args: string[]): Args { result.noSession = true; } else if (arg === "--session" && i + 1 < args.length) { result.session = args[++i]; + } else if (arg === "--session-id" && i + 1 < args.length) { + result.sessionId = args[++i]; } else if (arg === "--fork" && i + 1 < args.length) { result.fork = args[++i]; } else if (arg === "--session-dir" && i + 1 < args.length) { @@ -224,6 +227,7 @@ ${chalk.bold("Options:")} --continue, -c Continue previous session --resume, -r Select a session to resume --session Use specific session file or partial UUID + --session-id Use exact project session ID, creating it if missing --fork Fork specific session file or partial UUID into a new session --session-dir Directory for session storage and lookup --no-session Don't save session (ephemeral) diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index e7d12314..22c1f436 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -202,6 +202,14 @@ function createSessionId(): string { return uuidv7(); } +export function assertValidSessionId(id: string): void { + if (!/^[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9])?$/.test(id)) { + throw new Error( + "Session id must be non-empty, contain only alphanumeric characters, '-', '_', and '.', and start and end with an alphanumeric character", + ); + } +} + /** Generate a unique short ID (8 hex chars, collision-checked) */ function generateId(byId: { has(id: string): boolean }): string { for (let i = 0; i < 100; i++) { @@ -721,7 +729,13 @@ export class SessionManager { private labelTimestampsById: Map = new Map(); private leafId: string | null = null; - private constructor(cwd: string, sessionDir: string, sessionFile: string | undefined, persist: boolean) { + private constructor( + cwd: string, + sessionDir: string, + sessionFile: string | undefined, + persist: boolean, + newSessionOptions?: NewSessionOptions, + ) { this.cwd = resolvePath(cwd); this.sessionDir = normalizePath(sessionDir); this.persist = persist; @@ -732,7 +746,7 @@ export class SessionManager { if (sessionFile) { this.setSessionFile(sessionFile); } else { - this.newSession(); + this.newSession(newSessionOptions); } } @@ -770,6 +784,9 @@ export class SessionManager { } newSession(options?: NewSessionOptions): string | undefined { + if (options?.id !== undefined) { + assertValidSessionId(options.id); + } this.sessionId = options?.id ?? createSessionId(); const timestamp = new Date().toISOString(); const header: SessionHeader = { @@ -845,14 +862,23 @@ export class SessionManager { const hasAssistant = this.fileEntries.some((e) => e.type === "message" && e.message.role === "assistant"); if (!hasAssistant) { - // Mark as not flushed so when assistant arrives, all entries get written - this.flushed = false; + if (this.flushed) { + appendFileSync(this.sessionFile, `${JSON.stringify(entry)}\n`); + } else { + // Mark as not flushed so when assistant arrives, all entries get written + this.flushed = false; + } return; } if (!this.flushed) { - for (const e of this.fileEntries) { - appendFileSync(this.sessionFile, `${JSON.stringify(e)}\n`); + const fd = openSync(this.sessionFile, "wx"); + try { + for (const e of this.fileEntries) { + writeFileSync(fd, `${JSON.stringify(e)}\n`); + } + } finally { + closeSync(fd); } this.flushed = true; } else { @@ -1308,9 +1334,9 @@ export class SessionManager { * @param cwd Working directory (stored in session header) * @param sessionDir Optional session directory. If omitted, uses default (~/.pi/agent/sessions//). */ - static create(cwd: string, sessionDir?: string): SessionManager { + static create(cwd: string, sessionDir?: string, options?: NewSessionOptions): SessionManager { const dir = sessionDir ? normalizePath(sessionDir) : getDefaultSessionDir(cwd); - return new SessionManager(cwd, dir, undefined, true); + return new SessionManager(cwd, dir, undefined, true, options); } /** @@ -1356,7 +1382,12 @@ export class SessionManager { * @param targetCwd Target working directory (where the new session will be stored) * @param sessionDir Optional session directory. If omitted, uses default for targetCwd. */ - static forkFrom(sourcePath: string, targetCwd: string, sessionDir?: string): SessionManager { + static forkFrom( + sourcePath: string, + targetCwd: string, + sessionDir?: string, + options?: NewSessionOptions, + ): SessionManager { const resolvedSourcePath = resolvePath(sourcePath); const resolvedTargetCwd = resolvePath(targetCwd); const sourceEntries = loadEntriesFromFile(resolvedSourcePath); @@ -1375,7 +1406,10 @@ export class SessionManager { } // Create new session file with new ID but forked content - const newSessionId = createSessionId(); + if (options?.id !== undefined) { + assertValidSessionId(options.id); + } + const newSessionId = options?.id ?? createSessionId(); const timestamp = new Date().toISOString(); const fileTimestamp = timestamp.replace(/[:.]/g, "-"); const newSessionFile = join(dir, `${fileTimestamp}_${newSessionId}.jsonl`); @@ -1389,7 +1423,7 @@ export class SessionManager { cwd: resolvedTargetCwd, parentSession: resolvedSourcePath, }; - appendFileSync(newSessionFile, `${JSON.stringify(newHeader)}\n`); + writeFileSync(newSessionFile, `${JSON.stringify(newHeader)}\n`, { flag: "wx" }); // Copy all non-header entries from source for (const entry of sourceEntries) { diff --git a/packages/coding-agent/src/main.ts b/packages/coding-agent/src/main.ts index 5395e045..7e69d933 100644 --- a/packages/coding-agent/src/main.ts +++ b/packages/coding-agent/src/main.ts @@ -37,7 +37,7 @@ import { MissingSessionCwdError, type SessionCwdIssue, } from "./core/session-cwd.ts"; -import { SessionManager } from "./core/session-manager.ts"; +import { assertValidSessionId, SessionManager } from "./core/session-manager.ts"; import { SettingsManager } from "./core/settings-manager.ts"; import { printTimings, resetTimings, time } from "./core/timings.ts"; import { runMigrations, showDeprecationWarnings } from "./migrations.ts"; @@ -145,6 +145,16 @@ type ResolvedSession = * Resolve a session argument to a file path. * If it looks like a path, use as-is. Otherwise try to match as session ID prefix. */ +async function findLocalSessionByExactId( + sessionId: string, + cwd: string, + sessionDir?: string, +): Promise<{ type: "local"; path: string } | undefined> { + const localSessions = await SessionManager.list(cwd, sessionDir); + const localMatch = localSessions.find((s) => s.id === sessionId); + return localMatch ? { type: "local", path: localMatch.path } : undefined; +} + async function resolveSessionPath(sessionArg: string, cwd: string, sessionDir?: string): Promise { // If it looks like a file path, resolve it before handing it to the session manager. if (sessionArg.includes("/") || sessionArg.includes("\\") || sessionArg.endsWith(".jsonl")) { @@ -153,19 +163,20 @@ async function resolveSessionPath(sessionArg: string, cwd: string, sessionDir?: // Try to match as session ID in current project first const localSessions = await SessionManager.list(cwd, sessionDir); - const localMatches = localSessions.filter((s) => s.id.startsWith(sessionArg)); + const localMatch = + localSessions.find((s) => s.id === sessionArg) ?? localSessions.find((s) => s.id.startsWith(sessionArg)); - if (localMatches.length >= 1) { - return { type: "local", path: localMatches[0].path }; + if (localMatch) { + return { type: "local", path: localMatch.path }; } // Try global search across all projects const allSessions = await SessionManager.listAll(); - const globalMatches = allSessions.filter((s) => s.id.startsWith(sessionArg)); + const globalMatch = + allSessions.find((s) => s.id === sessionArg) ?? allSessions.find((s) => s.id.startsWith(sessionArg)); - if (globalMatches.length >= 1) { - const match = globalMatches[0]; - return { type: "global", path: match.path, cwd: match.cwd }; + if (globalMatch) { + return { type: "global", path: globalMatch.path, cwd: globalMatch.cwd }; } // Not found anywhere @@ -202,9 +213,33 @@ function validateForkFlags(parsed: Args): void { } } -function forkSessionOrExit(sourcePath: string, cwd: string, sessionDir?: string): SessionManager { +function validateSessionIdFlags(parsed: Args): void { + if (parsed.sessionId === undefined) return; + + const conflictingFlags = [ + parsed.session ? "--session" : undefined, + parsed.continue ? "--continue" : undefined, + parsed.resume ? "--resume" : undefined, + parsed.noSession ? "--no-session" : undefined, + ].filter((flag): flag is string => flag !== undefined); + + if (conflictingFlags.length > 0) { + console.error(chalk.red(`Error: --session-id cannot be combined with ${conflictingFlags.join(", ")}`)); + process.exit(1); + } + try { - return SessionManager.forkFrom(sourcePath, cwd, sessionDir); + assertValidSessionId(parsed.sessionId); + } catch (error: unknown) { + const message = error instanceof Error ? error.message : String(error); + console.error(chalk.red(`Error: ${message}`)); + process.exit(1); + } +} + +function forkSessionOrExit(sourcePath: string, cwd: string, sessionDir?: string, sessionId?: string): SessionManager { + try { + return SessionManager.forkFrom(sourcePath, cwd, sessionDir, { id: sessionId }); } catch (error: unknown) { const message = error instanceof Error ? error.message : String(error); console.error(chalk.red(`Error: ${message}`)); @@ -218,18 +253,26 @@ async function createSessionManager( sessionDir: string | undefined, settingsManager: SettingsManager, ): Promise { - if (parsed.noSession) { - return SessionManager.inMemory(); + if (parsed.noSession || parsed.help || parsed.listModels !== undefined) { + return SessionManager.inMemory(cwd); } if (parsed.fork) { + if (parsed.sessionId) { + const existingTarget = await findLocalSessionByExactId(parsed.sessionId, cwd, sessionDir); + if (existingTarget) { + console.error(chalk.red(`Session already exists with id '${parsed.sessionId}'`)); + process.exit(1); + } + } + const resolved = await resolveSessionPath(parsed.fork, cwd, sessionDir); switch (resolved.type) { case "path": case "local": case "global": - return forkSessionOrExit(resolved.path, cwd, sessionDir); + return forkSessionOrExit(resolved.path, cwd, sessionDir, parsed.sessionId); case "not_found": console.error(chalk.red(`No session found matching '${resolved.arg}'`)); @@ -282,7 +325,14 @@ async function createSessionManager( return SessionManager.continueRecent(cwd, sessionDir); } - return SessionManager.create(cwd, sessionDir); + if (parsed.sessionId) { + const existingSession = await findLocalSessionByExactId(parsed.sessionId, cwd, sessionDir); + if (existingSession) { + return SessionManager.open(existingSession.path, sessionDir); + } + } + + return SessionManager.create(cwd, sessionDir, { id: parsed.sessionId }); } function buildSessionOptions( @@ -483,6 +533,7 @@ export async function main(args: string[], options?: MainOptions) { } validateForkFlags(parsed); + validateSessionIdFlags(parsed); // Run migrations (pass cwd for project-local migrations) const { migratedAuthProviders: migratedProviders, deprecationWarnings } = runMigrations(process.cwd()); diff --git a/packages/coding-agent/test/args.test.ts b/packages/coding-agent/test/args.test.ts index 3bf41473..2f6d9ea9 100644 --- a/packages/coding-agent/test/args.test.ts +++ b/packages/coding-agent/test/args.test.ts @@ -130,6 +130,11 @@ describe("parseArgs", () => { expect(result.session).toBe("/path/to/session.jsonl"); }); + test("parses --session-id", () => { + const result = parseArgs(["--session-id", "orchestrated-session"]); + expect(result.sessionId).toBe("orchestrated-session"); + }); + test("parses --fork", () => { const result = parseArgs(["--fork", "1234abcd"]); expect(result.fork).toBe("1234abcd"); diff --git a/packages/coding-agent/test/session-id-readonly.test.ts b/packages/coding-agent/test/session-id-readonly.test.ts new file mode 100644 index 00000000..b6e97ce8 --- /dev/null +++ b/packages/coding-agent/test/session-id-readonly.test.ts @@ -0,0 +1,131 @@ +import { spawn } from "node:child_process"; +import { existsSync, mkdirSync, mkdtempSync, readdirSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join, resolve } from "node:path"; +import { afterEach, describe, expect, it } from "vitest"; +import { ENV_AGENT_DIR } from "../src/config.ts"; + +const cliPath = resolve(__dirname, "../src/cli.ts"); +const tempDirs: string[] = []; + +afterEach(() => { + for (const dir of tempDirs.splice(0)) { + rmSync(dir, { recursive: true, force: true }); + } +}); + +function createTempDir(): string { + const dir = mkdtempSync(join(tmpdir(), "pi-session-id-readonly-")); + tempDirs.push(dir); + return dir; +} + +function hasSessionWithId(root: string, sessionId: string): boolean { + if (!existsSync(root)) return false; + for (const entry of readdirSync(root, { withFileTypes: true })) { + const path = join(root, entry.name); + if (entry.isDirectory() && hasSessionWithId(path, sessionId)) return true; + if (!entry.isFile() || !entry.name.endsWith(".jsonl")) continue; + + try { + const firstLine = readFileSync(path, "utf8").split("\n", 1)[0]; + const header = JSON.parse(firstLine) as { type?: string; id?: string }; + if (header.type === "session" && header.id === sessionId) return true; + } catch { + // Ignore malformed session files. + } + } + return false; +} + +interface CliDirs { + agentDir: string; + projectDir: string; + sessionDir: string; +} + +async function runCli( + args: string[] | ((dirs: CliDirs) => string[]), + setup?: (dirs: CliDirs) => void, +): Promise<{ code: number | null; agentDir: string; stderr: string }> { + const tempRoot = createTempDir(); + const dirs: CliDirs = { + agentDir: join(tempRoot, "agent"), + projectDir: join(tempRoot, "project"), + sessionDir: join(tempRoot, "sessions"), + }; + mkdirSync(dirs.agentDir, { recursive: true }); + mkdirSync(dirs.projectDir, { recursive: true }); + setup?.(dirs); + const resolvedArgs = typeof args === "function" ? args(dirs) : args; + + let stderr = ""; + const code = await new Promise((resolvePromise, reject) => { + const child = spawn(process.execPath, [cliPath, ...resolvedArgs], { + cwd: dirs.projectDir, + env: { + ...process.env, + [ENV_AGENT_DIR]: dirs.agentDir, + PI_OFFLINE: "1", + TSX_TSCONFIG_PATH: resolve(__dirname, "../../../tsconfig.json"), + }, + stdio: ["ignore", "ignore", "pipe"], + }); + child.stderr.on("data", (chunk) => { + stderr += chunk.toString(); + }); + child.on("error", reject); + child.on("close", resolvePromise); + }); + + return { code, agentDir: dirs.agentDir, stderr }; +} + +function writeSession(sessionDir: string, cwd: string, id: string): void { + writeFileSync( + join(sessionDir, `${id}.jsonl`), + `${JSON.stringify({ type: "session", version: 3, id, timestamp: new Date().toISOString(), cwd })}\n`, + ); +} + +describe("--session-id read-only commands", () => { + it("does not reserve a session for --help", async () => { + const result = await runCli(["--session-id", "read-only-help", "--help"]); + + expect(result.code).toBe(0); + expect(hasSessionWithId(join(result.agentDir, "sessions"), "read-only-help")).toBe(false); + }); + + it("does not reserve a session for --list-models", async () => { + const result = await runCli(["--session-id", "read-only-models", "--list-models"]); + + expect(result.code).toBe(0); + expect(hasSessionWithId(join(result.agentDir, "sessions"), "read-only-models")).toBe(false); + }); + + it("rejects an existing fork target session id", async () => { + const result = await runCli( + (dirs) => ["--session-dir", dirs.sessionDir, "--fork", "source-id", "--session-id", "existing-id", "-p", "hi"], + (dirs) => { + mkdirSync(dirs.sessionDir, { recursive: true }); + writeSession(dirs.sessionDir, dirs.projectDir, "source-id"); + writeSession(dirs.sessionDir, dirs.projectDir, "existing-id"); + }, + ); + + expect(result.code).toBe(1); + expect(result.stderr).toContain("Session already exists with id 'existing-id'"); + }); +}); + +describe("--session-id validation", () => { + it("rejects ids invalid under SessionManager rules without stack traces", async () => { + for (const id of ["-bad", "bad id"]) { + const result = await runCli(["--session-id", id, "-p", "hi"]); + + expect(result.code).toBe(1); + expect(result.stderr).toContain("Session id must be non-empty"); + expect(result.stderr).not.toContain("SessionManager.create"); + } + }); +}); diff --git a/packages/coding-agent/test/session-manager/custom-session-id.test.ts b/packages/coding-agent/test/session-manager/custom-session-id.test.ts index f9c6ad0d..ee2da917 100644 --- a/packages/coding-agent/test/session-manager/custom-session-id.test.ts +++ b/packages/coding-agent/test/session-manager/custom-session-id.test.ts @@ -1,6 +1,6 @@ -import { mkdtempSync, writeFileSync } from "node:fs"; +import { existsSync, mkdtempSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; -import { join } from "node:path"; +import { basename, join } from "node:path"; import { describe, expect, it } from "vitest"; import { SessionManager } from "../../src/core/session-manager.ts"; @@ -13,6 +13,23 @@ describe("SessionManager.newSession with custom id", () => { expect(session.getSessionId()).toBe("my-custom-id"); }); + it("allows alphanumeric session ids with interior punctuation", () => { + const session = SessionManager.inMemory(); + session.newSession({ id: "abc-123_def.456" }); + expect(session.getSessionId()).toBe("abc-123_def.456"); + }); + + it("rejects invalid custom session ids", () => { + const invalidIds = ["", "-abc", "abc-", "_abc", "abc_", ".abc", "abc.", "abc/def", "abc\\def", "abc def"]; + + for (const id of invalidIds) { + const session = SessionManager.inMemory(); + expect(() => session.newSession({ id })).toThrow( + "Session id must be non-empty, contain only alphanumeric characters", + ); + } + }); + it("generates a UUIDv7 id when no id is provided", () => { const session = SessionManager.inMemory(); session.newSession(); @@ -46,6 +63,18 @@ describe("SessionManager.newSession with custom id", () => { expect(session.getHeader()!.id).toBe(session.getSessionId()); }); + it("uses the provided id when creating a persisted session", () => { + const tempDir = mkdtempSync(join(tmpdir(), "pi-session-manager-")); + const session = SessionManager.create(tempDir, tempDir, { id: "created-session-id" }); + + expect(session.getSessionId()).toBe("created-session-id"); + expect(session.getHeader()!.id).toBe("created-session-id"); + const sessionFile = session.getSessionFile()!; + expect(sessionFile).toContain("created-session-id"); + expect(basename(sessionFile)).toMatch(/^\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}-\d{3}Z_created-session-id\.jsonl$/); + expect(existsSync(sessionFile)).toBe(false); + }); + it("generates a UUIDv7 id when creating a branched session", () => { const session = SessionManager.inMemory(); const firstId = session.appendMessage({ @@ -106,4 +135,28 @@ describe("SessionManager.newSession with custom id", () => { expect(header!.id).toMatch(UUID_V7_RE); expect(header!.parentSession).toBe(sourcePath); }); + + it("uses the provided id when forking from another session file", () => { + const tempDir = mkdtempSync(join(tmpdir(), "pi-session-manager-")); + const sourcePath = join(tempDir, "source.jsonl"); + writeFileSync( + sourcePath, + `${JSON.stringify({ + type: "session", + version: 3, + id: "source-session-id", + timestamp: new Date().toISOString(), + cwd: tempDir, + })}\n`, + ); + + const forked = SessionManager.forkFrom(sourcePath, tempDir, tempDir, { id: "forked-session-id" }); + const header = forked.getHeader(); + expect(header).not.toBeNull(); + expect(header!.id).toBe("forked-session-id"); + expect(header!.parentSession).toBe(sourcePath); + const sessionFile = forked.getSessionFile()!; + expect(sessionFile).toContain("forked-session-id"); + expect(basename(sessionFile)).toMatch(/^\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}-\d{3}Z_forked-session-id\.jsonl$/); + }); });