fix(pi): settle foreground tools before runtime shutdown

Wait for native work to abort while the event stream is still open so Pi persists terminal tool results before restart history is loaded.

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<void>();
+  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);

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<void> {
+    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<void> {
     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]);

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<void> | null = null;
   closeCalls = 0;
   readonly canceledExtensionUiRequests: string[] = [];
   readonly extensionUiResponses: Array<{
@@ -295,6 +296,7 @@ export class FakePiSession implements PiRuntimeSession {
   async abort(): Promise<void> {
     this.controlRequests.push("abort");
     this.abortRequested = true;
+    await this.abortCompletion;
   }
 
   async getState(): Promise<PiSessionState> {
@@ -427,6 +429,7 @@ export class FakePiSession implements PiRuntimeSession {
 
   async close(): Promise<void> {
     this.closeCalls += 1;
+    this.subscribers.clear();
   }
 
   emit(event: PiRuntimeEvent): void {