diff --git a/packages/server/src/server/agent/agent-manager.test.ts b/packages/server/src/server/agent/agent-manager.test.ts index 41f784349e..1e8a95eb81 100644 --- a/packages/server/src/server/agent/agent-manager.test.ts +++ b/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 { + 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 { + 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); diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index 5c01be317e..dd6b7cfbb7 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/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"); diff --git a/packages/server/src/server/agent/providers/claude/agent.test.ts b/packages/server/src/server/agent/providers/claude/agent.test.ts index 758fad2e8f..a5d15bdb24 100644 --- a/packages/server/src/server/agent/providers/claude/agent.test.ts +++ b/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) => 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>((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((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((resolve) => { + optionsStarted = resolve; + }); + const allowed = new Promise((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({ diff --git a/packages/server/src/server/agent/providers/claude/agent.ts b/packages/server/src/server/agent/providers/claude/agent.ts index cc30825f94..9b560747fc 100644 --- a/packages/server/src/server/agent/providers/claude/agent.ts +++ b/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(); 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();