diff --git a/apps/server/src/claudeHistoryWorker.ts b/apps/server/src/claudeHistoryWorker.ts index 3fd707f3c7a7d..ca7401df6f768 100644 --- a/apps/server/src/claudeHistoryWorker.ts +++ b/apps/server/src/claudeHistoryWorker.ts @@ -24,7 +24,12 @@ export async function runClaudeHistoryWorker( method === "getSessionMessages" ? await getSessionMessages(sessionId, options) : method === "forkSession" - ? await forkSession(sessionId, options) + ? await forkSession(sessionId, options).then(async (forked) => { + const headMessageId = await getSessionMessages(forked.sessionId, options) + .then((messages) => messages.at(-1)?.uuid) + .catch(() => undefined); + return headMessageId === undefined ? forked : { ...forked, headMessageId }; + }) : (() => { throw new Error("Unknown Claude history operation."); })(); diff --git a/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.test.ts b/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.test.ts index 19311662d756d..57994165593a7 100644 --- a/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.test.ts +++ b/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.test.ts @@ -2536,6 +2536,8 @@ describe("ClaudeAdapterV2 background wake turns", () => { readonly nativeSessionId?: string; readonly settings?: ClaudeSettings; readonly freshQueueOnReopen?: boolean; + readonly forkSession?: ClaudeAdapterV2.ClaudeAgentSdkQueryRunnerShape["forkSession"]; + readonly allocateSessionId?: Effect.Effect; }) => Effect.gen(function* () { const fileSystem = yield* FileSystem.FileSystem; @@ -2600,7 +2602,9 @@ describe("ClaudeAdapterV2 background wake turns", () => { }).pipe(Effect.andThen(Queue.offer(continuationReceipts, request))), }, queryRunner: { - allocateSessionId: Effect.succeed(options?.nativeSessionId ?? WAKE_NATIVE_SESSION), + allocateSessionId: + options?.allocateSessionId ?? + Effect.succeed(options?.nativeSessionId ?? WAKE_NATIVE_SESSION), open: (input) => Effect.gen(function* () { openedOptions = input.options; @@ -2662,7 +2666,7 @@ describe("ClaudeAdapterV2 background wake turns", () => { close: options?.close?.(sdkMessages) ?? Effect.void, }; }), - forkSession: () => Effect.die("unused forkSession"), + forkSession: options?.forkSession ?? (() => Effect.die("unused forkSession")), subagentLaunchToolUseId: () => Effect.succeed(options?.subagentLaunchToolUseId ?? null), assertComplete: Effect.void, }, @@ -3748,6 +3752,104 @@ describe("ClaudeAdapterV2 background wake turns", () => { close: (sdkMessages) => Queue.shutdown(sdkMessages), }); + it.effect.each([ + { kind: "fork with inherited head", fork: true, headMessageId: "remapped-inherited-head" }, + { kind: "fork without recorded head", fork: true, headMessageId: undefined }, + { kind: "ordinary thread", fork: false, headMessageId: undefined }, + ])("rolls back the first run of a $kind", ({ fork, headMessageId }) => + Effect.gen(function* () { + let allocations = 0; + const harness = yield* makeWakeHarnessWithOptions({ + close: (sdkMessages) => Queue.shutdown(sdkMessages), + freshQueueOnReopen: true, + allocateSessionId: Effect.sync(() => + allocations++ === 0 ? WAKE_NATIVE_SESSION : "fresh-native-session", + ), + forkSession: () => + Effect.succeed({ + sessionId: WAKE_NATIVE_SESSION, + ...(headMessageId === undefined ? {} : { headMessageId }), + }), + }); + const providerThread = fork + ? yield* harness.runtime.forkThread({ + sourceProviderThread: { + ...harness.providerThread, + nativeThreadRef: { + driver: ClaudeAdapterV2.CLAUDE_PROVIDER, + nativeId: "source-native-session", + strength: "strong", + }, + }, + targetThreadId: harness.threadId, + }) + : harness.providerThread; + const turnInput = makeClaudeTestTurnInput({ + threadId: harness.threadId, + providerThread, + now: yield* DateTime.now, + attemptId: RunAttemptId.make("first-run"), + text: "This run will be rolled back.", + attachments: [], + nativeThreadHasTurns: false, + }); + yield* harness.runtime.startTurn(turnInput); + yield* harness.offerAndWait( + makeAssistantTextFrame({ uuid: "own-run-message", text: "Discard this answer." }), + ); + yield* harness.offerAndWait( + makeResultFrame({ uuid: "own-run-result", result: "Discard this answer." }), + ); + assert.equal((yield* Queue.take(harness.terminalReceipts)).status, "completed"); + const completed = harness.events.findLast( + (event) => event.type === "provider_thread.updated", + ); + assert.equal(completed?.type, "provider_thread.updated"); + if (completed?.type !== "provider_thread.updated") return; + assert.equal(completed.providerThread.firstRunOrdinal, 1); + + const rolledBack = yield* harness.runtime.rollbackThread({ + providerThread: completed.providerThread, + target: { + type: "thread_start", + checkpointId: CheckpointId.make("checkpoint-before-first-run"), + appRunOrdinal: 0, + }, + providerThreadTurns: [lastProviderTurn(harness.events)], + }); + assert.equal( + rolledBack.providerThread.forkedFrom?.providerThreadId, + providerThread.forkedFrom?.providerThreadId, + ); + yield* harness.runtime.startTurn({ + ...turnInput, + attemptId: RunAttemptId.make("after-rollback"), + providerThread: rolledBack.providerThread, + providerTurnOrdinal: 2, + nativeThreadHasTurns: false, + }); + const options = harness.getOpenedOptions(); + assert.equal(harness.processQueues.length, 2); + if (headMessageId !== undefined) { + assert.equal(rolledBack.providerThread.nativeThreadRef?.nativeId, WAKE_NATIVE_SESSION); + assert.equal(rolledBack.providerThread.nativeConversationHeadRef?.nativeId, headMessageId); + assert.equal(options?.resume, WAKE_NATIVE_SESSION); + assert.equal(options?.resumeSessionAt, headMessageId); + assert.isUndefined(options?.sessionId); + } else { + assert.equal(rolledBack.providerThread.nativeThreadRef?.nativeId, "fresh-native-session"); + assert.isNull(rolledBack.providerThread.nativeConversationHeadRef); + assert.equal(options?.sessionId, "fresh-native-session"); + assert.isUndefined(options?.resume); + assert.isUndefined(options?.resumeSessionAt); + } + }).pipe( + Effect.provide( + Layer.mergeAll(IdAllocator.layer, McpProviderSessions.layer, NodeServices.layer), + ), + ), + ); + it.effect( "rolls back to the main-thread assistant message when a subagent frame comes last", () => diff --git a/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.ts b/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.ts index d6aada46a9025..c82ed54afaf7d 100644 --- a/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.ts +++ b/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.ts @@ -26,6 +26,7 @@ import { type EffortLevel, type ForkSessionOptions, type ForkSessionResult, + getSessionMessages, getSubagentMessages, query, type Options as ClaudeQueryOptions, @@ -497,7 +498,10 @@ export interface ClaudeAgentSdkQueryRunnerShape { ) => Effect.Effect; readonly forkSession: ( input: ClaudeAgentSdkSessionForkInput, - ) => Effect.Effect; + ) => Effect.Effect< + ForkSessionResult & { readonly headMessageId?: string }, + ClaudeAgentSdkQueryRunnerError + >; readonly subagentLaunchToolUseId: ( input: ClaudeAgentSdkSubagentLookupInput, ) => Effect.Effect; @@ -748,7 +752,9 @@ export function makeClaudeAgentSdkProtocolLogger(input: { const encodeClaudeHistoryOptions = Schema.encodeSync(Schema.fromJsonString(Schema.Unknown)); const decodeClaudeHistoryFork = Schema.decodeUnknownEffect( - Schema.fromJsonString(Schema.Struct({ sessionId: Schema.String })), + Schema.fromJsonString( + Schema.Struct({ sessionId: Schema.String, headMessageId: Schema.optionalKey(Schema.String) }), + ), ); const claudeConfigDirOf = ( @@ -1080,14 +1086,18 @@ export const layerQueryRunner: Layer.Layer< }); const hostEnvironment = yield* HostProcess.Environment; const homeDirectory = yield* HostProcess.HomeDirectory; - const fork: Effect.Effect = + const usesHostConfig = claudeConfigDirOf(path, input.environment, homeDirectory) === - claudeConfigDirOf(path, hostEnvironment, homeDirectory) - ? Effect.tryPromise(() => forkClaudeSession(input.sessionId, input.options)) - : forkClaudeSessionInWorker(input).pipe( - Effect.provideService(ChildProcessSpawner.ChildProcessSpawner, spawner), - Effect.provideService(Path.Path, path), - ); + claudeConfigDirOf(path, hostEnvironment, homeDirectory); + const fork: Effect.Effect< + ForkSessionResult & { readonly headMessageId?: string }, + unknown + > = usesHostConfig + ? Effect.tryPromise(() => forkClaudeSession(input.sessionId, input.options)) + : forkClaudeSessionInWorker(input).pipe( + Effect.provideService(ChildProcessSpawner.ChildProcessSpawner, spawner), + Effect.provideService(Path.Path, path), + ); const result = yield* fork.pipe( Effect.mapError((cause) => queryRunnerError(cause, "forkSession")), ); @@ -1099,7 +1109,22 @@ export const layerQueryRunner: Layer.Layer< sessionId: result.sessionId, }, }); - return result; + // The SDK remaps copied UUIDs. Capture the inherited head before the + // fork's first run appends messages. + const headMessageId = usesHostConfig + ? yield* Effect.tryPromise({ + try: () => + getSessionMessages( + result.sessionId, + input.options.dir === undefined ? {} : { dir: input.options.dir }, + ), + catch: (cause) => queryRunnerError(cause, "getSessionMessages"), + }).pipe( + Effect.map((messages) => messages.at(-1)?.uuid), + Effect.orElseSucceed(() => undefined), + ) + : result.headMessageId; + return headMessageId === undefined ? result : { ...result, headMessageId }; }), subagentLaunchToolUseId: Effect.fn("ClaudeAgentSdkQueryRunner.subagentLaunchToolUseId")( function* (input: ClaudeAgentSdkSubagentLookupInput) { @@ -1479,6 +1504,7 @@ function makeProviderThread(input: { readonly providerSessionId: OrchestrationV2ProviderThread["providerSessionId"]; readonly nativeThreadId: string; readonly forkedFrom?: NonNullable; + readonly nativeMetadata?: NonNullable; readonly now: DateTime.Utc; }): OrchestrationV2ProviderThread { return { @@ -1502,6 +1528,7 @@ function makeProviderThread(input: { lastRunOrdinal: null, handoffIds: [], forkedFrom: input.forkedFrom ?? null, + ...(input.nativeMetadata === undefined ? {} : { nativeMetadata: input.nativeMetadata }), createdAt: input.now, updatedAt: input.now, }; @@ -1607,7 +1634,7 @@ const resolveClaudeRollbackResumeSessionAt = Effect.fn( )(function* (input: ProviderAdapter.ProviderAdapterV2RollbackThreadInput) { switch (input.target.type) { case "thread_start": - return null; + return input.providerThread.nativeMetadata?.forkHeadMessageId ?? null; case "provider_turn": { const target = input.target; if (target.providerTurn.providerThreadId !== input.providerThread.id) { @@ -9788,8 +9815,9 @@ export const makeClaudeAdapterV2 = Effect.fn("makeClaudeAdapterV2")(function* ( const nativeThreadId = yield* getNativeThreadId(rollbackInput.providerThread); yield* closeLiveQueryForNativeThread(nativeThreadId); const now = yield* DateTime.now; + const resumeSessionAt = yield* resolveClaudeRollbackResumeSessionAt(rollbackInput); - if (rollbackInput.target.type === "thread_start") { + if (rollbackInput.target.type === "thread_start" && resumeSessionAt === null) { const resetNativeThreadId = yield* queryRunner.allocateSessionId; return { providerThread: { @@ -9815,7 +9843,6 @@ export const makeClaudeAdapterV2 = Effect.fn("makeClaudeAdapterV2")(function* ( }; } - const resumeSessionAt = yield* resolveClaudeRollbackResumeSessionAt(rollbackInput); return { providerThread: { ...rollbackInput.providerThread, @@ -9893,6 +9920,9 @@ export const makeClaudeAdapterV2 = Effect.fn("makeClaudeAdapterV2")(function* ( ? {} : { providerTurnId: forkInput.providerTurnId }), }, + ...(forked.headMessageId === undefined + ? {} + : { nativeMetadata: { forkHeadMessageId: forked.headMessageId } }), now, }); }, diff --git a/apps/server/src/orchestration-v2/Adapters/ClaudeAgentSdkQueryRunner.test.ts b/apps/server/src/orchestration-v2/Adapters/ClaudeAgentSdkQueryRunner.test.ts index c9c091d0e75cb..693e2cbab09f0 100644 --- a/apps/server/src/orchestration-v2/Adapters/ClaudeAgentSdkQueryRunner.test.ts +++ b/apps/server/src/orchestration-v2/Adapters/ClaudeAgentSdkQueryRunner.test.ts @@ -151,6 +151,13 @@ it.effect("forks a session stored under the instance's own CLAUDE_CONFIG_DIR", ( ); assert.include(transcript, '"content":"hello"'); assert.include(transcript, '"text":"hi"'); + const inheritedHead = transcript + .trim() + .split("\n") + .findLast((line) => line.includes('"type":"assistant"')) + ?.match(/"uuid"\s*:\s*"([^"]+)"/)?.[1]; + assert.isString(inheritedHead); + assert.equal(forked.headMessageId, inheritedHead); assert.equal(process.env.CLAUDE_CONFIG_DIR, hostConfigDir); assert.equal(inProcessFork.mock.calls.length, 0); }).pipe(Effect.provide(forkRunnerLayer)), @@ -187,6 +194,52 @@ const realQueryRunner = layerQueryRunner.pipe( Layer.provide(NodeServices.layer), ); +it.effect.each(["readable", "unreadable"] as const)( + "forks with a %s transcript, retaining the inherited head when available", + (transcript) => + Effect.gen(function* () { + let forkCompleted = false; + const fork = vi.spyOn(ClaudeSdk, "forkSession").mockImplementation(async () => { + forkCompleted = true; + return { sessionId: "forked-session" }; + }); + const read = vi.spyOn(ClaudeSdk, "getSessionMessages").mockImplementation(async () => { + assert.isTrue(forkCompleted); + if (transcript === "unreadable") throw new Error("transcript unreadable"); + return ["remapped-user", "remapped-assistant"].map((uuid) => ({ + type: "assistant" as const, + uuid, + session_id: "forked-session", + message: { role: "assistant", content: uuid }, + parent_tool_use_id: null, + parent_agent_id: null, + })); + }); + yield* Effect.addFinalizer(() => + Effect.sync(() => { + fork.mockRestore(); + read.mockRestore(); + }), + ); + const runner = yield* ClaudeAgentSdkQueryRunner; + const forked = yield* runner.forkSession({ + sessionId: "source-session", + options: { dir: "/workspace", upToMessageId: "source-assistant" }, + environment: process.env, + threadId: ThreadId.make("thread-fork-head"), + providerSessionId: ProviderSessionId.make("session-fork-head"), + }); + assert.deepEqual(forked, { + sessionId: "forked-session", + ...(transcript === "readable" ? { headMessageId: "remapped-assistant" } : {}), + }); + assert.deepEqual(fork.mock.calls, [ + ["source-session", { dir: "/workspace", upToMessageId: "source-assistant" }], + ]); + assert.deepEqual(read.mock.calls, [["forked-session", { dir: "/workspace" }]]); + }).pipe(Effect.provide(realQueryRunner)), +); + function fakeClaudeCli(script: string): string { const dir = NodeFS.mkdtempSync(NodePath.join(NodeOS.tmpdir(), "t3-claude-cli-")); const cli = NodePath.join(dir, "claude"); diff --git a/packages/contracts/src/orchestrationV2.ts b/packages/contracts/src/orchestrationV2.ts index 60c05ab7ede38..e6709630a6232 100644 --- a/packages/contracts/src/orchestrationV2.ts +++ b/packages/contracts/src/orchestrationV2.ts @@ -845,6 +845,7 @@ export const OrchestrationV2ProviderThreadNativeMetadata = Schema.Struct({ title: Schema.optional(Schema.NullOr(TrimmedNonEmptyString)), updatedAt: Schema.optional(Schema.NullOr(TrimmedNonEmptyString)), itemIdentityVersion: Schema.optional(Schema.Literal(2)), + forkHeadMessageId: Schema.optional(TrimmedNonEmptyString), }); export type OrchestrationV2ProviderThreadNativeMetadata = typeof OrchestrationV2ProviderThreadNativeMetadata.Type;