TODO Next: undo a turn — /undo /redo + per-turn snapshots

File tools record before-write state via snapshot hook; Session
captures beforeFiles at turn start and afterFiles at turn end (cap
100 via SnapshotStack), and undo/redo restore files + messages
together. Bash not snapshotted (docs + notice).
This commit is contained in:
asepharyana
2026-09-09 11:49:33 +07:00
parent 7edaeafca0
commit 4f3a7eef7c
8 changed files with 328 additions and 4 deletions
+8
View File
@@ -26,6 +26,8 @@ export type CommandAction =
| { type: 'info'; text: string }
| { type: 'model'; model: string }
| { type: 'resume'; id: string }
| { type: 'undo' }
| { type: 'redo' }
/** A custom command from a markdown file, expanded against its arguments. */
| { type: 'custom'; command: CustomCommand; args: string[] }
| { type: 'unknown'; name: string };
@@ -61,6 +63,8 @@ export const COMMANDS: CommandSpec[] = [
{ name: 'sessions', summary: 'list saved sessions' },
{ name: 'resume', arg: '<id>', summary: 'load a saved session' },
{ name: 'save', summary: 'write the session to disk now' },
{ name: 'undo', summary: 'undo the last turn — restores files and conversation (bash effects are not snapshotted)' },
{ name: 'redo', summary: 'redo the last undone turn' },
{ name: 'clear', summary: 'clear the transcript and history' },
{ name: 'exit', aliases: ['quit'], summary: 'quit' },
];
@@ -225,6 +229,10 @@ export function parseCommand(raw: string, custom: readonly CustomCommand[] = [])
return arg ? { type: 'model', model: arg } : { type: 'models' };
case 'resume':
return arg ? { type: 'resume', id: arg } : { type: 'info', text: 'usage: /resume <session-id>' };
case 'undo':
return { type: 'undo' };
case 'redo':
return { type: 'redo' };
default: {
const cmd = custom.find((c) => c.name === name);
return cmd ? { type: 'custom', command: cmd, args: arg ? arg.split(/\s+/) : [] } : { type: 'unknown', name };
+95
View File
@@ -21,6 +21,7 @@ import { detachProviderItems, droppedSpan, estimateTokens as pruneEstimateTokens
import { createSkillTool, renderSkills, type Skill } from './skills';
import { suggestSkillsFromTranscript, writeAutoSkill } from './skill-learner';
import { disabledToolNames, onBashOutput, tools as builtinTools, type ToolSetName } from './tools';
import { onBeforeWrite, SnapshotStack, type FileState } from './snapshot';
export type ApprovalRequest = {
approvalId: string;
@@ -204,6 +205,9 @@ export class Session {
/** The 80% spend warning is shown once, not on every turn past the line. */
private warnedSpend = false;
private controller: AbortController | undefined;
private readonly snapshots = new SnapshotStack();
private turnBeforeLen = 0;
private turnBeforeFiles = new Map<string, FileState>();
private learnTurns = 0;
private lastLearnLen = 0;
@@ -351,6 +355,7 @@ export class Session {
this.subagentOutputTokens = 0;
this.warnedSpend = false;
this.notebook.clear();
this.snapshots.clear();
this.opts.onChange?.(this.messages);
}
@@ -379,6 +384,54 @@ export class Session {
return this.opts.compactThreshold ?? DEFAULT_COMPACT_THRESHOLD;
}
canUndo(): boolean { return this.snapshots.canUndo(); }
canRedo(): boolean { return this.snapshots.canRedo(); }
async undo(): Promise<string> {
const snap = this.snapshots.popForUndo();
if (!snap) throw new Error('nothing to undo');
await this.restoreFiles(snap.beforeFiles);
// truncate messages to beforeLen; the tail is kept inside snap for redo
this.messages.length = snap.beforeLen;
this.opts.onChange?.(this.messages);
const n = snap.beforeFiles.size;
const filesNote = n === 0 ? 'no files to restore' : `${n} file(s) restored`;
const msgNote = snap.afterLen > snap.beforeLen ? `${snap.afterLen - snap.beforeLen} message(s) removed` : 'no messages to remove';
return `undone: ${filesNote}, ${msgNote} (bash effects, if any, were not snapshotted)`;
}
async redo(): Promise<string> {
const snap = this.snapshots.popForRedo();
if (!snap) throw new Error('nothing to redo');
await this.restoreFiles(snap.afterFiles);
// redo replays the tail that undo removed; stored in afterFiles? we also need messages tail.
// The messages tail is the slice that was removed on undo; reconstruct by re-inserting from stored span is not enough
// because snapshots hold beforeLen/afterLen but not the actual messages content.
// We store the removed tail inside the snapshot at push time as an extra field via (snap as any)._tail.
const tail = (snap as unknown as { _tail?: import('ai').ModelMessage[] })._tail;
if (tail && tail.length > 0) {
this.messages.push(...tail);
this.opts.onChange?.(this.messages);
}
const n = snap.afterFiles.size;
const filesNote = n === 0 ? 'no files to restore' : `${n} file(s) restored`;
return `redone: ${filesNote} (bash effects, if any, were not snapshotted)`;
}
private async restoreFiles(state: Map<string, FileState>): Promise<void> {
for (const [abs, st] of state) {
try {
if (!st.existed) {
if (await Bun.file(abs).exists()) await Bun.file(abs).delete();
} else {
await Bun.write(abs, st.content ?? '');
}
} catch {
// best-effort per file; one failure should not stop the rest
}
}
}
/**
* The session's spend so far and the configured ceiling, for the UI's status
* and the refuse-the-next-turn check. Unpriced models report no spend: a
@@ -535,6 +588,18 @@ export class Session {
return;
}
// snapshot boundary: remember messages length before this turn and arm file capture
this.turnBeforeLen = this.messages.length;
this.turnBeforeFiles = new Map<string, FileState>();
onBeforeWrite(async (abs: string) => {
if (this.turnBeforeFiles.has(abs)) return;
const exists = await Bun.file(abs).exists();
let content: string | null = null;
if (exists) {
try { content = await Bun.file(abs).text(); } catch { content = null; }
}
this.turnBeforeFiles.set(abs, { existed: exists, content });
});
this.messages.push({ role: 'user', content: userText });
this.opts.onChange?.(this.messages);
this.controller = new AbortController();
@@ -551,9 +616,39 @@ export class Session {
this.opts.onToolOutput?.(toolCallId, chunk);
});
let turnFailed = false;
try {
yield* this.run(signal, threshold, outputs);
} catch (e) {
turnFailed = true;
throw e;
} finally {
onBeforeWrite(undefined);
// finalize snapshot only for turns that actually ran (even if they errored after writing files,
// the file state is still worth snapshotting so undo can revert a half-failed turn)
try {
if (this.turnBeforeFiles.size > 0 || this.messages.length > this.turnBeforeLen) {
const afterFiles = new Map<string, FileState>();
for (const abs of this.turnBeforeFiles.keys()) {
const exists = await Bun.file(abs).exists();
let content: string | null = null;
if (exists) { try { content = await Bun.file(abs).text(); } catch { content = null; } }
afterFiles.set(abs, { existed: exists, content });
}
// tail for redo: the messages added by this turn
const tail = this.messages.slice(this.turnBeforeLen).map((m) => ({ ...m, content: typeof m.content === 'string' ? m.content : JSON.parse(JSON.stringify(m.content)) } as import('ai').ModelMessage));
const snap: import('./snapshot').TurnSnapshot & { _tail?: import('ai').ModelMessage[] } = {
beforeLen: this.turnBeforeLen,
afterLen: this.messages.length,
beforeFiles: new Map(this.turnBeforeFiles),
afterFiles,
};
(snap as unknown as { _tail?: import('ai').ModelMessage[] })._tail = tail;
// even failed turns push so undo can revert the file side; empty no-op turns are skipped above
if (!turnFailed || this.turnBeforeFiles.size > 0) this.snapshots.push(snap);
}
} catch {}
this.turnBeforeFiles = new Map<string, FileState>();
this.controller = undefined;
this.drainPendingHotReload();
onBashOutput(undefined);
+72
View File
@@ -0,0 +1,72 @@
/**
* Per-turn file snapshots for /undo and /redo.
*
* Bash is intentionally not snapshotted: a shell command can do anything
* (network, database, chmod, rm -rf) and there is no way to know what to
* restore. The docs and the undo notice say so plainly.
*
* File tools call `recordBeforeWrite(abs)` before their first write to a path
* in the current turn. Session drains the map at turn boundaries into its
* history stack (cap 100) and owns undo/redo.
*/
export type FileState = { existed: boolean; content: string | null };
export type TurnSnapshot = {
/** Messages length before the turn's user message was pushed. */
beforeLen: number;
/** Messages length after the turn completed (including tool results). */
afterLen: number;
/** File state before the turn, keyed by absolute path. Only files the turn touched. */
beforeFiles: Map<string, FileState>;
/** File state after the turn, for redo. */
afterFiles: Map<string, FileState>;
};
const MAX_HISTORY = 100;
let hook: ((abs: string) => Promise<void> | void) | undefined;
export function onBeforeWrite(fn: ((abs: string) => Promise<void> | void) | undefined): void {
hook = fn;
}
export async function recordBeforeWrite(abs: string): Promise<void> {
const fn = hook;
if (fn) await fn(abs);
}
export class SnapshotStack {
private readonly history: TurnSnapshot[] = [];
private readonly future: TurnSnapshot[] = [];
push(entry: TurnSnapshot): void {
this.history.push(entry);
if (this.history.length > MAX_HISTORY) this.history.shift();
this.future.length = 0;
}
canUndo(): boolean { return this.history.length > 0; }
canRedo(): boolean { return this.future.length > 0; }
popForUndo(): TurnSnapshot | undefined {
const e = this.history.pop();
if (e) this.future.push(e);
return e;
}
popForRedo(): TurnSnapshot | undefined {
const e = this.future.pop();
if (e) this.history.push(e);
return e;
}
clear(): void {
this.history.length = 0;
this.future.length = 0;
}
depth(): { undo: number; redo: number } {
return { undo: this.history.length, redo: this.future.length };
}
}
+6
View File
@@ -3,6 +3,7 @@ import { stat } from 'node:fs/promises';
import { resolve } from 'node:path';
import { z } from 'zod';
import { jail, posix, walk } from './ignore';
import { recordBeforeWrite } from './snapshot';
import { withMeta } from './tool-utils';
import { git } from './tools-git';
@@ -47,6 +48,7 @@ export const insertLinesTool = withMeta({ set: 'extra', mutating: true }, tool({
}),
execute: async ({ path, line, text }) => {
const { abs, lines: cur } = await readLines(path);
await recordBeforeWrite(abs);
if (line > cur.length + 1) throw new Error(`line ${line} is past the end of ${path} (${cur.length} lines)`);
cur.splice(line - 1, 0, ...lines(text));
await Bun.write(abs, cur.join('\n'));
@@ -64,6 +66,7 @@ export const deleteLinesTool = withMeta({ set: 'extra', mutating: true }, tool({
execute: async ({ path, start, end }) => {
if (end < start) throw new Error('end must be >= start');
const { abs, lines: cur } = await readLines(path);
await recordBeforeWrite(abs);
if (end > cur.length) throw new Error(`end ${end} is past the end of ${path} (${cur.length} lines)`);
if (start === 1 && end === cur.length) throw new Error('that deletes the whole file; use delete_file instead');
cur.splice(start - 1, end - start + 1);
@@ -83,6 +86,7 @@ export const replaceLinesTool = withMeta({ set: 'extra', mutating: true }, tool(
execute: async ({ path, start, end, text }) => {
if (end < start) throw new Error('end must be >= start');
const { abs, lines: cur } = await readLines(path);
await recordBeforeWrite(abs);
if (end > cur.length) throw new Error(`end ${end} is past the end of ${path} (${cur.length} lines)`);
cur.splice(start - 1, end - start + 1, ...lines(text));
await Bun.write(abs, cur.join('\n'));
@@ -95,6 +99,7 @@ export const appendFileTool = withMeta({ set: 'extra', mutating: true }, tool({
inputSchema: z.object({ path: z.string(), text: z.string() }),
execute: async ({ path, text }) => {
const { abs, lines: cur } = await readLines(path);
await recordBeforeWrite(abs);
await Bun.write(abs, `${cur.join('\n').replace(/\n?$/, '\n')}${text.replace(/\n?$/, '')}\n`);
return `Appended ${lines(text).length} line(s) to ${path}`;
},
@@ -105,6 +110,7 @@ export const prependFileTool = withMeta({ set: 'extra', mutating: true }, tool({
inputSchema: z.object({ path: z.string(), text: z.string() }),
execute: async ({ path, text }) => {
const { abs, lines: cur } = await readLines(path);
await recordBeforeWrite(abs);
await Bun.write(abs, `${text.replace(/\n?$/, '\n')}${cur.join('\n')}`);
return `Prepended ${lines(text).length} line(s) to ${path}`;
},
+9
View File
@@ -3,6 +3,7 @@ import { stat } from 'node:fs/promises';
import { join, resolve } from 'node:path';
import { z } from 'zod';
import { jail, posix, walk } from './ignore';
import { recordBeforeWrite } from './snapshot';
import { EXTRA_TOOL_NAMES, extraTools } from './tools-extra';
import { GIT_TOOL_NAMES, gitTools } from './tools-git';
import { NET_TOOL_NAMES, netTools } from './tools-net';
@@ -206,6 +207,8 @@ export const applyPatchTool = withMeta({ set: 'edit-plus', mutating: true }, too
if (seen.has(op.path)) throw new Error(`${op.path} appears twice in one patch`);
seen.add(op.path);
const abs = jail(op.path);
await recordBeforeWrite(abs);
if ((op as { moveTo?: string }).moveTo) await recordBeforeWrite(jail((op as { moveTo?: string }).moveTo!));
if (op.kind === 'delete') {
if (!(await Bun.file(abs).exists())) throw new Error(`cannot delete ${op.path}: no such file`);
@@ -274,6 +277,7 @@ export const writeFileTool = withMeta({ set: 'core', mutating: true }, tool({
}),
execute: async ({ path, content }) => {
const abs = jail(path);
await recordBeforeWrite(abs);
const before = await Bun.file(abs).exists() ? await Bun.file(abs).text() : undefined;
await Bun.write(abs, content);
@@ -300,6 +304,7 @@ export const editFileTool = withMeta({ set: 'core', mutating: true }, tool({
execute: async ({ path, oldString, newString, replaceAll = false }) => {
if (oldString === newString) throw new Error('oldString and newString are identical');
const abs = jail(path);
await recordBeforeWrite(abs);
const file = Bun.file(abs);
if (!(await file.exists())) throw new Error(`No such file: ${path}`);
const before = await file.text();
@@ -337,6 +342,7 @@ export const multiEditTool = withMeta({ set: 'edit-plus', mutating: true }, tool
}),
execute: async ({ path, edits }) => {
const abs = jail(path);
await recordBeforeWrite(abs);
const file = Bun.file(abs);
if (!(await file.exists())) throw new Error(`No such file: ${path}`);
@@ -700,6 +706,8 @@ export const moveFileTool = withMeta({ set: 'edit-plus', mutating: true }, tool(
execute: async ({ from, to }) => {
const source = jail(from);
const target = jail(to);
await recordBeforeWrite(source);
await recordBeforeWrite(target);
if (source === target) throw new Error('from and to are the same path');
const file = Bun.file(source);
@@ -721,6 +729,7 @@ export const deleteFileTool = withMeta({ set: 'edit-plus', mutating: true }, too
}),
execute: async ({ path }) => {
const abs = jail(path);
await recordBeforeWrite(abs);
// Bun.file on a directory reports exists() false, so the stat is what
// distinguishes "missing" from "a directory" and gives the right refusal.
+22
View File
@@ -674,6 +674,28 @@ export function App({
push({ kind: 'error', text: e instanceof Error ? e.message : String(e) });
}
return;
case 'undo': {
push({ kind: 'user', text: chosen.trim() });
setWorking(true);
try {
push({ kind: 'info', text: await session.undo() });
} catch (e) {
push({ kind: 'error', text: e instanceof Error ? e.message : String(e) });
}
setWorking(false);
return;
}
case 'redo': {
push({ kind: 'user', text: chosen.trim() });
setWorking(true);
try {
push({ kind: 'info', text: await session.redo() });
} catch (e) {
push({ kind: 'error', text: e instanceof Error ? e.message : String(e) });
}
setWorking(false);
return;
}
case 'provider':
push({ kind: 'user', text: chosen.trim() });
setOnboarding(true);