diff --git a/src/main/git/checkpoints/worktree-checkpoint-git.ts b/src/main/git/checkpoints/worktree-checkpoint-git.ts new file mode 100644 index 00000000000..973abd6550c --- /dev/null +++ b/src/main/git/checkpoints/worktree-checkpoint-git.ts @@ -0,0 +1,228 @@ +// Git plumbing for worktree checkpoints. Every command here reads, or writes only objects, refs +// under refs/orca/checkpoints/ and private temp index files: the user's HEAD, branch, index and +// stash list are never touched by a snapshot. + +import { randomBytes } from 'node:crypto' +import type { Stats } from 'node:fs' +import { copyFile, readdir, rm, stat, utimes } from 'node:fs/promises' +import { join } from 'node:path' +import { gitExecFileAsync } from '../command-runner/git-exec-file' + +export const CHECKPOINT_REF_ROOT = 'refs/orca/checkpoints/' + +/** Identity for checkpoint commits, so repos without user.name/user.email still snapshot. */ +const CHECKPOINT_IDENTITY = { + GIT_AUTHOR_NAME: 'Orca', + GIT_AUTHOR_EMAIL: 'checkpoints@orca.invalid', + GIT_COMMITTER_NAME: 'Orca', + GIT_COMMITTER_EMAIL: 'checkpoints@orca.invalid' +} + +export type CheckpointGitOptions = { + env?: Record + stdin?: string + signal?: AbortSignal + timeoutMs?: number +} + +export type CheckpointGit = (args: string[], options?: CheckpointGitOptions) => Promise + +export type CheckpointRepo = { + /** The worktree root; every path below is relative to it. */ + root: string + /** The worktree's own git dir (a linked worktree has its own index and temp files there). */ + gitDir: string + /** The user's index file. Copied, never written, by a snapshot. */ + indexFile: string + git: CheckpointGit +} + +export class WorktreeCheckpointError extends Error { + constructor( + readonly reason: + | 'not-a-repo' + | 'not-found' + | 'merge-conflicts' + | 'blocked-path' + | 'unsafe-path' + | 'snapshot-failed' + /** Files changed before the restore failed; `undoRef` puts them back. */ + | 'partial', + message: string, + readonly paths: string[] = [], + readonly undoRef?: string + ) { + super(message) + this.name = 'WorktreeCheckpointError' + } +} + +export function gitExitCode(error: unknown): number | null { + const code = error && typeof error === 'object' && 'code' in error ? error.code : null + return typeof code === 'number' ? code : null +} + +export function defaultCheckpointGit(cwd: string): CheckpointGit { + return async (args, options = {}) => { + const { stdout } = await gitExecFileAsync(args, { + cwd, + // Plumbing only: no optional locks, so a snapshot never contends with the user's own git. + env: { + ...process.env, + GIT_OPTIONAL_LOCKS: '0', + ...CHECKPOINT_IDENTITY, + ...options.env + }, + ...(options.stdin !== undefined ? { stdin: options.stdin } : {}), + ...(options.signal ? { signal: options.signal } : {}), + timeout: options.timeoutMs ?? 120_000, + maxBuffer: 256 * 1024 * 1024 + }) + return stdout + } +} + +/** Resolves the worktree that contains `directory`; throws `not-a-repo` for plain folders and bare repos. */ +export async function openCheckpointRepo( + directory: string, + gitFor: (cwd: string) => CheckpointGit = defaultCheckpointGit +): Promise { + let out: string + try { + out = await gitFor(directory)([ + 'rev-parse', + '--path-format=absolute', + '--show-toplevel', + '--git-dir', + '--git-path', + 'index' + ]) + } catch (error) { + throw new WorktreeCheckpointError( + 'not-a-repo', + `${directory} is not inside a git worktree: ${error instanceof Error ? error.message : String(error)}` + ) + } + const [root, gitDir, indexFile] = out.split('\n').map((line) => line.trim()) + if (!root || !gitDir || !indexFile) { + throw new WorktreeCheckpointError('not-a-repo', `${directory} has no worktree root.`) + } + return { root, gitDir, indexFile, git: gitFor(root) } +} + +export async function resolveCommit(repo: CheckpointRepo, rev: string): Promise { + try { + const out = await repo.git(['rev-parse', '-q', '--verify', `${rev}^{commit}`]) + return out.trim() || null + } catch { + return null + } +} + +const TEMP_INDEX_PREFIX = 'orca-checkpoint-' +const STALE_TEMP_INDEX_MS = 60 * 60 * 1000 +const swept = new Set() + +/** Removes temp indexes a crashed process left behind (once per git dir per process). */ +async function sweepStaleTempIndexes(gitDir: string): Promise { + if (swept.has(gitDir)) { + return + } + swept.add(gitDir) + const names = await readdir(gitDir).catch(() => []) + for (const name of names) { + if (!name.startsWith(TEMP_INDEX_PREFIX)) { + continue + } + const path = join(gitDir, name) + const info = await stat(path).catch(() => null) + // ctime, not mtime: a live temp index carries the user index's (possibly old) mtime. + if (info && Date.now() - Math.max(info.ctimeMs, info.mtimeMs) > STALE_TEMP_INDEX_MS) { + await rm(path, { force: true }) + } + } +} + +/** Runs `run` against a private copy of the user's index (stat data included, so `add -A` stays + * incremental on large repos). The copy is deleted afterwards; the user's index is only read. */ +export async function withTempIndex( + repo: CheckpointRepo, + seedFromUserIndex: boolean, + run: (env: Record) => Promise +): Promise { + await sweepStaleTempIndexes(repo.gitDir) + const file = `${repo.gitDir}/${TEMP_INDEX_PREFIX}${process.pid}-${randomBytes(6).toString('hex')}.index` + try { + if (seedFromUserIndex) { + await seedTempIndex(repo.indexFile, file) + } + return await run({ GIT_INDEX_FILE: file }) + } finally { + await rm(file, { force: true }) + await rm(`${file}.lock`, { force: true }) + } +} + +/** Copies the user's index keeping its mtime. Git treats an entry as "racily clean" (and rehashes + * the file) only when the entry is not older than the index file itself; a copy stamped "now" + * would make a same-size edit made in the index's last second look unchanged and drop it from + * the snapshot. The mtime is read before the copy, so a concurrent rewrite only makes the copy + * look older, which means more rehashing, never less. */ +async function seedTempIndex(indexFile: string, file: string): Promise { + let source: Stats + try { + source = await stat(indexFile) + } catch (error) { + // An unborn repo has no index yet; the empty temp index is the right seed. + if ((error as NodeJS.ErrnoException).code === 'ENOENT') { + return + } + throw error + } + await copyFile(indexFile, file) + await utimes(file, source.atime, source.mtime) +} + +export type TreeChange = { + status: 'A' | 'D' | 'M' | 'T' + path: string +} + +/** File-level differences from `fromTree` to `toTree`, renames split into delete + add. + * Submodule pointers are left out: a checkpoint never moves a submodule. */ +export async function diffTrees( + repo: CheckpointRepo, + fromTree: string, + toTree: string +): Promise { + if (fromTree === toTree) { + return [] + } + const out = await repo.git(['diff-tree', '-r', '-z', '--no-renames', '--raw', fromTree, toTree]) + const fields = out.split('\0') + const changes: TreeChange[] = [] + for (let index = 0; index + 1 < fields.length; index += 2) { + const header = fields[index] + const path = fields[index + 1] + if (!header.startsWith(':') || path === undefined) { + continue + } + const [oldMode, newMode, , , status] = header.slice(1).split(' ') + if (oldMode === '160000' || newMode === '160000') { + continue + } + const kind = status?.charAt(0) + if (kind === 'A' || kind === 'D' || kind === 'M' || kind === 'T') { + changes.push({ status: kind, path }) + } + } + return changes +} + +/** Paths git would refuse anyway; checked again so a bad tree can never reach outside the root. */ +export function isSafeRelativePath(path: string): boolean { + if (!path || path.startsWith('/') || path.includes('\0') || /^[A-Za-z]:/.test(path)) { + return false + } + return path.split(/[\\/]/).every((part) => part !== '' && part !== '..' && part !== '.git') +} diff --git a/src/main/git/checkpoints/worktree-checkpoint-real-git.test.ts b/src/main/git/checkpoints/worktree-checkpoint-real-git.test.ts new file mode 100644 index 00000000000..fec43add553 --- /dev/null +++ b/src/main/git/checkpoints/worktree-checkpoint-real-git.test.ts @@ -0,0 +1,414 @@ +import { chmod, mkdir, mkdtemp, readFile, rm, stat, symlink, writeFile } from 'node:fs/promises' +import { existsSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { runProcess } from '@orca/process-host' +import { afterEach, beforeEach, describe, expect, it } from 'vitest' +import { + openCheckpointRepo, + type CheckpointRepo, + WorktreeCheckpointError +} from './worktree-checkpoint-git' +import { + captureWorktreeCheckpoint, + CHECKPOINT_MAX_UNTRACKED_BYTES, + readWorktreeCheckpoint +} from './worktree-checkpoint-snapshot' +import { previewWorktreeRestore, restoreWorktreeCheckpoint } from './worktree-checkpoint-restore' +import { assertUndoCovers } from './worktree-checkpoint-restore-guards' +import { listCheckpointRefs, pruneCheckpointRefs } from './worktree-checkpoint-retention' + +let dir: string +let repo: CheckpointRepo + +async function git(...args: string[]): Promise { + const result = await runProcess({ program: 'git', args, cwd: dir }) + if (result.code !== 0) { + throw new Error(result.stderr) + } + return result.stdout.trim() +} +const put = async (path: string, text: string | Buffer) => { + await mkdir(join(dir, path, '..'), { recursive: true }) + await writeFile(join(dir, path), text) +} +const read = (path: string) => readFile(join(dir, path), 'utf8') +const T1 = 'refs/orca/checkpoints/s1/turn/t1' +const NOW = 1_800_000_000_000 +const capture = (ref = T1) => captureWorktreeCheckpoint(repo, ref, { kind: 'turn', createdAt: NOW }) +const restore = async (ref = T1) => { + const target = await readWorktreeCheckpoint(repo, ref) + return restoreWorktreeCheckpoint(repo, target!, { + ref: 'refs/orca/checkpoints/s1/undo/u1', + now: 2 + }) +} +async function userState() { + return { + head: await git('rev-parse', 'HEAD'), + branch: await git('branch', '--show-current'), + staged: await git('diff', '--cached', '--name-status'), + stash: await git('stash', 'list'), + status: await git('status', '--porcelain=v1', '-uall') + } +} + +beforeEach(async () => { + dir = await mkdtemp(join(tmpdir(), 'orca-checkpoint-')) + await git('init', '-q') + await git('symbolic-ref', 'HEAD', 'refs/heads/main') + await git('config', 'user.name', 'Audit') + await git('config', 'user.email', 'audit@example.test') + await git('config', 'commit.gpgsign', 'true') + await git('config', 'core.hooksPath', '/nonexistent') + await put('.gitignore', 'ignored/\n*.log\n') + await put('a.txt', 'a1\n') + await put('keep/b.txt', 'b1\n') + await git('add', '.') + await git('-c', 'commit.gpgsign=false', 'commit', '-qm', 'base') + repo = await openCheckpointRepo(dir) +}) +afterEach(async () => { + await rm(dir, { recursive: true, force: true }) +}) + +describe('worktree checkpoints', () => { + it('snapshots without touching HEAD, branch, index, stash or status', async () => { + await put('a.txt', 'a2\n') + await git('add', 'a.txt') + await put('a.txt', 'a3\n') + await put('new.txt', 'untracked\n') + const before = await userState() + const checkpoint = await capture() + expect(await userState()).toEqual(before) + expect(checkpoint.head).toBe(before.head) + expect(await git('show', `${checkpoint.commit}:a.txt`)).toBe('a3') + expect(await git('show', `${checkpoint.commit}:new.txt`)).toBe('untracked') + expect(await git('show', `${checkpoint.indexTree}:a.txt`)).toBe('a2') + expect(await git('for-each-ref', '--format=%(refname)')).toContain(T1) + expect(existsSync(join(dir, '.git', 'index.lock'))).toBe(false) + }) + + it('sees a same-size edit made in the same second the index was written (racy git)', async () => { + // Git rehashes an entry whose stat matches only when it is not older than the index file. A + // temp index stamped "now" (a later second) would trust the stale stat and drop the edit. + // Racy git is about real file-system seconds: start early in one so the edit shares it. + // ast-grep-ignore: test-no-polling-loop, test-no-wall-clock -- the file system clock is real + while (Date.now() % 1000 > 300) { + // ast-grep-ignore: test-no-real-time-wait -- the file system clock can't be faked + await new Promise((resolve) => setTimeout(resolve, 20)) + } + await put('a.txt', 'r1\n') + await git('add', 'a.txt') + await put('a.txt', 's1\n') + // ast-grep-ignore: test-no-real-time-wait -- the snapshot must run in a later second + await new Promise((resolve) => setTimeout(resolve, 1100)) + const checkpoint = await capture() + expect(await git('show', `${checkpoint.commit}:a.txt`)).toBe('s1') + }) + + it('restores untracked, deleted, binary, renamed and modified files', async () => { + const binary = Buffer.from([0, 1, 2, 255, 254, 0, 10]) + await put('bin.dat', binary) + await put('untracked.txt', 'u1\n') + await capture() + await put('a.txt', 'changed\n') + await rm(join(dir, 'untracked.txt')) + await put('bin.dat', Buffer.from([9, 9])) + await mkdir(join(dir, 'moved')) + await git('mv', 'keep/b.txt', 'moved/b.txt') + await put('agent-new/deep.txt', 'agent\n') + const preview = await previewWorktreeRestore(repo, (await readWorktreeCheckpoint(repo, T1))!) + expect(preview.changes.map((c) => `${c.status} ${c.path}`).sort()).toEqual([ + 'A keep/b.txt', + 'A untracked.txt', + 'D agent-new/deep.txt', + 'D moved/b.txt', + 'M a.txt', + 'M bin.dat' + ]) + const result = await restore() + expect(result.changes).toHaveLength(6) + expect(await read('a.txt')).toBe('a1\n') + expect(await read('untracked.txt')).toBe('u1\n') + expect(await readFile(join(dir, 'bin.dat'))).toEqual(binary) + expect(await read('keep/b.txt')).toBe('b1\n') + expect(existsSync(join(dir, 'moved'))).toBe(false) + expect(existsSync(join(dir, 'agent-new'))).toBe(false) + // The staged rename is gone from the index too: HEAD never moved, so staging comes back. + expect(result.indexRestored).toBe(true) + expect(await git('diff', '--cached', '--name-status')).toBe('') + expect(await git('status', '--porcelain=v1')).toBe('?? bin.dat\n?? untracked.txt'.trim()) + expect( + (await previewWorktreeRestore(repo, (await readWorktreeCheckpoint(repo, T1))!)).changes + ).toEqual([]) + }) + + it('leaves gitignored files alone', async () => { + await put('ignored/cache.bin', 'before\n') + await put('debug.log', 'log1\n') + await capture() + await put('ignored/cache.bin', 'after\n') + await put('ignored/new.bin', 'new\n') + await put('debug.log', 'log2\n') + await put('a.txt', 'changed\n') + await restore() + expect(await read('ignored/cache.bin')).toBe('after\n') + expect(await read('ignored/new.bin')).toBe('new\n') + expect(await read('debug.log')).toBe('log2\n') + expect(await read('a.txt')).toBe('a1\n') + }) + + it('brings back a dirty index while HEAD is unchanged', async () => { + await put('a.txt', 'staged\n') + await git('add', 'a.txt') + await put('a.txt', 'staged-then-edited\n') + await capture() + await git('reset', '-q') + await put('a.txt', 'agent\n') + const result = await restore() + expect(result.indexRestored).toBe(true) + expect(await git('show', ':a.txt')).toBe('staged') + expect(await read('a.txt')).toBe('staged-then-edited\n') + }) + + it('keeps commits made after the checkpoint and leaves the index when HEAD moved', async () => { + await capture() + const base = await git('rev-parse', 'HEAD') + await put('a.txt', 'committed\n') + await git('-c', 'commit.gpgsign=false', 'commit', '-qam', 'agent commit') + const agentHead = await git('rev-parse', 'HEAD') + const result = await restore() + expect(result.headMoved).toBe(true) + expect(result.indexRestored).toBe(false) + expect(await git('rev-parse', 'HEAD')).toBe(agentHead) + expect(agentHead).not.toBe(base) + expect(await read('a.txt')).toBe('a1\n') + }) + + it('undoes a restore, including an ignored file the restore overwrote', async () => { + await put('later-ignored.txt', 'v1\n') + await capture() + await put('.gitignore', 'ignored/\n*.log\nlater-ignored.txt\n') + await put('later-ignored.txt', 'precious\n') + await put('a.txt', 'agent\n') + const result = await restore() + expect(await read('later-ignored.txt')).toBe('v1\n') + const undo = await readWorktreeCheckpoint(repo, result.undo.ref) + await restoreWorktreeCheckpoint(repo, undo!, { + ref: 'refs/orca/checkpoints/s1/undo/u2', + now: 3 + }) + expect(await read('later-ignored.txt')).toBe('precious\n') + expect(await read('a.txt')).toBe('agent\n') + expect(await read('.gitignore')).toContain('later-ignored.txt') + }) + + it('restores file modes and symlinks', async () => { + await put('run.sh', '#!/bin/sh\n') + await chmod(join(dir, 'run.sh'), 0o755) + await symlink('a.txt', join(dir, 'link')) + await capture() + await chmod(join(dir, 'run.sh'), 0o644) + await rm(join(dir, 'link')) + await put('link', 'now a file\n') + await restore() + expect((await stat(join(dir, 'run.sh'))).mode & 0o111).not.toBe(0) + expect(await read('link')).toBe('a1\n') + }) + + it('refuses while merge conflicts are unresolved, before writing anything', async () => { + await capture() + await git('checkout', '-qb', 'other') + await put('a.txt', 'other\n') + await git('-c', 'commit.gpgsign=false', 'commit', '-qam', 'other') + await git('checkout', '-q', 'main') + await put('a.txt', 'main\n') + await git('-c', 'commit.gpgsign=false', 'commit', '-qam', 'main') + await runProcess({ program: 'git', args: ['merge', 'other'], cwd: dir }) + const conflicted = await read('a.txt') + await expect(restore()).rejects.toMatchObject({ + reason: 'merge-conflicts' + }) + expect(await read('a.txt')).toBe(conflicted) + expect(await git('for-each-ref', 'refs/orca/checkpoints/s1/undo/')).toBe('') + }) + + it('refuses when a folder now sits where the checkpoint has a file', async () => { + await capture() + await rm(join(dir, 'a.txt')) + await put('a.txt/inside.log', 'ignored content\n') + await expect(restore()).rejects.toMatchObject({ reason: 'blocked-path' }) + expect(await read('a.txt/inside.log')).toBe('ignored content\n') + }) + + it('refuses when an ignored file now sits where the checkpoint has a folder', async () => { + await put('cache/x.txt', 'tracked then\n') + await capture() + await rm(join(dir, 'cache'), { recursive: true }) + await put('.gitignore', 'ignored/\n*.log\ncache\n') + await put('cache', 'precious ignored file\n') + await expect(restore()).rejects.toMatchObject({ reason: 'blocked-path' }) + expect(await read('cache')).toBe('precious ignored file\n') + }) + + it('snapshots past an unreadable file and refuses to overwrite it', async () => { + await put('secret.txt', 'v1\n') + await capture() + await put('secret.txt', 'unreadable now\n') + await chmod(join(dir, 'secret.txt'), 0o000) + await put('a.txt', 'agent\n') + try { + const later = await capture('refs/orca/checkpoints/s1/turn/t2') + expect(await git('show', `${later.commit}:a.txt`)).toBe('agent') + await expect(restore()).rejects.toMatchObject({ reason: 'blocked-path' }) + expect(await read('a.txt')).toBe('agent\n') + } finally { + await chmod(join(dir, 'secret.txt'), 0o644) + } + expect(await read('secret.txt')).toBe('unreadable now\n') + }) + + it('leaves large untracked files out of snapshots and alone on restore', async () => { + await capture() + await put('big.bin', Buffer.alloc(CHECKPOINT_MAX_UNTRACKED_BYTES + 1, 7)) + await put('a.txt', 'agent\n') + const later = await capture('refs/orca/checkpoints/s1/turn/t2') + await expect(git('cat-file', '-e', `${later.commit}:big.bin`)).rejects.toThrow() + await restore() + expect((await stat(join(dir, 'big.bin'))).size).toBe(CHECKPOINT_MAX_UNTRACKED_BYTES + 1) + expect(await read('a.txt')).toBe('a1\n') + }) + + it('leaves skip-worktree and assume-unchanged files alone', async () => { + await capture() + await git('update-index', '--skip-worktree', 'a.txt') + await git('update-index', '--assume-unchanged', 'keep/b.txt') + await put('a.txt', 'local only\n') + await put('keep/b.txt', 'local too\n') + await put('new.txt', 'agent\n') + const result = await restore() + expect(result.changes.map((c) => c.path)).toEqual(['new.txt']) + expect(await read('a.txt')).toBe('local only\n') + expect(await read('keep/b.txt')).toBe('local too\n') + expect(await git('ls-files', '-v', 'a.txt')).toBe('S a.txt') + }) + + it('names its undo point when a restore fails partway', async () => { + await put('locked/keep.txt', 'k\n') + await put('locked/gone.txt', 'g\n') + await capture() + await rm(join(dir, 'locked/gone.txt')) + await put('extra.txt', 'agent\n') + await chmod(join(dir, 'locked'), 0o555) + try { + const error = await restore().then( + () => null, + (failure: unknown) => failure + ) + expect(error).toBeInstanceOf(WorktreeCheckpointError) + const undoRef = error instanceof WorktreeCheckpointError ? error.undoRef : undefined + expect(error).toMatchObject({ reason: 'partial' }) + expect(existsSync(join(dir, 'extra.txt'))).toBe(false) + await chmod(join(dir, 'locked'), 0o755) + const undo = await readWorktreeCheckpoint(repo, undoRef ?? '') + await restoreWorktreeCheckpoint(repo, undo!, { + ref: 'refs/orca/checkpoints/s1/undo/u9', + now: NOW + }) + expect(await read('extra.txt')).toBe('agent\n') + } finally { + await chmod(join(dir, 'locked'), 0o755) + } + }) + + it('never writes into a nested repository the agent created', async () => { + await put('sub/a.txt', 'checkpointed\n') + await capture() + await git('-C', join(dir, 'sub'), 'init', '-q') + await put('sub/a.txt', 'uncommitted work in the nested repo\n') + await put('a.txt', 'agent\n') + await expect(restore()).rejects.toMatchObject({ reason: 'blocked-path' }) + expect(await read('sub/a.txt')).toBe('uncommitted work in the nested repo\n') + expect(await read('a.txt')).toBe('agent\n') + }) + + it('refuses when a path on disk is missing from the undo tree', async () => { + await put('present.txt', 'on disk\n') + const empty = (await git('hash-object', '-t', 'tree', '-w', '--stdin')).trim() + await expect( + assertUndoCovers(repo, empty, [{ status: 'M', path: 'present.txt' }]) + ).rejects.toMatchObject({ reason: 'blocked-path', paths: ['present.txt'] }) + await expect( + assertUndoCovers(repo, empty, [{ status: 'A', path: 'absent.txt' }]) + ).resolves.toBeUndefined() + // A saved file next to an unsaved one whose name holds a newline: answers must not shift. + await put('x', 'saved\n') + await put('x\ny', 'not saved\n') + await git('add', 'x') + const tree = await git('write-tree') + await expect( + assertUndoCovers(repo, tree, [ + { status: 'M', path: 'x\ny' }, + { status: 'M', path: 'x' } + ]) + ).rejects.toMatchObject({ reason: 'blocked-path', paths: ['x\ny'] }) + }) + + it('works on an unborn branch', async () => { + await rm(join(dir, '.git'), { recursive: true }) + await git('init', '-q') + repo = await openCheckpointRepo(dir) + await capture() + await put('a.txt', 'agent\n') + await put('fresh.txt', 'agent\n') + await restore() + expect(await read('a.txt')).toBe('a1\n') + expect(existsSync(join(dir, 'fresh.txt'))).toBe(false) + }) + + it('prunes beyond the newest N and anything past the max age', async () => { + for (let index = 0; index < 5; index += 1) { + await capture(`refs/orca/checkpoints/s1/turn/t${index}`) + } + await capture('refs/orca/checkpoints/s2/turn/x') + const deleted = await pruneCheckpointRefs(repo, { + prefix: 'refs/orca/checkpoints/s1/turn/', + keep: 3, + maxAgeMs: 60_000, + now: NOW + }) + expect(deleted).toHaveLength(2) + expect(await listCheckpointRefs(repo, 'refs/orca/checkpoints/s1/turn/')).toHaveLength(3) + expect(await listCheckpointRefs(repo, 'refs/orca/checkpoints/s2/')).toHaveLength(1) + await pruneCheckpointRefs(repo, { + prefix: 'refs/orca/checkpoints/s1/turn/', + keep: 50, + maxAgeMs: -60_000, + now: NOW + }) + expect(await listCheckpointRefs(repo, 'refs/orca/checkpoints/')).toEqual([]) + await capture('refs/orca/checkpoints/s1/turn/fresh') + await pruneCheckpointRefs(repo, { + prefix: 'refs/orca/checkpoints/s1/turn/', + keep: 0, + maxAgeMs: -60_000, + now: NOW, + protect: ['refs/orca/checkpoints/s1/turn/fresh'] + }) + expect(await listCheckpointRefs(repo, 'refs/orca/checkpoints/')).toEqual([ + 'refs/orca/checkpoints/s1/turn/fresh' + ]) + }) + + it('rejects plain folders', async () => { + const plain = await mkdtemp(join(tmpdir(), 'orca-checkpoint-plain-')) + try { + await expect(openCheckpointRepo(plain)).rejects.toMatchObject({ + reason: 'not-a-repo' + }) + } finally { + await rm(plain, { recursive: true, force: true }) + } + }) +}) diff --git a/src/main/git/checkpoints/worktree-checkpoint-restore-guards.ts b/src/main/git/checkpoints/worktree-checkpoint-restore-guards.ts new file mode 100644 index 00000000000..f5b5985343c --- /dev/null +++ b/src/main/git/checkpoints/worktree-checkpoint-restore-guards.ts @@ -0,0 +1,191 @@ +// Checks a restore runs before it writes anything. Each one refuses (or narrows the restore) where +// writing would lose a file nothing else has a copy of. + +import { constants } from 'node:fs' +import { access, lstat } from 'node:fs/promises' +import { join } from 'node:path' +import { + type CheckpointRepo, + diffTrees, + isSafeRelativePath, + type TreeChange, + WorktreeCheckpointError +} from './worktree-checkpoint-git' +import type { WorktreeCheckpoint } from './worktree-checkpoint-snapshot' + +export async function assertNoConflicts(repo: CheckpointRepo): Promise { + const unmerged = await repo.git(['ls-files', '-u', '-z']) + if (unmerged.length > 0) { + throw new WorktreeCheckpointError( + 'merge-conflicts', + 'The worktree has unresolved merge conflicts; resolve or abort them first.' + ) + } +} + +/** Entries git keeps off the disk (sparse checkout, skip-worktree) or doesn't look at + * (assume-unchanged). A snapshot holds their index version, not the file, so restore leaves them. */ +export async function hiddenIndexPaths(repo: CheckpointRepo): Promise> { + const out = await repo.git(['ls-files', '-v', '-z']) + const hidden = new Set() + for (const entry of out.split('\0')) { + const tag = entry.charAt(0) + if (entry.length > 2 && (tag === 'S' || (tag >= 'a' && tag <= 'z'))) { + hidden.add(entry.slice(2)) + } + } + return hidden +} + +export async function restoreChanges( + repo: CheckpointRepo, + fromTree: string, + toTree: string, + hidden: Set +): Promise { + return (await diffTrees(repo, fromTree, toTree)).filter((change) => !hidden.has(change.path)) +} + +/** Paths the checkpoint writes that exist on disk but not in the working tree's snapshot (ignored, + * or too large to snapshot): the undo checkpoint must hold them before they are overwritten. */ +export function overwrittenUnsnapshottedPaths(repo: CheckpointRepo, target: WorktreeCheckpoint) { + return async (tree: string): Promise => { + const added = (await diffTrees(repo, tree, target.tree)).filter((c) => c.status === 'A') + const present: string[] = [] + for (const change of added) { + const absolute = join(repo.root, change.path) + // An unreadable file can't be saved; assertWritable refuses to overwrite it instead. + const stat = isSafeRelativePath(change.path) ? await lstat(absolute).catch(() => null) : null + if ( + stat && + (stat.isSymbolicLink() || + (await access(absolute, constants.R_OK).then( + () => true, + () => false + ))) + ) { + present.push(change.path) + } + } + return present + } +} + +async function hasGitDir( + repo: CheckpointRepo, + folder: string, + cache: Map +): Promise { + let known = cache.get(folder) + if (known === undefined) { + known = (await lstat(join(repo.root, folder, '.git')).catch(() => null)) !== null + cache.set(folder, known) + } + return known +} + +/** Refuses before any write when a path to write is a folder or an unreadable file now, when a + * leading folder of it is now a file the restore doesn't delete itself (an ignored file would be + * replaced with no copy), or when any path is unsafe. */ +export async function assertWritable(repo: CheckpointRepo, changes: TreeChange[]): Promise { + const unsafe = changes.filter((change) => !isSafeRelativePath(change.path)) + if (unsafe.length > 0) { + throw new WorktreeCheckpointError( + 'unsafe-path', + 'The checkpoint names paths outside the worktree.', + unsafe.map((change) => change.path) + ) + } + const deleted = new Set(changes.filter((c) => c.status === 'D').map((c) => c.path)) + const folders = new Map() + const isNonFolder = async (path: string): Promise => { + let known = folders.get(path) + if (known === undefined) { + const stat = await lstat(join(repo.root, path)).catch(() => null) + known = stat !== null && !stat.isDirectory() + folders.set(path, known) + } + return known + } + const gitDirs = new Map() + const blocked: string[] = [] + for (const change of changes) { + const parts = change.path.split('/') + // A nested repository owns its files; the restore never writes or deletes inside one. + if (change.status === 'D') { + for (let depth = 1; depth < parts.length; depth += 1) { + if (await hasGitDir(repo, parts.slice(0, depth).join('/'), gitDirs)) { + blocked.push(change.path) + break + } + } + continue + } + const absolute = join(repo.root, change.path) + const stat = await lstat(absolute).catch(() => null) + const unreadable = + stat?.isFile() === true && + (await access(absolute, constants.R_OK).then( + () => false, + () => true + )) + let leadingFile = false + for (let depth = 1; depth < parts.length && !leadingFile; depth += 1) { + const prefix = parts.slice(0, depth).join('/') + leadingFile = !deleted.has(prefix) && (await isNonFolder(prefix)) + } + let nestedRepo = false + for (let depth = 1; depth < parts.length && !nestedRepo; depth += 1) { + const prefix = parts.slice(0, depth).join('/') + nestedRepo = await hasGitDir(repo, prefix, gitDirs) + } + if (stat?.isDirectory() || unreadable || leadingFile || nestedRepo) { + blocked.push(change.path) + } + } + if (blocked.length > 0) { + throw new WorktreeCheckpointError( + 'blocked-path', + 'Something now sits where the checkpoint writes a file.', + blocked + ) + } +} + +/** Known gap: a tracked file that is unreadable now and that the restore deletes is covered by its + * stale index blob, so the undo point holds its last staged content, not what is on disk. + * + * Every path the restore writes or deletes that exists on disk must be in the undo tree, or the + * restore could destroy work nothing saved. `git add` can skip a path without failing (a nested + * repo with no commits, a file it couldn't read), so this checks the tree itself. */ +export async function assertUndoCovers( + repo: CheckpointRepo, + undoTree: string, + changes: TreeChange[] +): Promise { + const present: string[] = [] + for (const change of changes) { + if (await lstat(join(repo.root, change.path)).catch(() => null)) { + present.push(change.path) + } + } + if (present.length === 0) { + return + } + const lined = present.filter((path) => !path.includes('\n')) + const out = await repo.git(['cat-file', '--batch-check=%(objecttype)'], { + stdin: lined.map((path) => `${undoTree}:${path}\n`).join('') + }) + const answers = out.split('\n') + // batch-check is line-based: a path with a newline would shift every answer after it, so such a + // path is never asked about and counts as unsaved. + const saved = new Set(lined.filter((_, index) => answers[index]?.endsWith('blob'))) + const missing = present.filter((path) => !saved.has(path)) + if (missing.length > 0) { + throw new WorktreeCheckpointError( + 'blocked-path', + 'Some files the restore would change could not be saved first.', + missing + ) + } +} diff --git a/src/main/git/checkpoints/worktree-checkpoint-restore.ts b/src/main/git/checkpoints/worktree-checkpoint-restore.ts new file mode 100644 index 00000000000..78bb60ca115 --- /dev/null +++ b/src/main/git/checkpoints/worktree-checkpoint-restore.ts @@ -0,0 +1,181 @@ +// Puts the working tree back to a checkpoint. Before anything is written, every file the restore +// would change or delete is saved in an "undo" checkpoint (ignored files it would overwrite too), +// so the restore itself can always be undone. HEAD and the branch never move: commits made after +// the checkpoint stay on the branch. The staged state comes back only while HEAD is unchanged. + +import { realpath, rm, rmdir } from 'node:fs/promises' +import { dirname, join, relative, sep } from 'node:path' +import { + type CheckpointRepo, + resolveCommit, + type TreeChange, + withTempIndex, + WorktreeCheckpointError +} from './worktree-checkpoint-git' +import { + assertNoConflicts, + assertUndoCovers, + assertWritable, + hiddenIndexPaths, + overwrittenUnsnapshottedPaths, + restoreChanges +} from './worktree-checkpoint-restore-guards' +import { + captureWorktreeTree, + type WorktreeCheckpoint, + writeWorktreeCheckpoint +} from './worktree-checkpoint-snapshot' + +export type WorktreeRestorePreview = { + /** Files that differ between the working tree now and the checkpoint. */ + changes: TreeChange[] + /** HEAD moved since the checkpoint (for example the agent committed); those commits stay. */ + headMoved: boolean +} + +export type WorktreeRestoreResult = WorktreeRestorePreview & { + /** The checkpoint of the working tree as it was just before this restore. */ + undo: WorktreeCheckpoint + indexRestored: boolean +} + +export async function previewWorktreeRestore( + repo: CheckpointRepo, + target: WorktreeCheckpoint +): Promise { + const now = await captureWorktreeTree(repo) + return { + changes: await restoreChanges(repo, now.tree, target.tree, await hiddenIndexPaths(repo)), + headMoved: now.head !== target.head + } +} + +/** Deletes a file the checkpoint doesn't have, refusing to follow a symlinked parent folder. */ +async function removeWorktreeFile(rootReal: string, root: string, path: string): Promise { + const absolute = join(root, path) + const parentReal = await realpath(dirname(absolute)).catch(() => null) + if (!parentReal) { + return + } + const fromRoot = relative(rootReal, parentReal) + if (fromRoot.startsWith('..') || fromRoot.startsWith(sep)) { + throw new WorktreeCheckpointError('unsafe-path', `${path} resolves outside the worktree.`, [ + path + ]) + } + await rm(absolute, { force: true }) + // Drop folders the deletion left empty, never the root and never a folder with anything left. + let folder = dirname(absolute) + while (folder.length > root.length && folder.startsWith(root)) { + try { + await rmdir(folder) + } catch { + break + } + folder = dirname(folder) + } +} + +async function restoreIndexEntries( + repo: CheckpointRepo, + fromTree: string, + toTree: string, + hidden: Set +): Promise { + if (fromTree === toTree) { + return + } + const out = await repo.git(['diff-tree', '-r', '-z', '--no-renames', '--raw', fromTree, toTree]) + const fields = out.split('\0') + const records: string[] = [] + for (let index = 0; index + 1 < fields.length; index += 2) { + const header = fields[index] + if (!header.startsWith(':')) { + continue + } + const [, newMode, , newObject] = header.slice(1).split(' ') + const path = fields[index + 1] + // Rewriting an entry drops its skip-worktree/assume-unchanged bit; those stay as they are. + if (path === undefined || hidden.has(path)) { + continue + } + // A path the staged tree lacks arrives as mode 000000 with an all-zero id: --index-info removes it. + records.push(`${newMode === '000000' ? '0' : newMode} ${newObject}\t${path}`) + } + if (records.length > 0) { + await repo.git(['update-index', '-z', '--index-info'], { + stdin: `${records.join('\0')}\0` + }) + } +} + +async function applyChanges( + repo: CheckpointRepo, + target: WorktreeCheckpoint, + changes: TreeChange[] +): Promise { + const rootReal = await realpath(repo.root) + // Deletes first, so a file can take the place of a folder (or a case-only rename) it replaces. + for (const change of changes) { + if (change.status === 'D') { + await removeWorktreeFile(rootReal, repo.root, change.path) + } + } + const writes = changes.filter((change) => change.status !== 'D').map((change) => change.path) + if (writes.length > 0) { + await withTempIndex(repo, false, async (env) => { + await repo.git(['read-tree', target.tree], { env }) + // checkout-index refuses symlinked leading paths itself; -f replaces what is there. + await repo.git(['checkout-index', '-f', '-z', '--stdin'], { + env, + stdin: `${writes.join('\0')}\0` + }) + }) + } +} + +export async function restoreWorktreeCheckpoint( + repo: CheckpointRepo, + target: WorktreeCheckpoint, + undo: { ref: string; labels?: Record; now: number } +): Promise { + await assertNoConflicts(repo) + const hidden = await hiddenIndexPaths(repo) + // Saving the undo point first is what makes the restore safe: a failure below leaves it to restore. + const before = await captureWorktreeTree(repo, { + forceInclude: overwrittenUnsnapshottedPaths(repo, target) + }) + const changes = await restoreChanges(repo, before.tree, target.tree, hidden) + await assertWritable(repo, changes) + await assertUndoCovers(repo, before.tree, changes) + const undoCheckpoint = await writeWorktreeCheckpoint(repo, undo.ref, before, { + kind: 'undo', + createdAt: undo.now, + ...(undo.labels ? { labels: undo.labels } : {}) + }) + try { + await applyChanges(repo, target, changes) + } catch (error) { + throw new WorktreeCheckpointError( + 'partial', + `Some files were not restored: ${error instanceof Error ? error.message : String(error)}`, + [], + undoCheckpoint.ref + ) + } + const head = await resolveCommit(repo, 'HEAD') + const headMoved = head !== target.head + let indexRestored = false + if (!headMoved && target.indexTree && before.indexTree) { + // Same commit as at the checkpoint: put back exactly what was staged then. Only entries that + // differ are rewritten, so the rest keep their stat data and the next status stays fast. Files + // are already restored, so a busy index (another git holding index.lock) only loses staging. + try { + await restoreIndexEntries(repo, before.indexTree, target.indexTree, hidden) + indexRestored = true + } catch { + indexRestored = false + } + } + return { changes, headMoved, undo: undoCheckpoint, indexRestored } +} diff --git a/src/main/git/checkpoints/worktree-checkpoint-retention.ts b/src/main/git/checkpoints/worktree-checkpoint-retention.ts new file mode 100644 index 00000000000..a1b66e8ceb5 --- /dev/null +++ b/src/main/git/checkpoints/worktree-checkpoint-retention.ts @@ -0,0 +1,67 @@ +// Keeps checkpoint refs bounded: the newest N under one prefix, and nothing older than a max age +// anywhere under refs/orca/checkpoints/. Deleting a ref only drops its reachability; git gc frees +// the objects later on its own schedule. + +import { type CheckpointRepo, CHECKPOINT_REF_ROOT } from './worktree-checkpoint-git' + +type RefAge = { ref: string; committedAt: number } + +async function listRefs(repo: CheckpointRepo, prefix: string): Promise { + const out = await repo.git([ + 'for-each-ref', + '--sort=-committerdate', + '--format=%(committerdate:unix) %(refname)', + prefix + ]) + return out + .split('\n') + .filter(Boolean) + .map((line) => { + const space = line.indexOf(' ') + return { + committedAt: Number(line.slice(0, space)) * 1000, + ref: line.slice(space + 1) + } + }) +} + +export async function listCheckpointRefs(repo: CheckpointRepo, prefix: string): Promise { + return (await listRefs(repo, prefix)).map((entry) => entry.ref) +} + +export async function deleteCheckpointRefs(repo: CheckpointRepo, refs: string[]): Promise { + const owned = refs.filter((ref) => ref.startsWith(CHECKPOINT_REF_ROOT)) + if (owned.length === 0) { + return + } + await repo.git(['update-ref', '--stdin'], { + stdin: owned.map((ref) => `delete ${ref}\n`).join('') + }) +} + +/** Deletes all but the newest `keep` refs under `prefix`, plus any checkpoint older than `maxAgeMs`. */ +export async function pruneCheckpointRefs( + repo: CheckpointRepo, + input: { + prefix: string + keep: number + maxAgeMs: number + now: number + /** Never deleted here, whatever its date says (a clock that ran ahead must not prune the + * checkpoint just written). */ + protect?: string[] + } +): Promise { + if (!input.prefix.startsWith(CHECKPOINT_REF_ROOT)) { + throw new Error(`checkpoint prefixes live under ${CHECKPOINT_REF_ROOT}: ${input.prefix}`) + } + const overCount = (await listRefs(repo, input.prefix)).slice(input.keep).map((entry) => entry.ref) + const cutoff = input.now - input.maxAgeMs + const tooOld = (await listRefs(repo, CHECKPOINT_REF_ROOT)) + .filter((entry) => entry.committedAt < cutoff) + .map((entry) => entry.ref) + const kept = new Set(input.protect ?? []) + const doomed = [...new Set([...overCount, ...tooOld])].filter((ref) => !kept.has(ref)) + await deleteCheckpointRefs(repo, doomed) + return doomed +} diff --git a/src/main/git/checkpoints/worktree-checkpoint-snapshot.ts b/src/main/git/checkpoints/worktree-checkpoint-snapshot.ts new file mode 100644 index 00000000000..2baabad10bf --- /dev/null +++ b/src/main/git/checkpoints/worktree-checkpoint-snapshot.ts @@ -0,0 +1,237 @@ +// A checkpoint is a stash-like commit kept under refs/orca/checkpoints/: its tree is the working +// tree (tracked, plus untracked files .gitignore doesn't exclude), its first parent is HEAD and its +// second parent holds the staged tree. Writing one never touches HEAD, the branch, the user's index +// or the stash list. + +import { z } from 'zod' +import { lstat } from 'node:fs/promises' +import { join } from 'node:path' +import { + type CheckpointGitOptions, + type CheckpointRepo, + CHECKPOINT_REF_ROOT, + gitExitCode, + resolveCommit, + withTempIndex, + WorktreeCheckpointError +} from './worktree-checkpoint-git' + +const META_PREFIX = 'orca-checkpoint: ' + +export type WorktreeCheckpointKind = 'turn' | 'undo' + +export type WorktreeCheckpointMeta = { + v: 1 + kind: WorktreeCheckpointKind + /** HEAD when the checkpoint was taken; null on an unborn branch. */ + head: string | null + /** The staged tree; null while the index had unmerged entries. */ + indexTree: string | null + createdAt: number + /** Caller's own labels (session and turn ids), opaque to git. */ + labels?: Record +} + +export type WorktreeTree = { + tree: string + head: string | null + indexTree: string | null +} + +export type WorktreeCheckpoint = WorktreeCheckpointMeta & { + ref: string + commit: string + tree: string +} + +export type CaptureOptions = { + signal?: AbortSignal + timeoutMs?: number + /** Given the working tree's tree, returns ignored paths that must be saved too (restore uses it + * to keep a copy of every file it is about to overwrite). */ + forceInclude?: (tree: string) => Promise +} + +/** Untracked files above this are left out of snapshots (like ignored files): hashing them every + * turn would bloat the object store. A restore that would overwrite one still saves it first. */ +export const CHECKPOINT_MAX_UNTRACKED_BYTES = 10 * 1024 * 1024 + +async function largeUntrackedPaths( + repo: CheckpointRepo, + options: CheckpointGitOptions, + deadline: number | null +): Promise { + const out = await repo.git(['ls-files', '-o', '--exclude-standard', '-z'], options) + const large: string[] = [] + for (const path of out.split('\0')) { + if (deadline !== null && Date.now() > deadline) { + throw new WorktreeCheckpointError('snapshot-failed', 'The snapshot ran out of time.') + } + const stat = path ? await lstat(join(repo.root, path)).catch(() => null) : null + if (stat?.isFile() && stat.size > CHECKPOINT_MAX_UNTRACKED_BYTES) { + large.push(path) + } + } + return large +} + +/** `add --ignore-errors` exits 1 after skipping files it couldn't read; everything else is in. */ +async function addSkippingUnreadable( + repo: CheckpointRepo, + args: string[], + options: CheckpointGitOptions +): Promise { + try { + await repo.git(['add', '--ignore-errors', ...args], options) + } catch (error) { + if (gitExitCode(error) !== 1) { + throw error + } + } +} + +/** Writes the working tree's tree into the object store without recording a ref. */ +export async function captureWorktreeTree( + repo: CheckpointRepo, + options: CaptureOptions = {} +): Promise { + // One deadline for the whole capture, not one per git command. + const deadline = options.timeoutMs === undefined ? null : Date.now() + options.timeoutMs + const run = (env?: Record): CheckpointGitOptions => ({ + ...(options.signal ? { signal: options.signal } : {}), + ...(deadline === null ? {} : { timeoutMs: Math.max(1, deadline - Date.now()) }), + ...(env ? { env } : {}) + }) + const head = await resolveCommit(repo, 'HEAD') + const large = await largeUntrackedPaths(repo, run(), deadline) + return withTempIndex(repo, true, async (env) => { + let indexTree: string | null = null + try { + // The temp index is a byte copy of the user's, so this is exactly what they had staged. + indexTree = (await repo.git(['write-tree'], run(env))).trim() || null + } catch { + indexTree = null + } + await addSkippingUnreadable(repo, ['-A', '--pathspec-from-file=-', '--pathspec-file-nul'], { + ...run(env), + stdin: `${[':/', ...large.map((path) => `:(exclude,literal)${path}`)].join('\0')}\0` + }) + let tree = (await repo.git(['write-tree'], run(env))).trim() + const forced = options.forceInclude ? await options.forceInclude(tree) : [] + if (forced.length > 0) { + await repo.git(['add', '-f', '--pathspec-from-file=-', '--pathspec-file-nul'], { + ...run(env), + stdin: `${forced.map(literalPathspec).join('\0')}\0` + }) + tree = (await repo.git(['write-tree'], run(env))).trim() + } + if (!tree) { + throw new WorktreeCheckpointError('snapshot-failed', 'git write-tree returned no tree.') + } + return { tree, head, indexTree } + }) +} + +function literalPathspec(path: string): string { + return `:(literal)${path}` +} + +export async function writeWorktreeCheckpoint( + repo: CheckpointRepo, + ref: string, + snapshot: WorktreeTree, + meta: Omit +): Promise { + if (!ref.startsWith(CHECKPOINT_REF_ROOT)) { + throw new Error(`checkpoint refs live under ${CHECKPOINT_REF_ROOT}: ${ref}`) + } + const full: WorktreeCheckpointMeta = { + v: 1, + kind: meta.kind, + head: snapshot.head, + indexTree: snapshot.indexTree, + createdAt: meta.createdAt, + ...(meta.labels ? { labels: meta.labels } : {}) + } + const headParent = snapshot.head ? ['-p', snapshot.head] : [] + // Commit dates follow the caller's clock, so retention and ordering read one timeline. + const stamp = `@${Math.floor(meta.createdAt / 1000)} +0000` + const env = { GIT_AUTHOR_DATE: stamp, GIT_COMMITTER_DATE: stamp } + // The staged tree rides as the second parent so gc keeps it as long as the checkpoint. + const indexCommit = snapshot.indexTree + ? ( + await repo.git( + [ + 'commit-tree', + '--no-gpg-sign', + snapshot.indexTree, + ...headParent, + '-m', + 'Orca checkpoint index' + ], + { env } + ) + ).trim() + : null + const commit = ( + await repo.git( + [ + 'commit-tree', + '--no-gpg-sign', + snapshot.tree, + ...headParent, + ...(indexCommit ? ['-p', indexCommit] : []), + '-m', + `Orca checkpoint (${meta.kind})\n\n${META_PREFIX}${JSON.stringify(full)}` + ], + { env } + ) + ).trim() + await repo.git(['update-ref', ref, commit]) + return { ...full, ref, commit, tree: snapshot.tree } +} + +export async function captureWorktreeCheckpoint( + repo: CheckpointRepo, + ref: string, + meta: Omit, + options: CaptureOptions = {} +): Promise { + return writeWorktreeCheckpoint(repo, ref, await captureWorktreeTree(repo, options), meta) +} + +export async function readWorktreeCheckpoint( + repo: CheckpointRepo, + ref: string +): Promise { + if (!ref.startsWith(CHECKPOINT_REF_ROOT)) { + return null + } + const commit = await resolveCommit(repo, ref) + if (!commit) { + return null + } + const out = await repo.git(['show', '-s', '--format=%T%x00%B', commit]) + const [tree, body = ''] = out.split('\0') + const line = body.split('\n').find((candidate) => candidate.startsWith(META_PREFIX)) + const meta = line ? parseMeta(line.slice(META_PREFIX.length)) : null + return meta ? { ...meta, ref, commit, tree: tree.trim() } : null +} + +const MetaSchema = z.object({ + v: z.literal(1), + kind: z.enum(['turn', 'undo']), + head: z.string().nullable(), + indexTree: z.string().nullable(), + createdAt: z.number(), + labels: z.record(z.string(), z.string()).optional() +}) + +function parseMeta(text: string): WorktreeCheckpointMeta | null { + try { + const parsed = MetaSchema.safeParse(JSON.parse(text)) + return parsed.success ? parsed.data : null + } catch { + return null + } +} diff --git a/src/main/native-chat/agent-session-wire/checkpoints/agent-session-checkpoint-queue.ts b/src/main/native-chat/agent-session-wire/checkpoints/agent-session-checkpoint-queue.ts new file mode 100644 index 00000000000..cc589390d5d --- /dev/null +++ b/src/main/native-chat/agent-session-wire/checkpoints/agent-session-checkpoint-queue.ts @@ -0,0 +1,34 @@ +// Per-worktree ordering shared by every chat's checkpoint work in this process. Process-local by +// design: the host that captures a chat's turns is the one that restores them (one main process +// or orcad per machine owns a chat). + +// One git writer per worktree at a time across every chat in this process: a capture never +// interleaves with a restore, so a send after a rewind snapshots the restored files. +const queues = new Map>() +export function serialize(root: string, task: () => Promise): Promise { + const run = (queues.get(root) ?? Promise.resolve()).then(task, task) + const tail = run.catch(() => undefined) + queues.set(root, tail) + void tail.then(() => queues.get(root) === tail && queues.delete(root)) + return run +} + +// Worktrees with a turn between its snapshot and its hand-off to the agent: the chat isn't +// "working" yet, but a restore now would rewrite files the agent is about to see. +export const startingTurns = new Map() +export function markTurnStarting(root: string): () => void { + startingTurns.set(root, (startingTurns.get(root) ?? 0) + 1) + let released = false + return () => { + if (released) { + return + } + released = true + const left = (startingTurns.get(root) ?? 1) - 1 + if (left > 0) { + startingTurns.set(root, left) + } else { + startingTurns.delete(root) + } + } +} diff --git a/src/main/native-chat/agent-session-wire/checkpoints/agent-session-checkpoints-host.ts b/src/main/native-chat/agent-session-wire/checkpoints/agent-session-checkpoints-host.ts new file mode 100644 index 00000000000..064c3668c05 --- /dev/null +++ b/src/main/native-chat/agent-session-wire/checkpoints/agent-session-checkpoints-host.ts @@ -0,0 +1,60 @@ +// Binds checkpoints to the host: the capture a send takes, and the host's preview/restore/undo. + +import type { StructuredAgentSessionDeliveryLoopDeps } from '../delivery/structured-agent-session-delivery-loop-deps' +import type { StructuredAgentSessionHostDeps } from '../host/structured-agent-session-host-types' +import { + type AgentSessionCheckpoints, + createAgentSessionCheckpoints +} from './agent-session-checkpoints' + +/** The snapshot a send takes before its message opens a turn; never throws. */ +export function structuredAgentSessionTurnCheckpoint( + deps: Pick< + StructuredAgentSessionDeliveryLoopDeps, + 'record' | 'resolveWorkspacePath' | 'logger' | 'now' + >, + sessionId: string +): (itemId: string, signal?: AbortSignal) => Promise<() => void> { + const { captureTurn } = createAgentSessionCheckpoints({ + record: deps.record, + ...(deps.resolveWorkspacePath ? { resolveWorkspacePath: deps.resolveWorkspacePath } : {}), + isBusy: () => false, + logger: deps.logger, + now: deps.now + }) + return (itemId, signal) => captureTurn(sessionId, itemId, signal) +} + +type CheckpointHost = { + deps: Pick + readStatusSummary: (sessionId: string) => { status: string | null } | undefined +} +const hostCheckpoints = new WeakMap() + +/** Preview, restore and undo for a host's chats; a restore waits for a running turn to end. */ +export function hostAgentSessionCheckpoints(host: CheckpointHost): AgentSessionCheckpoints { + let checkpoints = hostCheckpoints.get(host) + if (!checkpoints) { + const { deps } = host + checkpoints = createAgentSessionCheckpoints({ + record: (sessionId) => deps.store.getRecord(sessionId), + ...(deps.resolveWorkspacePath ? { resolveWorkspacePath: deps.resolveWorkspacePath } : {}), + // Any chat working in the same worktree could be writing the files a restore rewrites. + isBusy: (sessionId) => { + const location = deps.store.getRecord(sessionId)?.location + return deps.store + .listConversations() + .some( + (record) => + record.location.workspaceId === location?.workspaceId && + record.location.executionHostId === location.executionHostId && + host.readStatusSummary(record.sessionId)?.status === 'working' + ) + }, + logger: deps.logger, + now: () => deps.now?.() ?? Date.now() + }) + hostCheckpoints.set(host, checkpoints) + } + return checkpoints +} diff --git a/src/main/native-chat/agent-session-wire/checkpoints/agent-session-checkpoints.ts b/src/main/native-chat/agent-session-wire/checkpoints/agent-session-checkpoints.ts new file mode 100644 index 00000000000..23b48dd9af6 --- /dev/null +++ b/src/main/native-chat/agent-session-wire/checkpoints/agent-session-checkpoints.ts @@ -0,0 +1,258 @@ +// Ties worktree checkpoints to chats: one snapshot before each turn a message opens, keyed by that +// message's journal item, so "Rewind to here" can put the files back as they were when it was sent. +// Capture failures never block or fail a send; they only mean that message has no checkpoint. + +import { createHash, randomBytes } from 'node:crypto' +import type { AgentSessionRecord } from '../../../../shared/agent-session/lifecycle/agent-session-record' +import { + AGENT_SESSION_CHECKPOINT_UNDO_ID, + agentSessionCheckpointPartialMessage, + agentSessionCheckpointRefusalMessage, + type AgentSessionCheckpointPreview, + type AgentSessionCheckpointRestoreResult +} from '../../../../shared/agent-session/lifecycle/agent-session-checkpoint' +import { + CHECKPOINT_REF_ROOT, + type CheckpointRepo, + openCheckpointRepo, + WorktreeCheckpointError +} from '../../../git/checkpoints/worktree-checkpoint-git' +import { + captureWorktreeCheckpoint, + readWorktreeCheckpoint, + type WorktreeCheckpoint +} from '../../../git/checkpoints/worktree-checkpoint-snapshot' +import { + previewWorktreeRestore, + restoreWorktreeCheckpoint +} from '../../../git/checkpoints/worktree-checkpoint-restore' +import { pruneCheckpointRefs } from '../../../git/checkpoints/worktree-checkpoint-retention' +import type { StructuredAgentSessionLogger } from '../host/structured-agent-session-logger' +import { markTurnStarting, serialize, startingTurns } from './agent-session-checkpoint-queue' + +export const TURN_CHECKPOINTS_KEPT = 50 +export const UNDO_CHECKPOINTS_KEPT = 10 +export const CHECKPOINT_MAX_AGE_MS = 30 * 24 * 60 * 60 * 1000 +/** A worktree too large to snapshot in this long sends without a checkpoint. */ +export const CHECKPOINT_CAPTURE_TIMEOUT_MS = 15_000 + +export type AgentSessionCheckpointsDeps = { + record: (sessionId: string) => AgentSessionRecord | null + resolveWorkspacePath?: (workspaceId: string) => Promise + /** True while any chat in this chat's worktree has a turn running; restores refuse then. */ + isBusy: (sessionId: string) => boolean + logger: StructuredAgentSessionLogger + now: () => number + openRepo?: (directory: string) => Promise +} + +export type AgentSessionCheckpoints = ReturnType + +/** Ref-safe and stable: readable ids pass through, anything else is hashed. */ +function refComponent(value: string): string { + return /^[A-Za-z0-9_-][A-Za-z0-9._-]{0,127}$/.test(value) && !value.endsWith('.lock') + ? value + : createHash('sha256').update(value).digest('hex').slice(0, 32) +} + +function sessionPrefix(sessionId: string): string { + return `${CHECKPOINT_REF_ROOT}${refComponent(sessionId)}/` +} + +function turnRef(sessionId: string, itemId: string): string { + return `${sessionPrefix(sessionId)}turn/${createHash('sha256').update(itemId).digest('hex').slice(0, 32)}` +} + +function refuse(reason: Parameters[0]): never { + throw new Error(agentSessionCheckpointRefusalMessage(reason)) +} + +/** Something else in the workspace that may write files: a terminal agent, say. */ +/** Gets the worktree root the restore would write. */ +export type AgentSessionWorkspaceBusy = (root: string) => Promise + +export function createAgentSessionCheckpoints(deps: AgentSessionCheckpointsDeps) { + const open = deps.openRepo ?? ((directory: string) => openCheckpointRepo(directory)) + const repoFor = async (sessionId: string): Promise => { + const record = deps.record(sessionId) + // WSL worktrees route git through another filesystem; their temp index paths don't map. + if (!record || record.location.wslDistro) { + return null + } + const directory = + record.launchDirectory ?? (await deps.resolveWorkspacePath?.(record.location.workspaceId)) + return directory ? open(directory) : null + } + + const captureTurn = async ( + sessionId: string, + itemId: string, + signal?: AbortSignal + ): Promise<() => void> => { + const started = deps.now() + let release = () => {} + try { + const repo = await repoFor(sessionId) + if (!repo) { + return release + } + release = markTurnStarting(repo.root) + await serialize(repo.root, async () => { + await captureWorktreeCheckpoint( + repo, + turnRef(sessionId, itemId), + { + kind: 'turn', + createdAt: deps.now(), + labels: { sessionId, itemId } + }, + { signal, timeoutMs: CHECKPOINT_CAPTURE_TIMEOUT_MS } + ) + await pruneCheckpointRefs(repo, { + prefix: `${sessionPrefix(sessionId)}turn/`, + keep: TURN_CHECKPOINTS_KEPT, + maxAgeMs: CHECKPOINT_MAX_AGE_MS, + now: deps.now(), + protect: [turnRef(sessionId, itemId)] + }) + }) + } catch (error) { + if (!(error instanceof WorktreeCheckpointError && error.reason === 'not-a-repo')) { + deps.logger.warn('[agent-session] turn checkpoint skipped', { + scope: 'checkpoint-capture', + sessionId, + elapsedMs: deps.now() - started, + error: error instanceof Error ? error.message : String(error) + }) + } + } + return release + } + + const target = async ( + sessionId: string, + ref: string + ): Promise<{ + repo: CheckpointRepo + checkpoint: WorktreeCheckpoint + } | null> => { + const repo = await repoFor(sessionId) + const checkpoint = repo ? await readWorktreeCheckpoint(repo, ref) : null + return repo && checkpoint ? { repo, checkpoint } : null + } + + const preview = async ( + sessionId: string, + itemId: string, + workspaceBusy?: AgentSessionWorkspaceBusy + ): Promise => { + try { + const found = await target(sessionId, turnRef(sessionId, itemId)) + if (!found) { + return { available: false, reason: 'none' } + } + const { changes, headMoved } = await serialize(found.repo.root, () => + previewWorktreeRestore(found.repo, found.checkpoint) + ) + return { + available: true, + changedFiles: changes.length, + headMoved, + createdAt: found.checkpoint.createdAt, + busy: + deps.isBusy(sessionId) || + (startingTurns.get(found.repo.root) ?? 0) > 0 || + ((await workspaceBusy?.(found.repo.root).catch(() => false)) ?? false) + } + } catch (error) { + if (error instanceof WorktreeCheckpointError && error.reason === 'not-a-repo') { + return { available: false, reason: 'not-a-repo' } + } + deps.logger.warn('[agent-session] checkpoint preview failed', { + scope: 'checkpoint-preview', + sessionId, + error: error instanceof Error ? error.message : String(error) + }) + return { available: false, reason: 'unavailable' } + } + } + + const restoreRef = async ( + sessionId: string, + ref: string, + workspaceBusy?: AgentSessionWorkspaceBusy + ): Promise => { + const repo = await repoFor(sessionId).catch((error: unknown) => { + if (error instanceof WorktreeCheckpointError && error.reason === 'not-a-repo') { + refuse('not-a-repo') + } + throw error + }) + if (!repo) { + refuse('none') + } + const now = deps.now() + const undoId = `${now}-${randomBytes(4).toString('hex')}` + const undoRef = `${sessionPrefix(sessionId)}undo/${undoId}` + try { + // Look up and check busy inside the queue: a capture queued first runs first, and the turn + // it starts must make this restore refuse rather than rewrite files under a working agent. + const result = await serialize(repo.root, async () => { + if ( + deps.isBusy(sessionId) || + (startingTurns.get(repo.root) ?? 0) > 0 || + (await workspaceBusy?.(repo.root)) + ) { + refuse('busy') + } + const checkpoint = await readWorktreeCheckpoint(repo, ref) + if (!checkpoint) { + refuse('none') + } + return restoreWorktreeCheckpoint(repo, checkpoint, { + ref: undoRef, + now, + labels: { sessionId } + }) + }) + void pruneCheckpointRefs(repo, { + prefix: `${sessionPrefix(sessionId)}undo/`, + keep: UNDO_CHECKPOINTS_KEPT, + maxAgeMs: CHECKPOINT_MAX_AGE_MS, + now, + protect: [undoRef] + }).catch(() => undefined) + return { + changedFiles: result.changes.length, + headMoved: result.headMoved, + undoId + } + } catch (error) { + if (error instanceof WorktreeCheckpointError && error.reason === 'partial') { + throw new Error(agentSessionCheckpointPartialMessage(undoId)) + } + if ( + error instanceof WorktreeCheckpointError && + (error.reason === 'merge-conflicts' || + error.reason === 'blocked-path' || + error.reason === 'unsafe-path') + ) { + refuse(error.reason) + } + throw error + } + } + + return { + captureTurn, + preview, + restore: (sessionId: string, itemId: string, workspaceBusy?: AgentSessionWorkspaceBusy) => + restoreRef(sessionId, turnRef(sessionId, itemId), workspaceBusy), + undo: async (sessionId: string, undoId: string, workspaceBusy?: AgentSessionWorkspaceBusy) => { + if (!AGENT_SESSION_CHECKPOINT_UNDO_ID.test(undoId)) { + refuse('none') + } + return restoreRef(sessionId, `${sessionPrefix(sessionId)}undo/${undoId}`, workspaceBusy) + } + } +} diff --git a/src/main/native-chat/agent-session-wire/delivery/structured-agent-session-delivery-loop-deps.ts b/src/main/native-chat/agent-session-wire/delivery/structured-agent-session-delivery-loop-deps.ts index 040de3949c9..0e996777462 100644 --- a/src/main/native-chat/agent-session-wire/delivery/structured-agent-session-delivery-loop-deps.ts +++ b/src/main/native-chat/agent-session-wire/delivery/structured-agent-session-delivery-loop-deps.ts @@ -69,4 +69,6 @@ export type StructuredAgentSessionDeliveryLoopDeps = { now: () => number acquireAborts: Pick recordRanOptions?: StructuredAgentSessionRanRecorder + /** Where a chat's worktree is, for the file checkpoint taken before each turn. */ + resolveWorkspacePath?: ((workspaceId: string) => Promise) | undefined } diff --git a/src/main/native-chat/agent-session-wire/delivery/structured-agent-session-handover.ts b/src/main/native-chat/agent-session-wire/delivery/structured-agent-session-handover.ts index 896c310260a..7ede759ffff 100644 --- a/src/main/native-chat/agent-session-wire/delivery/structured-agent-session-handover.ts +++ b/src/main/native-chat/agent-session-wire/delivery/structured-agent-session-handover.ts @@ -113,9 +113,18 @@ export async function beginSubmissionHandover( } // Preserve handover time so turn-open cannot re-anchor elapsed work. const requestedAt = structuredAgentSessionHandoverOrigin(ctx.journal, submission) + let releaseCheckpoint: (() => void) | undefined return { run: async (commit) => { try { + // A message that opens a turn gets a file checkpoint first, taken before the agent can + // write; one that joins a running turn would snapshot the agent mid-edit, so it gets none. + if (turnScope.kind === 'thread' && ctx.captureTurnCheckpoint && !attempt?.signal.aborted) { + releaseCheckpoint = await ctx.captureTurnCheckpoint( + agentJournalSubmissionKey(clientMessageId), + attempt?.signal + ) + } let outcome: AgentSessionDispatchOutcome | null try { for (;;) { @@ -216,6 +225,7 @@ export async function beginSubmissionHandover( } }) } finally { + releaseCheckpoint?.() release?.() released() } diff --git a/src/main/native-chat/agent-session-wire/delivery/structured-agent-session-host-delivery.ts b/src/main/native-chat/agent-session-wire/delivery/structured-agent-session-host-delivery.ts index 7537b242ae6..04e3e74f955 100644 --- a/src/main/native-chat/agent-session-wire/delivery/structured-agent-session-host-delivery.ts +++ b/src/main/native-chat/agent-session-wire/delivery/structured-agent-session-host-delivery.ts @@ -220,6 +220,7 @@ export function createStructuredAgentSessionConversationDelivery(input: { stopping: input.clientDelivery.readStopping, now: () => deps.now?.() ?? Date.now(), acquireAborts: input.acquireAborts, + resolveWorkspacePath: deps.resolveWorkspacePath, recordRanOptions: (sessionId, fence) => createStructuredAgentSessionRanRecorder( { deps, sessions, now: () => deps.now?.() ?? Date.now() }, diff --git a/src/main/native-chat/agent-session-wire/delivery/structured-agent-session-input-handover.ts b/src/main/native-chat/agent-session-wire/delivery/structured-agent-session-input-handover.ts index f1d88d98181..faea4b45888 100644 --- a/src/main/native-chat/agent-session-wire/delivery/structured-agent-session-input-handover.ts +++ b/src/main/native-chat/agent-session-wire/delivery/structured-agent-session-input-handover.ts @@ -25,6 +25,7 @@ import { import { structuredAgentSessionEndedChildFailure } from '../termination/structured-agent-session-ended-child-failure' import { sameProviderChild } from '../host/structured-agent-session-provider-child' import { previousExitUnverifiableRefusal } from '../termination/structured-agent-session-child-close' +import { structuredAgentSessionTurnCheckpoint } from '../checkpoints/agent-session-checkpoints-host' type StartFailure = NonNullable> @@ -293,7 +294,8 @@ export async function beginStructuredAgentSessionInputHandover( }, get recordRanOptions() { return recorder - } + }, + captureTurnCheckpoint: structuredAgentSessionTurnCheckpoint(deps, sessionId) }, next, // Re-read right before the provider call, which runs outside the serialize. diff --git a/src/main/native-chat/agent-session-wire/host/structured-agent-session-host-types.ts b/src/main/native-chat/agent-session-wire/host/structured-agent-session-host-types.ts index 9673d8f4f10..db882e765d6 100644 --- a/src/main/native-chat/agent-session-wire/host/structured-agent-session-host-types.ts +++ b/src/main/native-chat/agent-session-wire/host/structured-agent-session-host-types.ts @@ -280,6 +280,9 @@ export type AgentSessionHandoverContext = StructuredAgentSessionCommandHandoverC settleEndedInput?: (owner: () => void) => Promise /** A message the agent took: a chat keeps, where it has none saved, the options it ran. */ recordRanOptions?: ReturnType + /** Snapshots the worktree before this message opens a turn; never throws. */ + /** Resolves to a release, called once the message has been handed to the agent. */ + captureTurnCheckpoint?: (itemId: string, signal?: AbortSignal) => Promise<() => void> } export type StructuredAgentSessionConversationLifetime = ReturnType< diff --git a/src/main/native-chat/agent-session-wire/host/testing/structured-agent-session-rewind-files.test.ts b/src/main/native-chat/agent-session-wire/host/testing/structured-agent-session-rewind-files.test.ts new file mode 100644 index 00000000000..98586ac9d60 --- /dev/null +++ b/src/main/native-chat/agent-session-wire/host/testing/structured-agent-session-rewind-files.test.ts @@ -0,0 +1,247 @@ +// End to end through the host: real sends snapshot a real git worktree before each turn, a real +// rewind goes back to a message, and the checkpoint restore (then its undo) puts the files back. + +import { mkdir, mkdtemp, readFile, rm, writeFile } from 'node:fs/promises' +import { existsSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { runProcess } from '@orca/process-host' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { agentJournalSubmissionKey } from '../../../../../shared/agent-session/journal/agent-session-journal-item-key' +import { computeAgentSessionPayloadFingerprint } from '../../../../../shared/agent-session/lifecycle/agent-session-mutation-envelope' +import type { AgentSessionRecordStore } from '../../../../runtime/agent-session/records/agent-session-record-store' +import { openTestAgentSessionRecordStore } from '../../../../runtime/agent-session/records/agent-session-record-store-test-harness' +import { StructuredAgentSessionHost } from '../structured-agent-session-host' +import type { AgentSessionDispatchOutcome } from '../../providers/structured-agent-session-adapter' +import { createAgentSessionCheckpoints } from '../../checkpoints/agent-session-checkpoints' +import { hostAgentSessionCheckpoints } from '../../checkpoints/agent-session-checkpoints-host' +import type { AgentSessionRecord } from '../../../../../shared/agent-session/lifecycle/agent-session-record' +import { + HOST_TEST_NOW, + HOST_TEST_SESSION, + HOST_TEST_THREAD, + hostTestAttachParams, + hostTestMessage, + hostTestOperationId, + resetHostTestOperationIds +} from './structured-agent-session-host-test-data' +import { + closeTestJournalHostDatabases, + openTestJournalHostDatabase +} from '../../../agent-session-journal/store-lifecycle/journal-host-database-test-support' +import { createStructuredAgentSessionLogger } from '../structured-agent-session-logger' +import { codexProviderHandle } from '../../../../../shared/agent-session/lifecycle/agent-session-provider-handle-encoding' +import { claudeAndCodexDeclared } from '../../providers/testing/structured-agent-session-adapter-router-test-support' +import { startsWhenPublished } from './structured-agent-session-instant-start.test-support' + +const caller = { callerKey: 'desktop' } +let directory: string +let worktree: string +let store: AgentSessionRecordStore +let host: StructuredAgentSessionHost +/** What "the agent" does to the worktree when a message reaches it. */ +let agentEdits: (() => Promise)[] + +async function git(...args: string[]): Promise { + const result = await runProcess({ program: 'git', args, cwd: worktree }) + if (result.code !== 0) { + throw new Error(result.stderr) + } + return result.stdout.trim() +} +const put = async (path: string, text: string) => { + await mkdir(join(worktree, path, '..'), { recursive: true }) + await writeFile(join(worktree, path), text) +} +const read = (path: string) => readFile(join(worktree, path), 'utf8') + +beforeEach(async () => { + resetHostTestOperationIds() + agentEdits = [] + directory = await mkdtemp(join(tmpdir(), 'orca-rewind-files-')) + worktree = await mkdtemp(join(tmpdir(), 'orca-rewind-files-tree-')) + await git('init', '-q') + await git('config', 'user.name', 'Audit') + await git('config', 'user.email', 'audit@example.test') + await put('.gitignore', 'node_modules/\n') + await put('app.ts', 'v0\n') + await git('add', '.') + await git('-c', 'commit.gpgsign=false', 'commit', '-qm', 'base') + store = await openTestAgentSessionRecordStore(directory) + let acquires = 0 + const adapter = startsWhenPublished( + { + supportsCreate: (_location, agent) => agent === 'codex', + supportsLocation: () => true, + acquire: async (input) => { + acquires += 1 + return { + process: { + hostId: 'local', + pid: 4000 + acquires, + processStartTimeMs: HOST_TEST_NOW, + spawnToken: input.spawnToken + }, + acquisitionGeneration: `generation-${acquires}`, + link: { + linkId: `link-${acquires}`, + mintedAtFence: input.fence, + observedAt: HOST_TEST_NOW, + origin: acquires === 1 ? 'created' : 'resumed', + handle: codexProviderHandle(HOST_TEST_THREAD) + } + } + }, + dispatch: vi.fn(async (input): Promise => { + await agentEdits.shift()?.() + return { + state: 'accepted', + providerIdentity: { + provider: 'codex', + threadId: HOST_TEST_THREAD, + turnId: input.clientMessageId, + ordinal: 0 + } + } + }), + cancelTurn: async () => ({ cancelled: false }), + answerPrompt: async () => {}, + setOption: async () => {}, + rewindSupport: () => ({ supported: true }), + releaseAcquisition: async () => true, + closeSession: async () => true + }, + () => host + ) + host = new StructuredAgentSessionHost({ + agents: claudeAndCodexDeclared(), + logger: createStructuredAgentSessionLogger(), + store, + adapter, + journalDatabase: openTestJournalHostDatabase(directory), + claimKeyId: 'key', + resolveWorkspacePath: async () => worktree, + now: () => HOST_TEST_NOW, + probeOwner: async () => ({ outcome: 'exit-observed' }) + }) +}) +afterEach(async () => { + await host.flushAllStreamedEvents() + await closeTestJournalHostDatabases() + await rm(directory, { recursive: true, force: true }) + await rm(worktree, { recursive: true, force: true }) +}) + +async function send(text: string, edit: () => Promise): Promise { + agentEdits.push(edit) + const clientOperationId = hostTestOperationId() + const body = hostTestMessage(text) + expect( + await host.send(caller, { + body, + envelope: { + sessionId: HOST_TEST_SESSION, + clientOperationId, + expectedRuntimeFence: store.getRecord(HOST_TEST_SESSION)!.lease.runtimeFence, + payloadFingerprint: computeAgentSessionPayloadFingerprint({ + method: 'agentSession.send', + sessionId: HOST_TEST_SESSION, + fields: { body } + }) + } + }) + ).toMatchObject({ ok: true }) + await vi.waitFor(async () => + expect( + (await host.journalSnapshot(HOST_TEST_SESSION)).submissions.find( + (entry) => entry.clientMessageId === clientOperationId + )?.dispatchState + ).toBe('accepted') + ) + return agentJournalSubmissionKey(clientOperationId) +} + +describe('rewind with file restore', () => { + it('snapshots before each turn and puts files back to the rewound message', async () => { + const attached = await host.attach(caller, hostTestAttachParams(null)) + expect(attached, JSON.stringify(attached)).toMatchObject({ ok: true }) + const first = await send('add a feature', async () => { + await put('app.ts', 'v1\n') + await put('feature.ts', 'feature\n') + }) + const second = await send('now refactor', async () => { + await put('app.ts', 'v2\n') + await rm(join(worktree, 'feature.ts')) + await put('refactor/util.ts', 'util\n') + await put('node_modules/dep/index.js', 'ignored\n') + }) + const refs = await git('for-each-ref', '--format=%(refname)', 'refs/orca/checkpoints/') + expect(refs.split('\n')).toHaveLength(2) + expect(await git('stash', 'list')).toBe('') + + const preview = await hostAgentSessionCheckpoints(host).preview(HOST_TEST_SESSION, second) + expect(preview).toMatchObject({ available: true, changedFiles: 3, headMoved: false }) + const restored = await hostAgentSessionCheckpoints(host).restore(HOST_TEST_SESSION, second) + expect(restored.changedFiles).toBe(3) + expect(await read('app.ts')).toBe('v1\n') + expect(await read('feature.ts')).toBe('feature\n') + expect(existsSync(join(worktree, 'refactor'))).toBe(false) + expect(await read('node_modules/dep/index.js')).toBe('ignored\n') + + await hostAgentSessionCheckpoints(host).undo(HOST_TEST_SESSION, restored.undoId) + expect(await read('app.ts')).toBe('v2\n') + expect(existsSync(join(worktree, 'feature.ts'))).toBe(false) + expect(await read('refactor/util.ts')).toBe('util\n') + + await hostAgentSessionCheckpoints(host).restore(HOST_TEST_SESSION, first) + expect(await read('app.ts')).toBe('v0\n') + expect(existsSync(join(worktree, 'feature.ts'))).toBe(false) + }) + + it('has no checkpoint for unknown messages and refuses a bad undo id', async () => { + const attached = await host.attach(caller, hostTestAttachParams(null)) + expect(attached, JSON.stringify(attached)).toMatchObject({ ok: true }) + expect(await hostAgentSessionCheckpoints(host).preview(HOST_TEST_SESSION, 'orca:nope')).toEqual( + { + available: false, + reason: 'none' + } + ) + await expect( + hostAgentSessionCheckpoints(host).restore(HOST_TEST_SESSION, 'orca:nope') + ).rejects.toThrow('agent_session_checkpoint_none') + await expect( + hostAgentSessionCheckpoints(host).undo(HOST_TEST_SESSION, '../../heads/main') + ).rejects.toThrow('agent_session_checkpoint_none') + }) +}) + +describe('restore admission', () => { + it('refuses while a turn is between its snapshot and its hand-off, or the workspace is busy', async () => { + // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the checkpoints read only location and launchDirectory. + const record = { + sessionId: 's', + location: { workspaceId: 'w', wslDistro: null }, + launchDirectory: worktree + } as unknown as AgentSessionRecord + const checkpoints = createAgentSessionCheckpoints({ + record: () => record, + isBusy: () => false, + logger: createStructuredAgentSessionLogger(), + now: () => HOST_TEST_NOW + }) + const release = await checkpoints.captureTurn('s', 'orca:one') + await put('app.ts', 'agent\n') + await expect(checkpoints.restore('s', 'orca:one')).rejects.toThrow( + 'agent_session_checkpoint_busy' + ) + expect(await checkpoints.preview('s', 'orca:one')).toMatchObject({ busy: true }) + release() + await expect( + checkpoints.restore('s', 'orca:one', async (root) => root.length > 0) + ).rejects.toThrow('agent_session_checkpoint_busy') + expect(await read('app.ts')).toBe('agent\n') + await checkpoints.restore('s', 'orca:one', async () => false) + expect(await read('app.ts')).toBe('v0\n') + }) +}) diff --git a/src/main/runtime/rpc/methods/structured-sessions/conversation/structured-agent-session-checkpoint-methods.ts b/src/main/runtime/rpc/methods/structured-sessions/conversation/structured-agent-session-checkpoint-methods.ts new file mode 100644 index 00000000000..b00c150f26f --- /dev/null +++ b/src/main/runtime/rpc/methods/structured-sessions/conversation/structured-agent-session-checkpoint-methods.ts @@ -0,0 +1,70 @@ +// Per-turn file checkpoints: what restoring a message's snapshot would change, the restore itself +// (which first saves an undo point), and undoing that restore. Same host code locally and on VMs. + +import { realpath } from 'node:fs/promises' +import { defineMethod } from '../../../core' +import { requireInstalledStructuredHost } from '../creation/structured-agent-session-gate' +import type { AgentSessionWorkspaceBusy } from '../../../../../native-chat/agent-session-wire/checkpoints/agent-session-checkpoints' +import { hostAgentSessionCheckpoints } from '../../../../../native-chat/agent-session-wire/checkpoints/agent-session-checkpoints-host' +import { getStructuredAgentSessionHost } from '../../../../../native-chat/agent-session-wire/host/structured-agent-session-registry' +import type { RpcContext } from '../../../core' +import { + CheckpointItemParams, + CheckpointUndoParams +} from '../../../../../../shared/rpc-contract/structured-agent-session-checkpoint-params' + +/** Any agent working in the chat's workspace (terminal agents included) may be writing files. + * Rows match by workspace id or by worktree path; a truncated listing can't prove the workspace + * idle, so it counts as busy. */ +function terminalAgentsWorking(ctx: RpcContext, sessionId: string): AgentSessionWorkspaceBusy { + return async (root) => { + const workspaceId = + getStructuredAgentSessionHost()?.deps.store.getRecord(sessionId)?.location.workspaceId + const ps = await ctx.runtime.getWorktreePs(10_000) + if (ps.truncated) { + return true + } + const rootReal = await realpath(root).catch(() => root) + for (const row of ps.worktrees) { + const sameRow = + row.worktreeId === workspaceId || + (await realpath(row.path).catch(() => row.path)) === rootReal + if (sameRow && row.agents.some((agent) => agent.state === 'working')) { + return true + } + } + return false + } +} + +export const STRUCTURED_AGENT_SESSION_CHECKPOINT_METHODS = [ + defineMethod({ + name: 'agentSession.checkpointPreview', + permission: 'workspace', + params: CheckpointItemParams, + handler: async (params, ctx) => + hostAgentSessionCheckpoints( + await requireInstalledStructuredHost(ctx, params.sessionId) + ).preview(params.sessionId, params.itemId, terminalAgentsWorking(ctx, params.sessionId)) + }), + defineMethod({ + name: 'agentSession.checkpointRestore', + permission: 'workspace', + params: CheckpointItemParams, + handler: async (params, ctx) => + hostAgentSessionCheckpoints( + await requireInstalledStructuredHost(ctx, params.sessionId) + ).restore(params.sessionId, params.itemId, terminalAgentsWorking(ctx, params.sessionId)) + }), + defineMethod({ + name: 'agentSession.checkpointUndo', + permission: 'workspace', + params: CheckpointUndoParams, + handler: async (params, ctx) => + hostAgentSessionCheckpoints(await requireInstalledStructuredHost(ctx, params.sessionId)).undo( + params.sessionId, + params.undoId, + terminalAgentsWorking(ctx, params.sessionId) + ) + }) +] diff --git a/src/main/runtime/rpc/methods/structured-sessions/structured-agent-session.ts b/src/main/runtime/rpc/methods/structured-sessions/structured-agent-session.ts index 6ae64ac2c10..0c6d92a95d8 100644 --- a/src/main/runtime/rpc/methods/structured-sessions/structured-agent-session.ts +++ b/src/main/runtime/rpc/methods/structured-sessions/structured-agent-session.ts @@ -28,6 +28,7 @@ import { STRUCTURED_AGENT_SESSION_PROMPT_RESPONSE_METHODS } from './conversation import { STRUCTURED_AGENT_SESSION_REVEAL_METHODS } from './conversation/structured-agent-session-reveal' import { STRUCTURED_AGENT_SESSION_REGENERATE_NAME_METHODS } from './conversation/structured-agent-session-regenerate-name' import { STRUCTURED_AGENT_SESSION_RENAME_METHODS } from './conversation/structured-agent-session-rename' +import { STRUCTURED_AGENT_SESSION_CHECKPOINT_METHODS } from './conversation/structured-agent-session-checkpoint-methods' import { STRUCTURED_AGENT_SESSION_QUEUED_METHODS } from './conversation/structured-agent-session-queued-methods' import { STRUCTURED_AGENT_SESSION_RESTART_RESUME_METHODS } from './conversation/structured-agent-session-restart-resume' import { STRUCTURED_AGENT_SESSION_STATUS_METHODS } from './subscriptions/structured-agent-session-status-stream' @@ -212,6 +213,7 @@ export const STRUCTURED_AGENT_SESSION_METHODS = [ ...STRUCTURED_AGENT_SESSION_REVEAL_METHODS, ...STRUCTURED_AGENT_SESSION_REGENERATE_NAME_METHODS, ...STRUCTURED_AGENT_SESSION_RENAME_METHODS, + ...STRUCTURED_AGENT_SESSION_CHECKPOINT_METHODS, ...STRUCTURED_AGENT_SESSION_RESTART_RESUME_METHODS, ...STRUCTURED_AGENT_SESSION_USAGE_LIMIT_CONTINUE_METHODS, ...STRUCTURED_AGENT_SESSION_STATUS_METHODS, diff --git a/src/renderer/src/components/confirmation-dialog-context.ts b/src/renderer/src/components/confirmation-dialog-context.ts index 65376347d3f..fba1252b34d 100644 --- a/src/renderer/src/components/confirmation-dialog-context.ts +++ b/src/renderer/src/components/confirmation-dialog-context.ts @@ -23,13 +23,24 @@ export type ConfirmationDialogOptions = { cancelVariant?: 'outline' | 'ghost' /** A notice, not a question: only the confirm button shows, and it settles true. */ acknowledgeOnly?: boolean - initialFocus?: 'confirm' + initialFocus?: 'confirm' | 'cancel' /** The short close-prompt look: small title, text and buttons. */ compact?: boolean /** Ignore repeated Enter and the opening double-click’s second click; fresh presses and Escape still settle the prompt. */ ignoresOpeningGesture?: boolean /** Renders a "Don't ask again" checkbox. `onConfirmed` runs only when the user confirms with it checked. */ dontAskAgain?: { label?: string; onConfirmed: () => void } + /** Renders a checkbox for an optional part of the action, with an optional line under it. + * `onConfirmed` gets the box's state, only on confirm. */ + option?: { + label: string + /** Describes the outcome while checked (or always, without `uncheckedDetail`). */ + detail?: string + uncheckedDetail?: string + defaultChecked: boolean + disabled?: boolean + onConfirmed: (checked: boolean) => void + } /** Runs when the cancel button settles the prompt, unlike Esc, an outside press or an abort. */ onCancelClicked?: () => void /** Renders text inputs; Enter in one confirms. `onConfirmed` gets their values, only on confirm. */ diff --git a/src/renderer/src/components/confirmation-dialog-option.tsx b/src/renderer/src/components/confirmation-dialog-option.tsx new file mode 100644 index 00000000000..744f5bdf55e --- /dev/null +++ b/src/renderer/src/components/confirmation-dialog-option.tsx @@ -0,0 +1,46 @@ +import React from 'react' +import { Checkbox } from '@/components/ui/checkbox' +import { Label } from '@/components/ui/label' +import type { ConfirmationDialogOptions } from '@/components/confirmation-dialog-context' + +/** A confirmation's optional part: sits in the body, right under the description, because it + * changes what the confirm button does. The line under it says what the current state means. */ +export function ConfirmationDialogOption({ + option, + checked, + onCheckedChange +}: { + option: NonNullable + checked: boolean + onCheckedChange: (checked: boolean) => void +}): React.JSX.Element { + const detail = + checked || option.uncheckedDetail === undefined ? option.detail : option.uncheckedDetail + return ( +
+ onCheckedChange(value === true)} + aria-describedby={detail ? 'confirmation-dialog-option-detail' : undefined} + /> +
+ + {detail ? ( +

+ {detail} +

+ ) : null} +
+
+ ) +} diff --git a/src/renderer/src/components/confirmation-dialog.test.tsx b/src/renderer/src/components/confirmation-dialog.test.tsx index b1e9e2f6237..09e1c694f9c 100644 --- a/src/renderer/src/components/confirmation-dialog.test.tsx +++ b/src/renderer/src/components/confirmation-dialog.test.tsx @@ -103,6 +103,37 @@ describe('ConfirmationDialogProvider', () => { await waitFor(() => expect(onSettled).toHaveBeenCalledWith(true)) }) + it('starts an option at its default and hands its state over only on confirm', async () => { + const onConfirmed = vi.fn() + const { onSettled } = renderDialog({ + title: 'Rewind to here?', + confirmLabel: 'Rewind', + option: { + label: 'Also restore 3 changed files', + detail: 'Later changes are discarded.', + uncheckedDetail: 'Files on disk stay as they are.', + defaultChecked: true, + onConfirmed + }, + initialFocus: 'cancel' + }) + + await userEvent.click(screen.getByRole('button', { name: 'ask' })) + const box = await screen.findByRole('checkbox', { name: 'Also restore 3 changed files' }) + await waitFor(() => + expect(document.activeElement).toBe(screen.getByRole('button', { name: 'Cancel' })) + ) + expect(box.getAttribute('aria-checked')).toBe('true') + expect(box.getAttribute('aria-describedby')).toBe('confirmation-dialog-option-detail') + expect(screen.getByText('Later changes are discarded.')).toBeTruthy() + await userEvent.click(box) + expect(screen.getByText('Files on disk stay as they are.')).toBeTruthy() + await userEvent.click(screen.getByRole('button', { name: 'Rewind' })) + + expect(onConfirmed).toHaveBeenCalledWith(false) + await waitFor(() => expect(onSettled).toHaveBeenCalledWith(true)) + }) + it('never saves the preference when the user backs out', async () => { const onConfirmed = vi.fn() const { onSettled } = renderDialog({ diff --git a/src/renderer/src/components/confirmation-dialog.tsx b/src/renderer/src/components/confirmation-dialog.tsx index 05dc78a514e..1404855494a 100644 --- a/src/renderer/src/components/confirmation-dialog.tsx +++ b/src/renderer/src/components/confirmation-dialog.tsx @@ -20,6 +20,7 @@ import { import { useAppStore } from '@/store' import { useTranslate } from '@/i18n/use-translate' import { cn } from '@/lib/utils' +import { ConfirmationDialogOption } from '@/components/confirmation-dialog-option' type ConfirmationDialogRequest = { id: number @@ -35,8 +36,11 @@ export function ConfirmationDialogProvider({ const t = useTranslate() const nextIdRef = useRef(0) const confirmButtonRef = useRef(null) + const cancelButtonRef = useRef(null) const [queue, setQueue] = useState([]) const [dontAskAgain, setDontAskAgain] = useState(false) + // null until the user touches the option, so each prompt starts from its own default. + const [optionChecked, setOptionChecked] = useState(null) const [fieldValues, setFieldValues] = useState>({}) const activeRequest = queue[0] ?? null const activeRequestRef = useRef(activeRequest) @@ -109,6 +113,10 @@ export function ConfirmationDialogProvider({ if (confirmed && dontAskAgain) { request.options.dontAskAgain?.onConfirmed() } + const option = request.options.option + if (confirmed && option) { + option.onConfirmed(optionChecked ?? option.defaultChecked) + } const fields = request.options.fields if (confirmed && fields) { const defaults = Object.fromEntries( @@ -118,6 +126,7 @@ export function ConfirmationDialogProvider({ } // Why: queued prompts must not inherit this request's preference or typed values. setDontAskAgain(false) + setOptionChecked(null) setFieldValues({}) request.resolve(confirmed) setQueue((currentQueue) => { @@ -127,7 +136,7 @@ export function ConfirmationDialogProvider({ return currentQueue.filter((queuedRequest) => queuedRequest.id !== request.id) }) }, - [dontAskAgain, fieldValues] + [dontAskAgain, fieldValues, optionChecked] ) return ( @@ -151,6 +160,9 @@ export function ConfirmationDialogProvider({ if (activeRequest?.options.initialFocus === 'confirm') { event.preventDefault() confirmButtonRef.current?.focus() + } else if (activeRequest?.options.initialFocus === 'cancel') { + event.preventDefault() + cancelButtonRef.current?.focus() } }} > @@ -178,6 +190,13 @@ export function ConfirmationDialogProvider({ ) : null} + {displayedRequest?.options.option ? ( + + ) : null} {displayedRequest?.options.fields ? (
{displayedRequest.options.fields.inputs.map((field) => ( @@ -229,6 +248,7 @@ export function ConfirmationDialogProvider({ {displayedRequest?.options.acknowledgeOnly ? null : (