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 039d3eaf22..501df2c1c2 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/packages/server/src/server/agent/providers/pi/agent.test.ts @@ -208,6 +208,52 @@ async function createSession( return { pi, session, events }; } +test("close preserves a foreground tool's terminal update before releasing Pi", async () => { + const { pi, session, events } = await createSession(); + const native = pi.latestSession(); + const stopped = Promise.withResolvers(); + native.abortCompletion = stopped.promise; + await session.startTurn("Run a foreground command"); + native.emit({ + type: "tool_execution_start", + toolCallId: "foreground-bash", + toolName: "bash", + args: { command: "sleep 120" }, + }); + + const closing = session.close(); + await flushTurnScheduling(); + // Pi's abort acknowledgement follows its final tool result and durable transcript write. + native.emit({ + type: "tool_execution_end", + toolCallId: "foreground-bash", + toolName: "bash", + result: { content: [{ type: "text", text: "Command aborted" }] }, + isError: true, + }); + stopped.resolve(); + await closing; + + expect( + events + .timelineItems() + .findLast((item) => item.type === "tool_call" && item.callId === "foreground-bash"), + ).toMatchObject({ + callId: "foreground-bash", + status: "failed", + }); +}); + +test("close releases Pi when stopping its active work fails", async () => { + const { pi, session } = await createSession(); + const native = pi.latestSession(); + await session.startTurn("Run a foreground command"); + native.clearQueueError = new Error("Pi control plane is unavailable"); + + await expect(session.close()).resolves.toBeUndefined(); + expect(native.closeCalls).toBe(1); +}); + test("forwards launch-context env to the Pi process launch", async () => { const pi = new FakePi(); const client = createClient(pi); diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index 198d596c9b..004bbfb053 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/packages/server/src/server/agent/providers/pi/agent.ts @@ -1540,16 +1540,7 @@ export class PiRpcAgentSession implements AgentSession { this.activeTurnId || this.activeTurnStarted ? { turnId, error: null } : null; this.interruptingTurn = interruption; try { - try { - await this.runtimeSession.clearQueue(); - } catch (error) { - // COMPAT(piClearQueueFallback): added in v0.5.0, remove after 2027-03-01 once - // the pi floor supports clear_queue (added in pi 0.84.4). - if (!isPiMissingClearQueueRpc(error)) { - throw error; - } - } - await this.runtimeSession.abort(); + await this.abortRuntimeWork(); } catch (error) { const terminalError = this.interruptingTurn === interruption ? interruption?.error : null; if (this.interruptingTurn === interruption) { @@ -1601,6 +1592,19 @@ export class PiRpcAgentSession implements AgentSession { } } + private async abortRuntimeWork(): Promise { + try { + await this.runtimeSession.clearQueue(); + } catch (error) { + // COMPAT(piClearQueueFallback): added in v0.5.0, remove after 2027-03-01 once + // the pi floor supports clear_queue (added in pi 0.84.4). + if (!isPiMissingClearQueueRpc(error)) { + throw error; + } + } + await this.runtimeSession.abort(); + } + async revertConversation(input: { messageId: string }): Promise { if (this.activeTurnId) { throw new Error("Cannot rewind the Pi conversation while a turn is active"); @@ -1655,6 +1659,13 @@ export class PiRpcAgentSession implements AgentSession { this.extensionHost.close(); this.usagePoller.close(); try { + if (this.activeTurnId || this.activeTurnStarted) { + // Pi persists terminal tool results before acknowledging abort. Keep its + // event stream alive until then so restart history cannot strand the tool. + await this.abortRuntimeWork().catch((err) => { + this.logger.warn({ err }, "Failed to abort Pi work before closing its runtime"); + }); + } await this.runtimeSession.close(); } finally { await Promise.all([this.pendingOutput, ...this.pendingExtensionHydrations]); 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 770f24625c..51eca92a41 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 { // Native user entries paired with the complete active-branch history. historyUserEntries: FakePiUserEntry[] = []; abortRequested = false; + abortCompletion: Promise | null = null; closeCalls = 0; readonly canceledExtensionUiRequests: string[] = []; readonly extensionUiResponses: Array<{ @@ -295,6 +296,7 @@ export class FakePiSession implements PiRuntimeSession { async abort(): Promise { this.controlRequests.push("abort"); this.abortRequested = true; + await this.abortCompletion; } async getState(): Promise { @@ -427,6 +429,7 @@ export class FakePiSession implements PiRuntimeSession { async close(): Promise { this.closeCalls += 1; + this.subscribers.clear(); } emit(event: PiRuntimeEvent): void {