diff --git a/packages/server/src/server/agent/providers/claude/agent.rewind-anchors.test.ts b/packages/server/src/server/agent/providers/claude/agent.rewind-anchors.test.ts new file mode 100644 index 0000000000..15a2b8eba9 --- /dev/null +++ b/packages/server/src/server/agent/providers/claude/agent.rewind-anchors.test.ts @@ -0,0 +1,231 @@ +import { describe, expect, test, vi } from "vitest"; + +import { createTestLogger } from "../../../../test-utils/test-logger.js"; +import { ClaudeAgentClient } from "./agent.js"; +import { FakeClaudeSdk } from "./test-rewind-claude-sdk.js"; +import { streamSession } from "../test-utils/session-stream-adapter.js"; +import type { AgentSession } from "../../agent-sdk-types.js"; + +const SESSION_ID = "session-1"; + +/** Claude Code replies to a turn either with an assistant message or with nothing at all. */ +type TurnOutcome = { assistantMessageId: string } | "no-response"; + +function initMessage(): Record { + return { + type: "system", + subtype: "init", + session_id: SESSION_ID, + permissionMode: "default", + model: "claude-sonnet-4-6", + }; +} + +/** Claude Code echoes the user message back under the uuid Paseo minted for it. */ +function userEcho(uuid: string): Record { + return { + type: "user", + uuid, + session_id: SESSION_ID, + message: { role: "user", content: [{ type: "text", text: "prompt" }] }, + }; +} + +function assistantReply(uuid: string): Record { + return { + type: "assistant", + uuid, + session_id: SESSION_ID, + message: { id: uuid, role: "assistant", content: [{ type: "text", text: "ok" }] }, + }; +} + +function successResult(): Record { + return { + type: "result", + subtype: "success", + uuid: "result-ok", + session_id: SESSION_ID, + duration_ms: 1, + duration_api_ms: 1, + is_error: false, + num_turns: 1, + result: "", + total_cost_usd: 0, + usage: { input_tokens: 1, cache_read_input_tokens: 0, output_tokens: 1 }, + permission_denials: [], + }; +} + +/** A turn that dies upstream before the model emits its first token. */ +function failedResult(): Record { + return { + type: "result", + subtype: "error_during_execution", + uuid: "result-failed", + session_id: SESSION_ID, + duration_ms: 1, + duration_api_ms: 1, + is_error: true, + num_turns: 1, + errors: ["upstream gateway returned 503"], + total_cost_usd: 0, + usage: { input_tokens: 1, cache_read_input_tokens: 0, output_tokens: 0 }, + permission_denials: [], + }; +} + +interface Conversation { + queryFactory: ReturnType; + /** The uuid Paseo minted for each prompt, in the order the turns ran. */ + userMessageIds: string[]; +} + +function createConversation(outcomes: TurnOutcome[]): Conversation { + const userMessageIds: string[] = []; + const queryFactory = vi.fn(({ prompt }: { prompt: AsyncIterable }) => { + const queued: Array> = []; + const waiters: Array<() => void> = []; + const closedRef = { value: false }; + let turnIndex = 0; + + const wake = () => waiters.shift()?.(); + const enqueue = (message: Record) => { + queued.push(message); + wake(); + }; + + void (async () => { + for await (const sent of prompt) { + const userMessageId = String((sent as { uuid?: unknown }).uuid); + userMessageIds.push(userMessageId); + const outcome = outcomes[turnIndex]; + turnIndex += 1; + if (turnIndex === 1) { + enqueue(initMessage()); + } + enqueue(userEcho(userMessageId)); + if (outcome === "no-response" || !outcome) { + enqueue(failedResult()); + continue; + } + enqueue(assistantReply(outcome.assistantMessageId)); + enqueue(successResult()); + } + closedRef.value = true; + wake(); + })(); + + return { + next: vi.fn(async () => { + while (queued.length === 0 && !closedRef.value) { + await new Promise((resolve) => { + waiters.push(resolve); + }); + } + if (queued.length === 0) { + return { done: true, value: undefined }; + } + return { done: false, value: queued.shift() }; + }), + interrupt: vi.fn(async () => undefined), + return: vi.fn(async () => { + closedRef.value = true; + wake(); + return undefined; + }), + close: vi.fn(() => { + closedRef.value = true; + wake(); + }), + setPermissionMode: vi.fn(async () => undefined), + setModel: vi.fn(async () => undefined), + supportedModels: vi.fn(async () => []), + supportedCommands: vi.fn(async () => []), + rewindFiles: vi.fn(async () => ({ canRewind: true })), + [Symbol.asyncIterator]() { + return this; + }, + }; + }); + return { queryFactory, userMessageIds }; +} + +async function createSession( + conversation: Conversation, + rewindSdk: FakeClaudeSdk, +): Promise { + const client = new ClaudeAgentClient({ + logger: createTestLogger(), + queryFactory: conversation.queryFactory as never, + resolveBinary: async () => "/test/claude/bin", + rewindSdk, + }); + return client.createSession({ + provider: "claude", + cwd: process.cwd(), + model: "claude-sonnet-4-6", + }); +} + +async function runTurns(session: AgentSession, count: number): Promise { + for (let index = 0; index < count; index += 1) { + for await (const _ of streamSession(session, `turn ${index + 1}`)) { + // drain to the terminal event + } + } +} + +describe("Claude rewind across a turn that produced no response", () => { + test("forks at the most recent turn that did produce a response", async () => { + const conversation = createConversation([ + { assistantMessageId: "assistant-1" }, + "no-response", + { assistantMessageId: "assistant-3" }, + ]); + const rewindSdk = new FakeClaudeSdk(); + const session = await createSession(conversation, rewindSdk); + + try { + await runTurns(session, 3); + await session.revertConversation?.({ messageId: conversation.userMessageIds[2] }); + } finally { + await session.close(); + } + + expect(rewindSdk.recordedForks).toEqual([{ upToMessageId: "assistant-1" }]); + }); + + test("starts a fresh session when no earlier turn produced a response", async () => { + const conversation = createConversation(["no-response", { assistantMessageId: "assistant-2" }]); + const rewindSdk = new FakeClaudeSdk(); + const session = await createSession(conversation, rewindSdk); + + try { + await runTurns(session, 2); + await session.revertConversation?.({ messageId: conversation.userMessageIds[1] }); + } finally { + await session.close(); + } + + expect(rewindSdk.recordedForks).toEqual([]); + }); + + test("still forks at the previous turn when every turn produced a response", async () => { + const conversation = createConversation([ + { assistantMessageId: "assistant-1" }, + { assistantMessageId: "assistant-2" }, + ]); + const rewindSdk = new FakeClaudeSdk(); + const session = await createSession(conversation, rewindSdk); + + try { + await runTurns(session, 2); + await session.revertConversation?.({ messageId: conversation.userMessageIds[1] }); + } finally { + await session.close(); + } + + expect(rewindSdk.recordedForks).toEqual([{ upToMessageId: "assistant-1" }]); + }); +}); diff --git a/packages/server/src/server/agent/providers/claude/agent.ts b/packages/server/src/server/agent/providers/claude/agent.ts index b23f58e244..30fdebd974 100644 --- a/packages/server/src/server/agent/providers/claude/agent.ts +++ b/packages/server/src/server/agent/providers/claude/agent.ts @@ -76,7 +76,12 @@ import { } from "./options.js"; import { renderPromptAttachmentAsText } from "../../prompt-attachments.js"; import { claudeQuery, type ClaudeOptions, type ClaudeQueryFactory } from "./query.js"; -import { realClaudeRewindSdk, revertClaudeConversation, revertClaudeFiles } from "./rewind.js"; +import { + realClaudeRewindSdk, + revertClaudeConversation, + revertClaudeFiles, + type ClaudeRewindSdk, +} from "./rewind.js"; import { normalizeProviderReplayTimestamp } from "../../provider-history-timestamps.js"; import { claudeProjectDirSync } from "./project-dir.js"; import { THINKING_APPLIES_NEXT_TURN_NOTICE } from "../../provider-notices.js"; @@ -405,6 +410,7 @@ interface ClaudeAgentClientOptions { resolveBinary?: () => Promise; resolveVersion?: (signal?: AbortSignal) => Promise; configDir?: string; + rewindSdk?: ClaudeRewindSdk; } interface ClaudeAgentSessionOptions { @@ -417,6 +423,7 @@ interface ClaudeAgentSessionOptions { logger: Logger; queryFactory?: ClaudeQueryFactory; resolveBinary: () => Promise; + rewindSdk?: ClaudeRewindSdk; } type ClaudeThinkingEffort = "low" | "medium" | "high" | "xhigh" | "max"; @@ -1500,6 +1507,7 @@ export class ClaudeAgentClient implements AgentClient { private readonly resolveBinary: () => Promise; private readonly resolveVersion: (signal?: AbortSignal) => Promise; private readonly configDir?: string; + private readonly rewindSdk: ClaudeRewindSdk; constructor(options: ClaudeAgentClientOptions) { this.defaults = options.defaults; @@ -1511,6 +1519,7 @@ export class ClaudeAgentClient implements AgentClient { options.resolveVersion ?? ((signal) => resolveClaudeCodeVersion(this.runtimeSettings, signal)); this.configDir = options.configDir; + this.rewindSdk = options.rewindSdk ?? realClaudeRewindSdk; } resolveConfiguredModel(model: AgentModelDefinition): AgentModelDefinition { @@ -1532,6 +1541,7 @@ export class ClaudeAgentClient implements AgentClient { logger: this.logger, queryFactory: this.queryFactory, resolveBinary: this.resolveBinary, + rewindSdk: this.rewindSdk, }); } @@ -1560,6 +1570,7 @@ export class ClaudeAgentClient implements AgentClient { logger: this.logger, queryFactory: this.queryFactory, resolveBinary: this.resolveBinary, + rewindSdk: this.rewindSdk, }); } @@ -2105,6 +2116,7 @@ class ClaudeAgentSession implements AgentSession { private userMessageIds: string[] = []; private readonly emittedUserMessageIds = new Set(); private readonly rewindTurnAnchors: ClaudeRewindTurnAnchor[] = []; + private readonly rewindSdk: ClaudeRewindSdk; private pendingFreshSessionId: string | null = null; private recentStderr = ""; private closed = false; @@ -2120,6 +2132,7 @@ class ClaudeAgentSession implements AgentSession { this.logger = options.logger.child({ agentId: this.agentId }); this.queryFactory = options.queryFactory; this.resolveBinary = options.resolveBinary; + this.rewindSdk = options.rewindSdk ?? realClaudeRewindSdk; this.contextUsage = new ClaudeContextUsageState( findClaudeModel(this.config.model)?.contextWindowMaxTokens, ); @@ -2760,7 +2773,7 @@ class ClaudeAgentSession implements AgentSession { return; } await revertClaudeConversation({ - sdk: realClaudeRewindSdk, + sdk: this.rewindSdk, sessionId: this.claudeSessionId, messageId: target.messageId, resolveMessageId: (messageId) => this.resolveClaudeMessageId(messageId), @@ -3076,17 +3089,18 @@ class ClaudeAgentSession implements AgentSession { throw new Error(`Claude rewind target ${messageId} is not in the tracked conversation`); } - if (index === 0) { - return { kind: "fresh-session" }; + // A turn the model never answered — an instant gateway error, or an abort before the first + // token — leaves its anchor without an assistant id, and carries no conversation state worth + // preserving. Fork at the most recent turn before the target that did answer. Reading only + // the immediately preceding anchor let one such turn fail every later rewind in the + // conversation. + for (let previous = index - 1; previous >= 0; previous -= 1) { + const assistantMessageId = this.rewindTurnAnchors[previous]?.assistantMessageId; + if (assistantMessageId) { + return { kind: "fork", messageId: assistantMessageId }; + } } - - const previousTurn = this.rewindTurnAnchors[index - 1]; - if (!previousTurn?.assistantMessageId) { - throw new Error( - `Claude rewind cannot preserve turn ${index} because its assistant response id was not observed`, - ); - } - return { kind: "fork", messageId: previousTurn.assistantMessageId }; + return { kind: "fresh-session" }; } private async ensureQuery(): Promise {