fix: release Pi snapshots on failed initialization

This commit is contained in:
2026-08-05 16:59:49 +02:00
parent 840db9bd4f
commit 7efaec434f
2 changed files with 25 additions and 4 deletions
+16 -3
View File
@@ -37,6 +37,7 @@ type SpawnFn = (
export class PiProcessManager {
private runtimes = new Map<string, SessionRuntime>();
private agentSnapshotCleanups = new WeakMap<ChildProcessWithoutNullStreams, () => 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;
}
}
+9 -1
View File
@@ -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);