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 index 15a2b8eba9..1b1932e9cd 100644 --- 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 @@ -9,7 +9,7 @@ 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"; +type TurnOutcome = { assistantMessageId: string; subagentMessageId?: string } | "no-response"; function initMessage(): Record { return { @@ -40,6 +40,21 @@ function assistantReply(uuid: string): Record { }; } +/** + * A message from a Task subagent. The CLI streams these on the same query as the main + * conversation, and their uuids live on the subagent's sidechain, not on the session's + * own message chain, so `forkSession` cannot resolve one. + */ +function subagentReply(uuid: string): Record { + return { + type: "assistant", + uuid, + session_id: SESSION_ID, + parent_tool_use_id: "toolu_subagent_1", + message: { id: uuid, role: "assistant", content: [{ type: "text", text: "subagent work" }] }, + }; +} + function successResult(): Record { return { type: "result", @@ -110,6 +125,9 @@ function createConversation(outcomes: TurnOutcome[]): Conversation { continue; } enqueue(assistantReply(outcome.assistantMessageId)); + if (outcome.subagentMessageId) { + enqueue(subagentReply(outcome.subagentMessageId)); + } enqueue(successResult()); } closedRef.value = true; @@ -229,3 +247,23 @@ describe("Claude rewind across a turn that produced no response", () => { expect(rewindSdk.recordedForks).toEqual([{ upToMessageId: "assistant-1" }]); }); }); + +describe("Claude rewind after a turn whose last assistant message came from a subagent", () => { + test("forks at the turn's own reply, not at the subagent message", async () => { + const conversation = createConversation([ + { assistantMessageId: "assistant-1", subagentMessageId: "subagent-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 30fdebd974..b4b516fe1e 100644 --- a/packages/server/src/server/agent/providers/claude/agent.ts +++ b/packages/server/src/server/agent/providers/claude/agent.ts @@ -3054,6 +3054,13 @@ class ClaudeAgentSession implements AgentSession { if (!messageId) { return; } + // Subagent frames ride the same stream as the conversation, but their uuids live on the + // subagent's sidechain, so forkSession cannot resolve one. Anchoring a turn on one leaves + // rewind throwing "Message not found in session". The persisted history path already + // skips sidechain entries; the live stream has to skip them too. + if (readClaudeParentToolUseId(message)) { + return; + } if ( message.type === "user" && !isSyntheticUserEntry(message) &&