diff --git a/packages/server/src/server/agent/providers/claude/agent.interrupt-restart-regression.test.ts b/packages/server/src/server/agent/providers/claude/agent.interrupt-restart-regression.test.ts index 02a09ba44b..e90907cad8 100644 --- a/packages/server/src/server/agent/providers/claude/agent.interrupt-restart-regression.test.ts +++ b/packages/server/src/server/agent/providers/claude/agent.interrupt-restart-regression.test.ts @@ -2,6 +2,7 @@ import { afterEach, expect, test, vi } from "vitest"; import { createTestLogger } from "../../../../test-utils/test-logger.js"; import { ClaudeAgentClient } from "./agent.js"; +import { AgentManager } from "../../agent-manager.js"; import { streamSession } from "../test-utils/session-stream-adapter.js"; import type { AgentSession, AgentStreamEvent } from "../../agent-sdk-types.js"; @@ -12,6 +13,7 @@ interface QueryMock { close: ReturnType; setPermissionMode: ReturnType; setModel: ReturnType; + applyFlagSettings: ReturnType; supportedModels: ReturnType; supportedCommands: ReturnType; rewindFiles: ReturnType; @@ -148,6 +150,7 @@ function createScriptedQuery(params: { close: vi.fn(() => undefined), setPermissionMode: vi.fn(async () => undefined), setModel: vi.fn(async () => undefined), + applyFlagSettings: vi.fn(async () => undefined), supportedModels: vi.fn(async () => [{ value: "opus", displayName: "Opus" }]), supportedCommands: vi.fn(async () => []), rewindFiles: vi.fn(async () => ({ canRewind: true })), @@ -276,6 +279,126 @@ afterEach(() => { queryFactory.mockReset(); }); +test.each([false, true])( + "submits after a thinking change with no live query (previous query ended: %s)", + async (previousQueryEnded) => { + const queries: ScriptedQuery[] = []; + queryFactory.mockImplementation(({ prompt }: { prompt: AsyncIterable }) => { + const scriptedQuery = createScriptedQuery({ + prompt, + sessionId: "thinking-restart-session", + handlePrompt: ({ query }) => query.emit(buildSuccessResult("thinking-restart-session")), + }); + queries.push(scriptedQuery); + return scriptedQuery; + }); + const session = await new ClaudeAgentClient({ + logger: createTestLogger(), + queryFactory, + resolveBinary: async () => "/test/claude/bin", + }).createSession({ provider: "claude", cwd: process.cwd() }); + + try { + if (previousQueryEnded) { + await collectUntilTerminal(streamSession(session, "first prompt")); + queries[0].end(); + // Let the pump retire the old query before changing thinking. + await new Promise((resolve) => setTimeout(resolve, 0)); + } + await session.setThinkingOption!("medium"); + const events = await collectUntilTerminal( + streamSession(session, "continue", { clientMessageId: "client-continue" }), + ); + expect(events.filter((event) => event.type === "turn_failed")).toEqual([]); + expect(events.at(-1)?.type).toBe("turn_completed"); + expect(queryFactory).toHaveBeenCalledTimes(previousQueryEnded ? 2 : 1); + expect(queries.at(-1)?.prompts.map((prompt) => prompt.text)).toEqual(["continue"]); + } finally { + await session.close(); + } + }, +); + +test("publishes the submitted prompt identity before an immediate provider failure", async () => { + queryFactory.mockImplementation(({ prompt }: { prompt: AsyncIterable }) => + createScriptedQuery({ + prompt, + sessionId: "immediate-failure-session", + handlePrompt: ({ query }) => + query.emit({ + type: "result", + subtype: "error_during_execution", + session_id: "immediate-failure-session", + errors: ["session limit reached"], + }), + }), + ); + const session = await new ClaudeAgentClient({ + logger: createTestLogger(), + queryFactory, + resolveBinary: async () => "/test/claude/bin", + }).createSession({ provider: "claude", cwd: process.cwd() }); + try { + const events = await collectUntilTerminal( + streamSession(session, "continue", { clientMessageId: "client-failed-prompt" }), + ); + expect(events.at(-1)?.type).toBe("turn_failed"); + expect(events).toContainEqual( + expect.objectContaining({ + type: "timeline", + item: expect.objectContaining({ + type: "user_message", + text: "continue", + clientMessageId: "client-failed-prompt", + messageId: expect.any(String), + }), + }), + ); + } finally { + await session.close(); + } +}); + +test("rewind resolves a failed prompt's client identity to its Claude checkpoint", async () => { + const queries: ScriptedQuery[] = []; + queryFactory.mockImplementation(({ prompt }: { prompt: AsyncIterable }) => { + const query = createScriptedQuery({ + prompt, + sessionId: "failed-prompt-rewind-session", + handlePrompt: ({ query: activeQuery }) => + activeQuery.emit({ + type: "result", + subtype: "error_during_execution", + session_id: "failed-prompt-rewind-session", + errors: ["session limit reached"], + }), + }); + queries.push(query); + return query; + }); + const client = new ClaudeAgentClient({ + logger: createTestLogger(), + queryFactory, + resolveBinary: async () => "/test/claude/bin", + }); + const manager = new AgentManager({ clients: { claude: client }, logger: createTestLogger() }); + const agent = await manager.createAgent({ provider: "claude", cwd: process.cwd() }, undefined, { + workspaceId: undefined, + }); + try { + await collectUntilTerminal( + manager.streamAgent(agent.id, "continue", { clientMessageId: "client-failed-prompt" }), + ); + expect(manager.getAgent(agent.id)?.lastError).toBe("session limit reached"); + await manager.rewind(agent.id, "client-failed-prompt", "files"); + expect(queries[0].rewindFiles).toHaveBeenCalledWith(queries[0].prompts[0].uuid, { + dryRun: false, + }); + } finally { + await manager.closeAgent(agent.id); + } +}); + test("interrupt only calls query.interrupt and leaves the query open", async () => { const logger = createTestLogger(); const queries: ScriptedQuery[] = []; diff --git a/packages/server/src/server/agent/providers/claude/agent.ts b/packages/server/src/server/agent/providers/claude/agent.ts index 9936969e7d..925a8412de 100644 --- a/packages/server/src/server/agent/providers/claude/agent.ts +++ b/packages/server/src/server/agent/providers/claude/agent.ts @@ -2403,13 +2403,9 @@ class ClaudeAgentSession implements AgentSession { } this.activeForegroundQuery = this.query; this.activeForegroundInput = this.input; - this.startQueryPump(); this.input.push(sdkMessage); - setTimeout(() => { - if (this.activeForegroundTurnId === turnId) { - this.emitSubmittedUserMessage(sdkMessage, turnId, options?.clientMessageId); - } - }, 0); + this.emitSubmittedUserMessage(sdkMessage, turnId, options?.clientMessageId); + this.startQueryPump(); } catch (error) { if (sdkUserMessageId) this.unstartedMessageUuids.delete(sdkUserMessageId); this.finishForegroundTurn( @@ -3275,6 +3271,9 @@ class ClaudeAgentSession implements AgentSession { } } + // A fresh query already uses the current settings, even when there was no old query to retire. + this.queryRestartNeeded = false; + // Preserve claudeSessionId across query recreation so buildOptions() passes // resume: sessionId and the new query continues the existing conversation. this.persistence = null; @@ -3877,11 +3876,12 @@ class ClaudeAgentSession implements AgentSession { } private startQueryPump(): void { - if (this.closed || this.queryPumpPromise) { + if (this.closed || this.queryPumpPromise || !this.query) { return; } - const pump = this.runQueryPump().catch((error) => { + // Drain the initialized query; a pending settings change belongs to the next turn. + const pump = this.runQueryPump(this.query).catch((error) => { this.logger.trace( { agentId: this.agentId, @@ -3902,25 +3902,7 @@ class ClaudeAgentSession implements AgentSession { }); } - private async runQueryPump(): Promise { - let activeQuery: Query; - try { - activeQuery = await this.ensureQuery(); - } catch (error) { - this.logger.trace( - { - agentId: this.agentId, - provider: "claude", - sessionId: this.claudeSessionId, - turnId: this.activeForegroundTurnId ?? this.autonomousTurn?.id ?? undefined, - err: error, - }, - "provider.claude.query_pump.init_failed", - ); - this.failActiveTurns(error instanceof Error ? error.message : "Claude stream failed"); - return; - } - + private async runQueryPump(activeQuery: Query): Promise { let consecutiveInterruptAbortRecoveries = 0; const logRawMessage = (message: SDKMessage): void => { this.logger.trace(