diff --git a/backend/src/pi/pi-process-manager.ts b/backend/src/pi/pi-process-manager.ts index 139a350d..226f967f 100644 --- a/backend/src/pi/pi-process-manager.ts +++ b/backend/src/pi/pi-process-manager.ts @@ -37,6 +37,7 @@ type SpawnFn = ( export class PiProcessManager { private runtimes = new Map(); + private agentSnapshotCleanups = new WeakMap void>(); private spawnFn: ( sessionId: string, author: string, provider: string | undefined, principal?: PrincipalContext, runtimeConfigPath?: string, @@ -58,6 +59,13 @@ export class PiProcessManager { } } + private cleanupAgentSnapshot(child: ChildProcessWithoutNullStreams): void { + const cleanup = this.agentSnapshotCleanups.get(child); + if (!cleanup) return; + this.agentSnapshotCleanups.delete(child); + cleanup(); + } + private spawnPi( spawnFn: SpawnFn, sessionId: string, author: string, provider: string | undefined, principal?: PrincipalContext, runtimeConfigPath?: string, @@ -104,16 +112,19 @@ export class PiProcessManager { cwd: this.cfg.harnessDir, env, }); - child.once("exit", agent.cleanup); - child.once("close", agent.cleanup); + this.agentSnapshotCleanups.set(child, agent.cleanup); + child.once("exit", () => this.cleanupAgentSnapshot(child!)); + child.once("close", () => this.cleanupAgentSnapshot(child!)); // Log stderr for debugging (was silently drained) child.stderr.on("data", (d: Buffer) => console.error(`[pi:${sessionId}] stderr:`, d.toString().trim())); return child; } catch (error) { if (child) { try { child.kill(); } catch { /* preserve the initialization error */ } + this.cleanupAgentSnapshot(child); + } else { + agent.cleanup(); } - agent.cleanup(); throw error; } } @@ -208,6 +219,7 @@ export class PiProcessManager { if (rt && this.runtimes.get(sessionId) === rt) this.runtimes.delete(sessionId); releaseRuntimeConfig(); try { child.kill(); } catch { /* preserve the initialization error */ } + this.cleanupAgentSnapshot(child); throw error; } } @@ -276,6 +288,7 @@ export class PiProcessManager { this.runtimes.delete(id); expected.releaseRuntimeConfig?.(); expected.child.kill(); + this.cleanupAgentSnapshot(expected.child); return true; } } diff --git a/backend/test/pi-process-manager.test.ts b/backend/test/pi-process-manager.test.ts index 3238ecb1..4227d231 100644 --- a/backend/test/pi-process-manager.test.ts +++ b/backend/test/pi-process-manager.test.ts @@ -306,10 +306,18 @@ test("createFor kills a spawned child when post-spawn initialization throws", () child.stdout = { on: () => { throw new Error("READER_INIT_SENTINEL"); }, }; - const mgr = new PiProcessManager(loadConfig({}), { spawnFn: () => child as any }); + let snapshotDir: string | undefined; + const mgr = new PiProcessManager(loadConfig({}), { + spawnFn: (_command, _args, options) => { + snapshotDir = options.env.PI_CODING_AGENT_DIR; + return child as any; + }, + }); expect(() => mgr.createFor("broken-init", {})).toThrow("READER_INIT_SENTINEL"); + expect(snapshotDir).toBeTruthy(); + expect(existsSync(snapshotDir!)).toBe(false); expect(child.kill).toHaveBeenCalledOnce(); expect(mgr.get("broken-init")).toBeUndefined(); expect(mgr.count()).toBe(0);