From beda7180c01c11180952177ae193c4ddbde530fc Mon Sep 17 00:00:00 2001 From: Ryan Winograd Date: Thu, 29 Oct 2026 16:15:11 -0500 Subject: [PATCH] fix(pi): close history-only runtimes after replay Archived timeline reads resume a Pi RPC process to reconstruct the saved conversation. The history-only session remained open after replay, so browsing or synchronizing archived agents left one idle Pi process per agent. Mark history-purpose resumes for closure and close the runtime in streamHistory's finally path. Interactive sessions remain resident, and idempotent close keeps later manager cleanup safe. --- .../server/agent/providers/pi/agent.test.ts | 40 +++++++++++++++++++ .../src/server/agent/providers/pi/agent.ts | 29 +++++++++----- .../agent/providers/pi/test-utils/fake-pi.ts | 5 ++- 3 files changed, 63 insertions(+), 11 deletions(-) diff --git a/packages/server/src/server/agent/providers/pi/agent.test.ts b/packages/server/src/server/agent/providers/pi/agent.test.ts index 7dc3fd184c..35e9bf7d92 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/packages/server/src/server/agent/providers/pi/agent.test.ts @@ -1600,6 +1600,46 @@ describe("PiRpcAgentSession", () => { ]); }); + test("closes a history-only Pi runtime after replaying its timeline", async () => { + const pi = new FakePi(); + pi.queueSessionSetup((runtimeSession) => { + runtimeSession.messages = [ + { role: "user", content: "Old question" }, + { role: "assistant", content: [{ type: "text", text: "Old answer" }] }, + ]; + }); + + const client: import("../../agent-sdk-types.js").AgentClient = createClient(pi); + const session = await client.resumeSession( + { + provider: "pi", + sessionId: "pi-session-1", + nativeHandle: "/tmp/native-pi-session", + metadata: { cwd: "/workspace/project" }, + }, + {}, + undefined, + { purpose: "history" }, + ); + + const history: AgentStreamEvent[] = []; + for await (const event of session.streamHistory()) history.push(event); + + expect(history).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + type: "timeline", + provider: "pi", + item: expect.objectContaining({ type: "assistant_message", text: "Old answer" }), + }), + ]), + ); + expect(pi.latestSession().closeCalls).toBe(1); + + await session.close(); + expect(pi.latestSession().closeCalls).toBe(1); + }); + test("resumes a session whose model was removed on the model Pi falls back to", async () => { const pi = new FakePi(); pi.removeModel("9router/deepseek-v4-flash"); diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index 8051c57fa9..c60a2b44b7 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/packages/server/src/server/agent/providers/pi/agent.ts @@ -26,6 +26,7 @@ import { type AgentRunOptions, type AgentRunResult, type AgentRuntimeInfo, + type AgentResumeSessionOptions, type AgentSelectOption, type AgentSession, type AgentSessionConfig, @@ -253,6 +254,7 @@ interface PiRpcAgentSessionOptions { currentModeId?: string | null; persistedSessionFile?: string; cleanup?: () => void; + closeAfterHistory?: boolean; extensionTimeoutMs?: number; logger: Logger; usagePollScheduler?: PiUsagePollScheduler; @@ -1255,6 +1257,7 @@ export class PiRpcAgentSession implements AgentSession { this.currentModeId = options.currentModeId ?? null; this.persistedSessionFile = options.persistedSessionFile; this.cleanup = options.cleanup; + this.closeAfterHistory = options.closeAfterHistory ?? false; this.extensionTimeoutMs = options.extensionTimeoutMs ?? DEFAULT_PI_EXTENSION_RESULT_TIMEOUT_MS; this.logger = options.logger; this.extensionHost = createPiExtensionHost(this.logger); @@ -1286,6 +1289,7 @@ export class PiRpcAgentSession implements AgentSession { private readonly config: AgentSessionConfig; private readonly persistedSessionFile?: string; private readonly cleanup?: () => void; + private readonly closeAfterHistory: boolean; private readonly extensionTimeoutMs: number; get id(): string | null { @@ -1454,16 +1458,20 @@ export class PiRpcAgentSession implements AgentSession { } async *streamHistory(): AsyncGenerator { - await this.requestEntryCapture("history"); - yield* streamPiHistory( - this.provider, - await this.runtimeSession.getMessages(), - this.contextUserEntries, - {}, - // At most eight 2 MiB child files per replay; later cards retain their summaries. - createPiExtensionHost(this.logger, undefined, 16 * 1024 * 1024), - this.closeController.signal, - ); + try { + await this.requestEntryCapture("history"); + yield* streamPiHistory( + this.provider, + await this.runtimeSession.getMessages(), + this.contextUserEntries, + {}, + // At most eight 2 MiB child files per replay; later cards retain their summaries. + createPiExtensionHost(this.logger, undefined, 16 * 1024 * 1024), + this.closeController.signal, + ); + } finally { + if (this.closeAfterHistory) await this.close(); + } } async getRuntimeInfo(): Promise { @@ -2738,6 +2746,7 @@ export class PiRpcAgentClient implements AgentClient { capabilities: capabilitiesForSession(mcp !== null), persistedSessionFile: sessionCopy ? sessionFile : undefined, cleanup: cleanupTempFiles, + closeAfterHistory: options?.purpose === "history", extensionTimeoutMs: providerOptions.extensionTimeoutMs, logger: this.logger, usagePollScheduler: this.usagePollScheduler, diff --git a/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts b/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts index 90a80e4c66..f80f5930e7 100644 --- a/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts +++ b/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts @@ -153,6 +153,7 @@ export class FakePiSession implements PiRuntimeSession { // The user entries on the current branch that getMessages() replays. contextUserEntries: FakePiUserEntry[] = []; abortRequested = false; + closeCalls = 0; readonly canceledExtensionUiRequests: string[] = []; readonly extensionUiResponses: Array<{ id: string; @@ -422,7 +423,9 @@ export class FakePiSession implements PiRuntimeSession { this.respondToExtensionUiRequest(id, { cancelled: true }); } - async close(): Promise {} + async close(): Promise { + this.closeCalls += 1; + } emit(event: PiRuntimeEvent): void { for (const subscriber of this.subscribers) {