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.

nix/npm-deps.hash

@@ -1 +1 @@
-sha256-iAi82L6Lm75hdg/9tJYyYcHMhCJLtXGSpcNqgv9QQaw=
+sha256-VdM4lXdgNEXPpTfvHSFvF9nm302rK/ZbVkZUp9tyOfg=

packages/server/src/server/agent/agent-manager.test.ts

@@ -10991,6 +10991,88 @@ test("concurrent explicit closes tear down the runtime once", async () => {
   }
 });
 
+test("close drains provider tool cleanup before retiring the runtime", async () => {
+  const workdir = mkdtempSync(join(tmpdir(), "agent-manager-close-cleanup-"));
+  const storage = new AgentStorage(join(workdir, "agents"), logger);
+  const timeline = new InMemoryAgentTimelineStore();
+  class CleanupSession extends TestAgentSession {
+    override async startTurn(): Promise<{ turnId: string }> {
+      const turnId = "cleanup-turn";
+      this.pushEvent({ type: "turn_started", provider: this.provider, turnId });
+      for (const callId of ["first-tool", "second-tool"]) {
+        this.pushEvent({
+          type: "timeline",
+          provider: this.provider,
+          turnId,
+          item: {
+            type: "tool_call",
+            callId,
+            name: "Bash",
+            detail: { type: "shell", command: "sleep 10" },
+            status: "running",
+            error: null,
+          },
+        });
+      }
+      return { turnId };
+    }
+    override async close(): Promise<void> {
+      for (const callId of ["first-tool", "second-tool"]) {
+        this.pushEvent({
+          type: "timeline",
+          provider: this.provider,
+          turnId: "cleanup-turn",
+          item: {
+            type: "tool_call",
+            callId,
+            name: "Bash",
+            detail: { type: "shell", command: "sleep 10" },
+            status: "canceled",
+            error: null,
+          },
+        });
+      }
+    }
+  }
+  const client = new (class extends TestAgentClient {
+    override async createSession(config: AgentSessionConfig): Promise<AgentSession> {
+      return new CleanupSession(config);
+    }
+  })();
+  const manager = new AgentManager({
+    clients: { codex: client },
+    registry: storage,
+    timelineStore: timeline,
+    logger,
+  });
+  const agent = await manager.createAgent({ provider: "codex", cwd: workdir }, undefined, {
+    workspaceId: undefined,
+  });
+  const run = drainAsyncGenerator(manager.streamAgent(agent.id, "run both tools"));
+  try {
+    await manager.waitForAgentRunStart(agent.id);
+    await manager.flush();
+    await manager.closeAgent(agent.id);
+    await run;
+    await manager.flush();
+    expect(
+      timeline
+        .getItems(agent.id)
+        .filter((item) => item.type === "tool_call")
+        .map((item) => ({ callId: item.callId, status: item.status })),
+    ).toEqual([
+      { callId: "first-tool", status: "canceled" },
+      { callId: "second-tool", status: "canceled" },
+    ]);
+    expect((await storage.get(agent.id))?.lastStatus).toBe("closed");
+  } finally {
+    await manager.closeAgent(agent.id).catch(() => undefined);
+    await run;
+    await manager.flush();
+    rmSync(workdir, { recursive: true, force: true });
+  }
+});
+
 test("provider close failure retains the runtime for cleanup instead of allowing another writer", async () => {
   const workdir = mkdtempSync(join(tmpdir(), "agent-manager-close-failure-"));
   const storage = new AgentStorage(join(workdir, "agents"), logger);

packages/server/src/server/agent/agent-manager.ts

@@ -1947,6 +1947,7 @@ export class AgentManager {
     // Retain ownership until shutdown succeeds. A failed close may still own a
     // native writer, so publishing a resumable closed snapshot would orphan it.
     await agent.session.close();
+    await this.drainSessionEvents(agentId);
     this.cancelRunningProviderSubagents(agentId);
     const closedAgent = this.prepareAgentForClosure(agent, "agent closed");
 

packages/server/src/server/agent/providers/claude/agent.test.ts

@@ -729,6 +729,196 @@ describe("ClaudeAgentSession features", () => {
     return { queryFactory, queryMock, launches };
   }
 
+  function createStreamingQueryMock() {
+    const mock = createQueryMock();
+    const messages: SDKMessage[] = [];
+    let pending: ((result: IteratorResult<SDKMessage, void>) => void) | null = null;
+    let ended = false;
+    const finish = () => {
+      ended = true;
+      pending?.({ done: true, value: undefined });
+      pending = null;
+    };
+    mock.queryMock.close.mockImplementation(finish);
+    mock.queryMock.return.mockImplementation(async () => finish());
+    mock.queryMock[Symbol.asyncIterator] = () => ({
+      next: async () => {
+        const message = messages.shift();
+        if (message) return { done: false, value: message };
+        if (ended) return { done: true, value: undefined };
+        return new Promise<IteratorResult<SDKMessage, void>>((resolve) => {
+          pending = resolve;
+        });
+      },
+    });
+    return {
+      ...mock,
+      push(message: SDKMessage) {
+        if (pending) {
+          pending({ done: false, value: message });
+          pending = null;
+        } else {
+          messages.push(message);
+        }
+      },
+    };
+  }
+
+  async function startToolForClosure() {
+    const mock = createStreamingQueryMock();
+    const session = await new ClaudeAgentClient({
+      logger,
+      queryFactory: mock.queryFactory,
+      resolveBinary: async () => "/test/claude/bin",
+    }).createSession({ provider: "claude", cwd: process.cwd(), modeId: "default" });
+    const events: AgentStreamEvent[] = [];
+    let toolStarted!: () => void;
+    const running = new Promise<void>((resolve) => {
+      toolStarted = resolve;
+    });
+    const unsubscribe = session.subscribe((event) => {
+      events.push(event);
+      if (
+        event.type === "timeline" &&
+        event.item.type === "tool_call" &&
+        event.item.status === "running"
+      )
+        toolStarted();
+    });
+    await session.startTurn("run the tool");
+    mock.push({
+      type: "stream_event",
+      session_id: "closure-session",
+      uuid: "00000000-0000-4000-8000-000000000001",
+      parent_tool_use_id: null,
+      event: {
+        type: "content_block_start",
+        index: 0,
+        content_block: { type: "tool_use", id: "closure-tool", name: "Bash", input: {} },
+      },
+    });
+    mock.push({
+      type: "stream_event",
+      session_id: "closure-session",
+      uuid: "00000000-0000-4000-8000-000000000002",
+      parent_tool_use_id: null,
+      event: { type: "content_block_stop", index: 0 },
+    });
+    await running;
+    return { ...mock, session, events, unsubscribe };
+  }
+
+  test("runtime close publishes tool cleanup without reporting a user turn cancellation", async () => {
+    const fixture = await startToolForClosure();
+    try {
+      await fixture.session.close();
+      expect(fixture.queryMock.close).toHaveBeenCalledTimes(1);
+      expect(
+        fixture.events.filter(
+          (event) =>
+            event.type === "timeline" &&
+            event.item.type === "tool_call" &&
+            event.item.status === "canceled",
+        ),
+      ).toHaveLength(1);
+      expect(
+        fixture.events.filter(
+          (event) =>
+            event.type === "turn_canceled" ||
+            event.type === "turn_completed" ||
+            event.type === "turn_failed",
+        ),
+      ).toEqual([]);
+    } finally {
+      fixture.unsubscribe();
+      await fixture.session.close();
+    }
+  });
+
+  test("explicit interruption still publishes tool cleanup and one turn cancellation", async () => {
+    const fixture = await startToolForClosure();
+    try {
+      await fixture.session.interrupt();
+      expect(
+        fixture.events.filter(
+          (event) =>
+            event.type === "timeline" &&
+            event.item.type === "tool_call" &&
+            event.item.status === "canceled",
+        ),
+      ).toHaveLength(1);
+      expect(fixture.events.filter((event) => event.type === "turn_canceled")).toHaveLength(1);
+    } finally {
+      fixture.unsubscribe();
+      await fixture.session.close();
+    }
+  });
+
+  test("runtime close resolves permissions once without reporting a user turn cancellation", async () => {
+    const { queryFactory } = createQueryMock();
+    const session = await new ClaudeAgentClient({
+      logger,
+      queryFactory,
+      resolveBinary: async () => "/test/claude/bin",
+    }).createSession({ provider: "claude", cwd: process.cwd(), modeId: "default" });
+    const events: AgentStreamEvent[] = [];
+    const unsubscribe = session.subscribe((event) => events.push(event));
+    const abort = new AbortController();
+    try {
+      await session.startTurn("run the tool");
+      const canUseTool = queryFactory.mock.calls[0]?.[0].options.canUseTool;
+      if (!canUseTool) throw new Error("Expected canUseTool callback");
+      const permission = canUseTool(
+        "Bash",
+        { command: "printf test" },
+        { signal: abort.signal, toolUseID: "tool-closed" },
+      );
+      const rejected = expect(permission).rejects.toThrow("Claude session closed");
+      await session.close();
+      abort.abort();
+      await rejected;
+      expect(session.getPendingPermissions()).toEqual([]);
+      expect(events.filter(isPermissionResolvedEvent)).toHaveLength(1);
+      expect(events.filter((event) => event.type === "turn_canceled")).toEqual([]);
+    } finally {
+      unsubscribe();
+      await session.close();
+    }
+  });
+
+  test("runtime close prevents a query from launching after pending options resolve", async () => {
+    const { queryFactory } = createQueryMock();
+    let optionsStarted!: () => void;
+    let allowOptions!: () => void;
+    const started = new Promise<void>((resolve) => {
+      optionsStarted = resolve;
+    });
+    const allowed = new Promise<void>((resolve) => {
+      allowOptions = resolve;
+    });
+    const session = await new ClaudeAgentClient({
+      logger,
+      queryFactory,
+      resolveBinary: async () => {
+        optionsStarted();
+        await allowed;
+        return "/test/claude/bin";
+      },
+    }).createSession({ provider: "claude", cwd: process.cwd(), modeId: "default" });
+    const turn = session.startTurn("must not run after closure");
+    try {
+      await started;
+      await session.close();
+      allowOptions();
+      await turn;
+      expect(queryFactory).not.toHaveBeenCalled();
+    } finally {
+      allowOptions();
+      await turn;
+      await session.close();
+    }
+  });
+
   test("publishes a resolution when the SDK aborts a permission callback", async () => {
     const { queryFactory } = createQueryMock();
     const session = await new ClaudeAgentClient({

packages/server/src/server/agent/providers/claude/agent.ts

@@ -3018,7 +3018,9 @@ class ClaudeAgentSession implements AgentSession {
     this.closed = true;
     await this.ensureHistoryReady();
     this.rejectAllPendingPermissions(new Error("Claude session closed"));
-    this.cancelCurrentTurn?.();
+    // AgentManager settles the run after runtime release. Preserve cleanup records
+    // without reporting this closure as an explicit turn cancellation.
+    this.flushPendingToolCalls();
     this.subscribers.clear();
     this.activeForegroundTurnId = null;
     this.activeForegroundQuery = null;
@@ -3525,6 +3527,8 @@ class ClaudeAgentSession implements AgentSession {
 
     const input = createAsyncMessageInput<SDKUserMessage>();
     const options = await this.buildOptions(launchMode);
+    // Closure can finish while launch options are being prepared.
+    if (this.closed) throw new Error("Claude session is closed");
     this.logger.debug({ options: summarizeClaudeOptionsForLog(options) }, "claude query");
     this.input = input;
     this.runtimeSafety.reset();

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 {