From e02e8f5bcbc6ed0390b3d5a672ab48ca5304248f Mon Sep 17 00:00:00 2001 From: Neil <4138956+nwparker@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:55:36 -0700 Subject: [PATCH] fix(chat): preserve Claude rewind context after rollback (#31231) * fix(chat): preserve Claude rewind context after rollback Fixes STA-11420. * fix(lint): merge duplicate journal-types import * fix(chat): skip torn Claude transcript lines when materializing a rewind fork --- ...vitest-real-agent-home-write-guard.test.ts | 12 ++ .../vitest-real-agent-home-write-guard.ts | 13 +- src/main/claude/claude-rewind-fork.ts | 11 +- .../claude/claude-rewind-transcript.test.ts | 120 ++++++++++++++++++ src/main/claude/claude-rewind-transcript.ts | 76 +++++++++++ .../claude/claude-stream-json-connection.ts | 2 +- .../claude-structured-real-cli-rewind.test.ts | 35 +++-- .../journal-epoch-controller.ts | 5 +- .../journal-epoch-replacement.ts | 6 + .../agent-session-journal/journal-store.ts | 7 +- .../structured-agent-session-adapter.ts | 5 +- ...ructured-agent-session-fork-rewind.test.ts | 109 +++++++++++++++- .../structured-agent-session-rewind.ts | 3 +- .../structured-rewind-recovery.ts | 14 +- .../agent-session-rewind-completion.ts | 56 +++++++- src/shared/agent-session-rewind.ts | 4 +- 16 files changed, 444 insertions(+), 34 deletions(-) create mode 100644 src/main/claude/claude-rewind-transcript.test.ts create mode 100644 src/main/claude/claude-rewind-transcript.ts diff --git a/config/scripts/vitest-real-agent-home-write-guard.test.ts b/config/scripts/vitest-real-agent-home-write-guard.test.ts index 9bd42217f67..e33f2779da2 100644 --- a/config/scripts/vitest-real-agent-home-write-guard.test.ts +++ b/config/scripts/vitest-real-agent-home-write-guard.test.ts @@ -49,6 +49,18 @@ describe('vitest real-agent-home write guard', () => { expect(existsSync(target)).toBe(false) }) + it('allows only Claude transcript writes for the real CLI opt-in', () => { + vi.stubEnv('ORCA_REAL_CLAUDE_CLI_TEST', '1') + const missingTranscript = missingRealFolder(join('.claude', 'projects')) + expect(() => rmSync(missingTranscript, { force: true })).not.toThrow() + expect(() => rmSync(missingRealFolder('.codex'), { force: true })).toThrow( + /real-agent-home guard/ + ) + expect(() => rmSync(missingRealFolder('.orca'), { force: true })).toThrow( + /real-agent-home guard/ + ) + }) + it('refuses a promise write under a ~/.claude.json sibling', async () => { const target = join(missingRealFolder('.claude.json'), 'x.tmp') await expect(writeFile(target, '{}')).rejects.toThrow(/real-agent-home guard/) diff --git a/config/scripts/vitest-real-agent-home-write-guard.ts b/config/scripts/vitest-real-agent-home-write-guard.ts index e1f3fc3ee39..ea92e89f80f 100644 --- a/config/scripts/vitest-real-agent-home-write-guard.ts +++ b/config/scripts/vitest-real-agent-home-write-guard.ts @@ -93,9 +93,20 @@ function targetPath(target: unknown): string | null { return null } +function realClaudeTranscriptWriteOptedIn(path: string): boolean { + if (process.env.ORCA_REAL_CLAUDE_CLI_TEST !== '1') { + return false + } + const homes = [ + ...realHomes().map((home) => join(home, '.claude')), + process.env.CLAUDE_CONFIG_DIR + ].filter((home): home is string => Boolean(home)) + return homes.some((home) => normalize(path).startsWith(normalize(join(home, 'projects')) + sep)) +} + function assertOutsideRealHome(state: GuardState, operation: string, target: unknown): void { const path = targetPath(target) - if (path === null || realAgentSuiteOptedIn()) { + if (path === null || realAgentSuiteOptedIn() || realClaudeTranscriptWriteOptedIn(path)) { return } const normalized = normalize(path) + sep diff --git a/src/main/claude/claude-rewind-fork.ts b/src/main/claude/claude-rewind-fork.ts index 7ea9a60faf1..1e84677ebb3 100644 --- a/src/main/claude/claude-rewind-fork.ts +++ b/src/main/claude/claude-rewind-fork.ts @@ -11,6 +11,7 @@ import type { StructuredAgentSessionAdapter } from '../native-chat/agent-session import { resolveSessionFilePath } from '../native-chat/session-file-resolver' import { splitTranscriptStreamLines } from '../native-chat/transcript-stream-lines' import { claudeSessionIdForRewind } from './claude-structured-session-id' +import { materializeClaudeRewindTranscript } from './claude-rewind-transcript' type TranscriptPromptRow = { type: 'user'; uuid: string; parentUuid?: unknown } @@ -77,12 +78,20 @@ export async function resolveClaudeRewindFork(input: { claudeConfigDir }) if (parent !== undefined) { - return { + const fork = { sessionId: claudeSessionIdForRewind(input.orcaSessionId, input.rewindOperationId), sourceSessionId: input.prompt.sessionId, ...(parent ? { throughEntryId: parent } : {}), replacesSessionId: input.head.nativeId } + const sourcePath = await resolveSessionFilePath('claude', input.prompt.sessionId, { + claudeProjectsDir: join(claudeConfigDir, 'projects'), + strictErrors: true + }) + if (!sourcePath) { + return null + } + return materializeClaudeRewindTranscript(sourcePath, fork) } } return null diff --git a/src/main/claude/claude-rewind-transcript.test.ts b/src/main/claude/claude-rewind-transcript.test.ts new file mode 100644 index 00000000000..66f7e0c3096 --- /dev/null +++ b/src/main/claude/claude-rewind-transcript.test.ts @@ -0,0 +1,120 @@ +import { randomUUID } from 'node:crypto' +import { mkdir, mkdtemp, readFile, readdir, rm, stat, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, beforeEach, describe, expect, it } from 'vitest' +import { claudeProviderHandle } from '../../shared/agent-session-provider-handle-encoding' +import { resolveClaudeRewindFork } from './claude-rewind-fork' + +let root: string +let home: string +let project: string +let source: string +let path: string +let firstPrompt: string +let firstAnswer: string +let secondPrompt: string +let original: string + +beforeEach(async () => { + root = await mkdtemp(join(tmpdir(), 'orca-rewind-transcript-')) + home = join(root, 'account') + project = join(home, 'projects', 'workspace') + await mkdir(project, { recursive: true }) + source = randomUUID() + firstPrompt = randomUUID() + firstAnswer = randomUUID() + secondPrompt = randomUUID() + path = join(project, `${source}.jsonl`) + const entries = [ + { + type: 'user', + uuid: firstPrompt, + parentUuid: null, + message: { role: 'user', content: 'ALPHA' } + }, + { + type: 'assistant', + uuid: firstAnswer, + parentUuid: firstPrompt, + message: { role: 'assistant', content: [{ type: 'text', text: 'ALPHA' }] } + }, + { + type: 'user', + uuid: secondPrompt, + parentUuid: firstAnswer, + message: { role: 'user', content: 'BRAVO' } + } + ] + .map((entry) => + JSON.stringify({ + ...entry, + sessionId: source, + timestamp: '2026-10-20T00:00:00.000Z', + isSidechain: false + }) + ) + .join('\n') + original = `${entries}\n` + await writeFile(path, original) +}) + +afterEach(async () => { + await rm(root, { recursive: true, force: true }) +}) + +function forkBefore(uuid: string) { + return resolveClaudeRewindFork({ + orcaSessionId: 'orca-chat', + rewindOperationId: randomUUID(), + head: claudeProviderHandle(source, secondPrompt), + prompt: { sessionId: source, uuid }, + claudeConfigDirs: [join(root, 'other-account'), home] + }) +} + +describe('durable Claude rewind transcripts', () => { + it('saves only kept context in its source account before the next send', async () => { + const fork = await forkBefore(secondPrompt) + expect(fork).toMatchObject({ + materialized: true, + sourceSessionId: source, + throughEntryId: firstAnswer + }) + if (!fork) { + throw new Error('fork missing') + } + const forkPath = join(project, `${fork.sessionId}.jsonl`) + const contents = await readFile(forkPath, 'utf8') + expect(contents).toContain('ALPHA') + expect(contents).not.toContain('BRAVO') + expect(contents).not.toContain(secondPrompt) + expect(contents).toContain(`"messageUuid":"${firstAnswer}"`) + expect(fork.resumeCursor).toBeTruthy() + expect(fork.resumeCursor).not.toBe(firstAnswer) + expect(await readFile(path, 'utf8')).toBe(original) + expect((await readdir(project)).sort()).toEqual( + [`${source}.jsonl`, `${fork.sessionId}.jsonl`].sort() + ) + if (process.platform !== 'win32') { + expect((await stat(forkPath)).mode & 0o777).toBe(0o600) + } + }) + + it('starts an empty fork when no turn remains', async () => { + expect(await forkBefore(firstPrompt)).toMatchObject({ materialized: true }) + expect(await readdir(project)).toEqual([`${source}.jsonl`]) + expect(await readFile(path, 'utf8')).toBe(original) + }) + + it('skips a torn transcript line instead of failing the rewind', async () => { + const damaged = `{broken-json}\n${original}{"type":` + await writeFile(path, damaged) + const fork = await forkBefore(secondPrompt) + expect(fork).toMatchObject({ materialized: true, throughEntryId: firstAnswer }) + const contents = await readFile(join(project, `${fork?.sessionId}.jsonl`), 'utf8') + expect(contents).toContain('ALPHA') + expect(contents).not.toContain('BRAVO') + expect(await readFile(path, 'utf8')).toBe(damaged) + }) +}) diff --git a/src/main/claude/claude-rewind-transcript.ts b/src/main/claude/claude-rewind-transcript.ts new file mode 100644 index 00000000000..18ead956751 --- /dev/null +++ b/src/main/claude/claude-rewind-transcript.ts @@ -0,0 +1,76 @@ +import type { SessionStoreEntry } from '@anthropic-ai/claude-agent-sdk' +import { createReadStream } from 'node:fs' +import { rm } from 'node:fs/promises' +import { dirname, join } from 'node:path' +import type { AgentSessionRewindFork } from '../../shared/agent-session-rewind' +import { durableWriteTempPath, renameDurable, writeTempFileDurable } from '../durable-file-write' +import { splitTranscriptStreamLines } from '../native-chat/transcript-stream-lines' +import { loadClaudeAgentSdk } from './claude-stream-json-connection' + +function isTranscriptEntry(value: unknown): value is SessionStoreEntry { + return ( + typeof value === 'object' && value !== null && 'type' in value && typeof value.type === 'string' + ) +} + +function parseTranscriptEntry(line: string): SessionStoreEntry | undefined { + try { + const entry: unknown = JSON.parse(line) + return isTranscriptEntry(entry) ? entry : undefined + } catch { + return undefined + } +} + +/** Use Claude's fork implementation for branch selection, compaction and UUID remapping. */ +export async function materializeClaudeRewindTranscript( + sourcePath: string, + fork: AgentSessionRewindFork +): Promise { + if (!fork.throughEntryId) { + return { ...fork, materialized: true } + } + const { forkSession } = await loadClaudeAgentSdk() + let resumeCursor: string | undefined + const result = await forkSession(fork.sourceSessionId, { + upToMessageId: fork.throughEntryId, + // A store pins both reads and writes to the resolved account, without changing process.env. + sessionStore: { + load: async () => { + const entries: SessionStoreEntry[] = [] + for await (const { line } of splitTranscriptStreamLines(createReadStream(sourcePath))) { + if (!line.trim()) { + continue + } + const entry = parseTranscriptEntry(line) + // Skip torn lines (crash mid-write) like the SDK's own local reader does. + if (entry) { + entries.push(entry) + } + } + return entries + }, + append: async (key, entries) => { + resumeCursor = entries.findLast( + (entry) => entry.type === 'assistant' || entry.type === 'user' + )?.uuid + if (!resumeCursor) { + throw new Error('claude rewind fork is missing its resume point') + } + const destination = join(dirname(sourcePath), `${key.sessionId}.jsonl`) + const temporary = durableWriteTempPath(destination) + try { + await writeTempFileDurable( + temporary, + `${entries.map((entry) => JSON.stringify(entry)).join('\n')}\n`, + 0o600 + ) + await renameDurable(temporary, destination) + } finally { + await rm(temporary, { force: true }) + } + } + } + }) + return { ...fork, sessionId: result.sessionId, resumeCursor, materialized: true } +} diff --git a/src/main/claude/claude-stream-json-connection.ts b/src/main/claude/claude-stream-json-connection.ts index dc446b3b1b3..6f26bba45a2 100644 --- a/src/main/claude/claude-stream-json-connection.ts +++ b/src/main/claude/claude-stream-json-connection.ts @@ -39,7 +39,7 @@ export const CLAUDE_READER_DRAIN_AFTER_EXIT_MS = 2_000 let claudeAgentSdk: Promise | null = null -function loadClaudeAgentSdk(): Promise { +export function loadClaudeAgentSdk(): Promise { claudeAgentSdk ??= import('@anthropic-ai/claude-agent-sdk') return claudeAgentSdk } diff --git a/src/main/claude/claude-structured-real-cli-rewind.test.ts b/src/main/claude/claude-structured-real-cli-rewind.test.ts index f593b418b52..6fc48cc7edb 100644 --- a/src/main/claude/claude-structured-real-cli-rewind.test.ts +++ b/src/main/claude/claude-structured-real-cli-rewind.test.ts @@ -15,7 +15,7 @@ import { realClaudeCommand, realClaudeLaunchHome } from './claude-real-cli-availability-test-support' -import { readClaudePromptParentUuid } from './claude-rewind-fork' +import { readClaudePromptParentUuid, resolveClaudeRewindFork } from './claude-rewind-fork' import { CLAUDE_STRUCTURED_BASE_OPTIONS, type ClaudeStructuredSdkOptions @@ -48,7 +48,6 @@ describe.skipIf(!realClaudeAvailable)('Claude structured real CLI rewind', () => 'forks the conversation at the rewound prompt and forgets the undone turn', async () => { const source = randomUUID() - const fork = randomUUID() const { claudeConfigDir, env } = realClaudeLaunchHome() const cwd = await mkdtemp(join(tmpdir(), 'orca-real-rewind-')) const results: string[] = [] @@ -111,14 +110,28 @@ describe.skipIf(!realClaudeAvailable)('Claude structured real CLI rewind', () => claudeConfigDir }) - const rewound = adapterFor(fork, { - sessionId: fork, - resume: source, - resumeSessionAt: throughEntryId ?? '', - forkSession: true + const fork = await resolveClaudeRewindFork({ + orcaSessionId: SESSION_ID, + rewindOperationId: randomUUID(), + head: claudeProviderHandle(source, null), + prompt: { sessionId: source, uuid: promptUuid ?? '' }, + claudeConfigDirs: [claudeConfigDir] }) + expect(fork).toMatchObject({ materialized: true, throughEntryId }) + if (!fork) { + throw new Error('rewind fork was not created') + } + // A rollback resumes the persisted head without knowing about Orca's pending-fork field. + const beforeSend = await transcript(fork.sessionId, claudeConfigDir) + expect(beforeSend).toContain(throughEntryId) + expect(beforeSend).not.toContain(promptUuid) + const rewound = adapterFor(fork.sessionId, { resume: fork.sessionId }) try { - await rewound.acquire({ identity: identity(fork), fence: 1, spawnToken: 'rewind-fork' }) + await rewound.acquire({ + identity: identity(fork.sessionId), + fence: 1, + spawnToken: 'rewind-fork' + }) await claudeStartupSettled(rewound, SESSION_ID) await send( rewound, @@ -132,9 +145,9 @@ describe.skipIf(!realClaudeAvailable)('Claude structured real CLI rewind', () => const answer = results[2] ?? '' expect(answer).toContain('ALPHA') expect(answer).not.toContain('BRAVO') - // The fork keeps the kept entries' ids, and the rewound conversation is untouched. - const forked = await transcript(fork, claudeConfigDir) - expect(forked).toContain(`"uuid":"${throughEntryId}"`) + // Claude remaps the kept entries' ids and records their provenance; the source is untouched. + const forked = await transcript(fork.sessionId, claudeConfigDir) + expect(forked).toContain(`"messageUuid":"${throughEntryId}"`) expect(forked).not.toContain(promptUuid) expect(await transcript(source, claudeConfigDir)).toBe(sourceBefore) }, diff --git a/src/main/native-chat/agent-session-journal/journal-epoch-controller.ts b/src/main/native-chat/agent-session-journal/journal-epoch-controller.ts index 2648515c2b7..2fb07dd53e1 100644 --- a/src/main/native-chat/agent-session-journal/journal-epoch-controller.ts +++ b/src/main/native-chat/agent-session-journal/journal-epoch-controller.ts @@ -11,6 +11,7 @@ import type { JournalLoad } from './journal-open' import type { AgentJournalEpochReason } from './journal-row-schema' import { assertJournalFence, assertJournalWritable } from './journal-write-guards' import type { JournalWriteBody } from './journal-write-queue' +import type { JournalOperationReceipt } from './journal-row-writer' export class JournalEpochController { constructor( @@ -60,7 +61,8 @@ export class JournalEpochController { replace( reason: AgentJournalEpochReason, fence: number, - items: readonly JournalReplacementItem[] + items: readonly JournalReplacementItem[], + receipt?: (cursor: AgentJournalCursor) => JournalOperationReceipt ): Promise { return this.deps.serialize(() => { assertJournalWritable(this.deps.readOnly(), this.deps.identity.sessionId) @@ -71,6 +73,7 @@ export class JournalEpochController { reason, fence, items, + receipt, previous: this.deps.state(), queuePause: this.deps.queuePauseRestatement(), now: this.deps.now, diff --git a/src/main/native-chat/agent-session-journal/journal-epoch-replacement.ts b/src/main/native-chat/agent-session-journal/journal-epoch-replacement.ts index 841e9968886..c5fb4644254 100644 --- a/src/main/native-chat/agent-session-journal/journal-epoch-replacement.ts +++ b/src/main/native-chat/agent-session-journal/journal-epoch-replacement.ts @@ -7,6 +7,7 @@ import { isQueuedAgentJournalSubmission } from '../../../shared/agent-session-qu import { agentJournalSubmissionKey } from '../../../shared/agent-session-journal-item-key' import type { + AgentJournalCursor, AgentJournalItemBody, AgentJournalItemIdentity, AgentJournalProducerLinkage, @@ -41,6 +42,7 @@ import { } from './journal-row-table' import type { AgentJournalEpochReason, JournalRow } from './journal-row-schema' import { assertJournalFence } from './journal-write-guards' +import type { JournalOperationReceipt } from './journal-row-writer' export type JournalReplacementItem = AgentJournalProducerLinkage & { identity: AgentJournalItemIdentity @@ -62,6 +64,7 @@ export function replaceJournalEpoch(input: { reason: AgentJournalEpochReason fence: number items: readonly JournalReplacementItem[] + receipt?: (cursor: AgentJournalCursor) => JournalOperationReceipt previous?: JournalReducerState /** Restated in the new epoch, or the rewind would release cards the person stopped, or bring * back a /clear pause they already lifted. */ @@ -145,6 +148,7 @@ export function replaceJournalEpoch(input: { } const { sessionId } = input.identity + const receipt = input.receipt?.({ epoch, sequence: state.lastSequence }) input.database.transaction((db) => { const retired = readJournalSessionEpoch(db, sessionId) if (retired !== null) { @@ -154,6 +158,7 @@ export function replaceJournalEpoch(input: { insertJournalRow(db, sessionId, row) } publishJournalSessionEpoch(db, input.identity, epoch) + receipt?.write(db) }) // COMMIT landed: on disk the superseded rows are gone and this epoch is the @@ -161,4 +166,5 @@ export function replaceJournalEpoch(input: { // live store writing into an epoch whose rows were just deleted. state.oldestSequence = 1 input.onPublished({ state, newer: null, damage: null }) + receipt?.committed() } diff --git a/src/main/native-chat/agent-session-journal/journal-store.ts b/src/main/native-chat/agent-session-journal/journal-store.ts index c169b71f908..eafbe4cac34 100644 --- a/src/main/native-chat/agent-session-journal/journal-store.ts +++ b/src/main/native-chat/agent-session-journal/journal-store.ts @@ -22,7 +22,6 @@ import { newestStructuredAgentSessionTurnBySequence, runningStructuredAgentSessionTurnItemBySequence } from '../../../shared/structured-agent-session-live-turn' -import type { JournalReplacementItem } from './journal-epoch-replacement' import { readJournalSince } from './journal-cursor' import type { JournalHostDatabase } from './journal-host-database' import { journalRowsAfterReader, type JournalLoad } from './journal-open' @@ -396,11 +395,9 @@ export class AgentSessionJournal { } replaceEpochItems( - reason: AgentJournalEpochReason, - fence: number, - items: readonly JournalReplacementItem[] + ...args: Parameters ): Promise { - return this.epochController.replace(reason, fence, items) + return this.epochController.replace(...args) } private adoptLoadedJournal(loaded: JournalLoad): void { diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-adapter.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-adapter.ts index 22fb7752ea5..3c1d810c9ce 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-adapter.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-adapter.ts @@ -305,9 +305,8 @@ export type StructuredAgentSessionAdapter = StructuredAgentSessionAdapterStop & fence: number beforeTurnId: string }): Promise<{ ok: true } | { ok: false; reason: AgentSessionRewindReason }> - /** An agent that rewinds by forking: the session a rewind to before `beforeMessage` continues the - * chat in, which its next launch creates. Read before anything is saved, so a refusal changes - * nothing, and the agent need not be running. */ + /** Saves the fork a rewind continues in, before the journal or record changes. + * The agent need not be running; a refusal leaves its conversation unchanged. */ forkForRewind?(input: { record: AgentSessionRecord beforeMessage: { sessionId: string; uuid: string } diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-fork-rewind.test.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-fork-rewind.test.ts index 2c2bd9a2798..29c9e6557b3 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-fork-rewind.test.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-fork-rewind.test.ts @@ -1,4 +1,4 @@ -// A rewind by an agent that forks (Claude) saves the fork for its next launch; no provider call. +// A rewind commits its durable fork's head with the shortened journal, including at rest. import { mkdtemp, rm } from 'node:fs/promises' import { tmpdir } from 'node:os' @@ -8,6 +8,7 @@ import { agentJournalItemKey } from '../../../shared/agent-session-journal-item- import { AGENT_JOURNAL_THREAD_SCOPE } from '../../../shared/agent-session-journal-types' import { computeAgentSessionPayloadFingerprint } from '../../../shared/agent-session-mutation-envelope' import { claudeProviderHandle } from '../../../shared/agent-session-provider-handle-encoding' +import { agentSessionProviderHandleChainHead } from '../../../shared/agent-session-provider-handle' import type { AgentSessionRewindFork } from '../../../shared/agent-session-rewind' import type { AgentSessionRecordStore } from '../../runtime/agent-session-record-store' import { openTestAgentSessionRecordStore } from '../../runtime/agent-session-record-store-test-harness' @@ -26,13 +27,16 @@ import { } from './structured-agent-session-host-test-data' import { startsWhenPublished } from './structured-agent-session-instant-start.test-support' import { createStructuredAgentSessionLogger } from './structured-agent-session-logger' +import { recoverStructuredRewind } from './structured-rewind-recovery' const caller = { callerKey: 'desktop' } const SOURCE = '11111111-1111-4111-8111-111111111111' const FORK: AgentSessionRewindFork = { sessionId: '22222222-2222-4222-8222-222222222222', sourceSessionId: SOURCE, - replacesSessionId: SOURCE + replacesSessionId: SOURCE, + materialized: true, + resumeCursor: 'fork-answer-1' } const key = (uuid: string) => agentJournalItemKey({ provider: 'claude', sessionId: SOURCE, uuid }) let directory: string @@ -56,6 +60,19 @@ const attachParams = (fence: number | null) => beforeEach(async () => { resetHostTestOperationIds() failAcquire = false + forkForRewind.mockReset() + closeSession.mockClear() + forkForRewind.mockImplementation(async ({ record, beforeMessage }) => ({ + ok: true, + fork: { + ...FORK, + sessionId: + beforeMessage.uuid === 'prompt-2' ? FORK.sessionId : '33333333-3333-4333-8333-333333333333', + ...(beforeMessage.uuid === 'prompt-1' ? { resumeCursor: undefined } : {}), + replacesSessionId: agentSessionProviderHandleChainHead(record.providerHandleChain)!.handle + .nativeId + } + })) directory = await mkdtemp(join(tmpdir(), 'orca-fork-rewind-')) store = await openTestAgentSessionRecordStore(directory) const adapter = startsWhenPublished( @@ -124,12 +141,33 @@ beforeEach(async () => { }) afterEach(async () => { + vi.restoreAllMocks() await host.flushAllStreamedEvents() await rm(directory, { recursive: true, force: true }) }) +async function rewindSecondPrompt() { + const fields = { + itemId: key('prompt-2'), + expectedEpoch: (await host.journalSnapshot(HOST_TEST_SESSION)).cursor.epoch + } + return host.rewind(caller, { + ...fields, + envelope: { + sessionId: HOST_TEST_SESSION, + clientOperationId: hostTestOperationId(), + expectedRuntimeFence: store.getRecord(HOST_TEST_SESSION)!.lease.runtimeFence, + payloadFingerprint: computeAgentSessionPayloadFingerprint({ + method: 'agentSession.rewind', + sessionId: HOST_TEST_SESSION, + fields + }) + } + }) +} + describe('rewinding a chat whose agent forks', () => { - it('saves the fork, keeps the prefix and rewinds again at rest without starting the agent', async () => { + it('commits the fork head with the prefix and rewinds again at rest without starting the agent', async () => { const epoch = (await host.journalSnapshot(HOST_TEST_SESSION)).cursor.epoch const fields = { itemId: key('prompt-2'), expectedEpoch: epoch } const envelope = { @@ -151,7 +189,14 @@ describe('rewinding a chat whose agent forks', () => { expect(closeSession).toHaveBeenCalled() expect(store.getRecord(HOST_TEST_SESSION)?.rewind).toMatchObject({ phase: 'completed', - providerFork: FORK + providerFork: { materialized: true } + }) + const record = store.getRecord(HOST_TEST_SESSION)! + const fork = record.rewind!.providerFork! + // These are the fields v1.4.233 reads; it ignores providerFork entirely. + expect(agentSessionProviderHandleChainHead(record.providerHandleChain)).toMatchObject({ + origin: 'forked', + handle: claudeProviderHandle(fork.sessionId, 'fork-answer-1') }) const snapshot = await host.journalSnapshot(HOST_TEST_SESSION) expect(snapshot.items.map((item) => item.itemId)).toEqual([key('prompt-1'), key('answer-1')]) @@ -174,5 +219,61 @@ describe('rewinding a chat whose agent forks', () => { }) ).toMatchObject({ ok: true }) expect((await host.journalSnapshot(HOST_TEST_SESSION)).items).toEqual([]) + expect( + agentSessionProviderHandleChainHead(store.getRecord(HOST_TEST_SESSION)!.providerHandleChain) + ?.handle + ).toEqual(claudeProviderHandle('33333333-3333-4333-8333-333333333333', null)) + }) + + it('keeps the original conversation when saving its fork fails', async () => { + const before = await host.journalSnapshot(HOST_TEST_SESSION) + forkForRewind.mockRejectedValueOnce(new Error('transcript write failed')) + await expect(rewindSecondPrompt()).rejects.toThrow('transcript write failed') + expect(await host.journalSnapshot(HOST_TEST_SESSION)).toEqual(before) + expect(store.getRecord(HOST_TEST_SESSION)?.rewind).toBeUndefined() + expect(closeSession).not.toHaveBeenCalled() + }) + + it('rolls back both history and handle when the receipt fails, then recovers the same fork', async () => { + const journal = host.collaboratorsForTests().sessions.get(HOST_TEST_SESSION)!.journal + const before = await host.journalSnapshot(HOST_TEST_SESSION) + const makeReceipt = store.conversationReceipts.rewind.bind(store.conversationReceipts) + const failed = vi + .spyOn(store.conversationReceipts, 'rewind') + .mockImplementationOnce((...args) => { + const receipt = makeReceipt(...args) + return { + ...receipt, + write: (db) => { + receipt.write(db) + throw new Error('commit failed') + } + } + }) + await expect(rewindSecondPrompt()).rejects.toThrow('commit failed') + expect(await host.journalSnapshot(HOST_TEST_SESSION)).toEqual(before) + expect( + agentSessionProviderHandleChainHead(store.getRecord(HOST_TEST_SESSION)!.providerHandleChain) + ?.handle.nativeId + ).toBe(SOURCE) + expect(store.getRecord(HOST_TEST_SESSION)?.rewind?.phase).toBe('provider-succeeded') + failed.mockRestore() + + await recoverStructuredRewind( + { store, logger: createStructuredAgentSessionLogger() }, + HOST_TEST_SESSION, + journal, + store.getRecord(HOST_TEST_SESSION)!.lease.runtimeFence + ) + expect(journal.snapshot().items.map((item) => item.itemId)).toEqual([ + key('prompt-1'), + key('answer-1') + ]) + expect( + agentSessionProviderHandleChainHead(store.getRecord(HOST_TEST_SESSION)!.providerHandleChain) + ?.handle.nativeId + ).toBe(FORK.sessionId) + expect(store.getRecord(HOST_TEST_SESSION)?.rewind?.phase).toBe('completed') + expect(forkForRewind).toHaveBeenCalledTimes(1) }) }) diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-rewind.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-rewind.ts index be63ea529bb..d71098c9c84 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-rewind.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-rewind.ts @@ -41,8 +41,7 @@ export async function rewindStructuredAgentSession( } = context const fields = { itemId: params.itemId, expectedEpoch: params.expectedEpoch } const operation = { method: 'agentSession.rewind', fields } - // A fork is made by the agent's next launch, so it rewinds at rest; any other rewind is the - // running provider's own. + // A fork can be saved without a running agent; other rewinds require the provider itself. const agent = store.getRecord(sessionId)?.provider const rewindsByFork = agent !== undefined && diff --git a/src/main/native-chat/agent-session-wire/structured-rewind-recovery.ts b/src/main/native-chat/agent-session-wire/structured-rewind-recovery.ts index 76abea25624..ac704cd2926 100644 --- a/src/main/native-chat/agent-session-wire/structured-rewind-recovery.ts +++ b/src/main/native-chat/agent-session-wire/structured-rewind-recovery.ts @@ -191,7 +191,7 @@ export async function applyStructuredRewindRecovery( throw new Error('agent_session_rewind:invalid-target') } if (rewind.phase === 'prepared' && rewind.providerFork) { - // The fork is the next launch's to make, so nothing outside Orca can have half-happened. + // New forks are durable already; older records leave the copy to their next launch. rewind = { ...rewind, phase: 'provider-succeeded' } await persistRewindRecord(store, sessionId, fence, rewind) } else if (rewind.phase === 'prepared') { @@ -224,6 +224,18 @@ export async function applyStructuredRewindRecovery( if (rewind.phase !== 'provider-succeeded') { return } + if (rewind.providerFork?.materialized && !rewind.contextClearOperationId && !alreadyReplaced) { + const completed = rewind + await journal.replaceEpochItems( + 'handle_forked', + fence, + cut!.kept, + (cursor) => + receipt?.(completed, fence, cursor) ?? + store.conversationReceipts.rewind(sessionId, fence, completed, cursor) + ) + return + } if (rewind.contextClearOperationId && rewind.contextClearSequence) { if ( store.getRecord(sessionId)?.providerContextBoundary?.operationId !== diff --git a/src/main/runtime/agent-session-rewind-completion.ts b/src/main/runtime/agent-session-rewind-completion.ts index 51adc4cbe13..e90dbd5aef6 100644 --- a/src/main/runtime/agent-session-rewind-completion.ts +++ b/src/main/runtime/agent-session-rewind-completion.ts @@ -2,6 +2,40 @@ import type { AgentJournalCursor } from '../../shared/agent-session-journal-type import type { AgentSessionRewindRecord } from '../../shared/agent-session-rewind' import type { AgentSessionStoreState } from './agent-session-store-state' import { settleAgentSessionOperationInto } from './agent-session-operation-admission' +import { + agentSessionProviderHandleChainHead, + agentSessionProviderHandleKey, + appendAgentSessionProviderHandleLink +} from '../../shared/agent-session-provider-handle' +import type { AgentSessionRecord } from '../../shared/agent-session-record' + +function rewindProviderChain( + record: AgentSessionRecord, + rewind: AgentSessionRewindRecord, + fence: number +) { + const fork = rewind.providerFork + const head = agentSessionProviderHandleChainHead(record.providerHandleChain)?.handle + if (!fork?.materialized || head?.nativeId === fork.sessionId) { + return record.providerHandleChain + } + if (!head || head.nativeId !== fork.replacesSessionId) { + throw new Error('agent_session_rewind:stale-provider') + } + return appendAgentSessionProviderHandleLink(record.providerHandleChain, { + linkId: `rewind-${fork.sessionId}`, + handle: { + transport: head.transport, + agent: head.agent, + nativeId: fork.sessionId, + ...(fork.resumeCursor ? { resumeCursor: fork.resumeCursor } : {}) + }, + origin: 'forked', + mintedAtFence: fence, + observedAt: Date.now(), + forkedFromKey: agentSessionProviderHandleKey(head) + }) +} export function commitAgentSessionRewindCompletion( draft: AgentSessionStoreState, @@ -14,18 +48,34 @@ export function commitAgentSessionRewindCompletion( if ( !record || record.lease.runtimeFence !== fence || - record.providerContextBoundary?.operationId !== rewind.contextClearOperationId + record.rewind?.operationId !== rewind.operationId || + (rewind.contextClearOperationId !== undefined && + record.providerContextBoundary?.operationId !== rewind.contextClearOperationId) ) { throw new Error('agent_session_rewind:stale-context') } const { retained: _retained, hydrationVerified: _hydrationVerified, ...current } = rewind draft.records.set(sessionId, { ...record, - rewind: { ...current, phase: 'completed', epoch: cursor.epoch, sequence: cursor.sequence } + providerHandleChain: rewindProviderChain(record, rewind, fence), + rewind: { + ...current, + phase: 'completed', + epoch: cursor.epoch, + ...(rewind.contextClearOperationId ? { sequence: cursor.sequence } : {}) + } }) settleAgentSessionOperationInto(draft, { callerKey: rewind.callerKey, operationId: rewind.operationId, - outcome: { status: 'succeeded', sessionId, rewind: { itemId: rewind.itemId, ...cursor } } + outcome: { + status: 'succeeded', + sessionId, + rewind: { + itemId: rewind.itemId, + epoch: cursor.epoch, + ...(rewind.contextClearOperationId ? { sequence: cursor.sequence } : {}) + } + } }) } diff --git a/src/shared/agent-session-rewind.ts b/src/shared/agent-session-rewind.ts index 96ff3623c32..9838f54bf1a 100644 --- a/src/shared/agent-session-rewind.ts +++ b/src/shared/agent-session-rewind.ts @@ -49,12 +49,14 @@ export const AgentSessionRewindRecordSchema = z.object({ /** Written by older builds only, as is `retained`: rewind now derives both from the journal. */ hydrationVerified: z.boolean().optional(), // A provider that rewinds by forking: the session the chat continues in, copied from `sourceSessionId` - // through `throughEntryId` (absent: nothing kept). Its next launch performs the fork. + // through `throughEntryId` (absent: nothing kept). Older records defer the copy to launch. providerFork: z .object({ sessionId: Key, sourceSessionId: Key, throughEntryId: Key.optional(), + materialized: z.boolean().optional(), + resumeCursor: Key.optional(), // The chat's session when it was rewound; a later head means the fork no longer applies. replacesSessionId: Key })