diff --git a/.changeset/warm-dragons-rewind.md b/.changeset/warm-dragons-rewind.md new file mode 100644 index 0000000000..5ddea5e8da --- /dev/null +++ b/.changeset/warm-dragons-rewind.md @@ -0,0 +1,5 @@ +--- +'@moonshot-ai/kimi-code': minor +--- + +Add `/rewind` to restore conversation history and tracked workspace files from local, Git-independent checkpoints. Run `/rewind` to review and confirm the pending file changes. diff --git a/apps/kimi-code/package.json b/apps/kimi-code/package.json index 2329f4409f..6077956a39 100644 --- a/apps/kimi-code/package.json +++ b/apps/kimi-code/package.json @@ -99,6 +99,7 @@ "chalk": "^5.4.1", "cli-highlight": "^2.1.11", "commander": "^13.1.0", + "ignore": "^5.3.2", "jimp": "^1.6.1", "pathe": "^2.0.3", "postject": "1.0.0-alpha.6", diff --git a/apps/kimi-code/src/tui/commands/dispatch.ts b/apps/kimi-code/src/tui/commands/dispatch.ts index b36951c784..18863119f9 100644 --- a/apps/kimi-code/src/tui/commands/dispatch.ts +++ b/apps/kimi-code/src/tui/commands/dispatch.ts @@ -12,6 +12,7 @@ import type { TasksBrowserController } from '../controllers/tasks-browser'; import { tryHandleDanceCommand } from '../easter-eggs/dance'; import type { ResolvedTheme } from '../theme/colors'; import type { TUIState } from '../tui-state'; +import type { WorkspaceCheckpointStore } from '../workspace-checkpoints'; import type { AppState, LoginProgressSpinnerHandle, @@ -64,6 +65,7 @@ import { } from './session'; import { handleSwarmCommand } from './swarm'; import { handleUndoCommand } from './undo'; +import { handleRewindCommand } from './rewind'; import { handleWebCommand } from './web'; // --------------------------------------------------------------------------- @@ -102,6 +104,7 @@ export { handleTitleCommand, } from './session'; export { handleUndoCommand } from './undo'; +export { handleRewindCommand } from './rewind'; export { handleWebCommand } from './web'; // --------------------------------------------------------------------------- @@ -166,6 +169,8 @@ export interface SlashCommandHost { /** Reset the client-side cache-break baseline after the context was cut * (/undo): the next step's cache-read drop is expected, not a break. */ noteContextCut?(): void; + /** Session-local filesystem before-images used by `/rewind`. */ + getWorkspaceCheckpointStore?(): WorkspaceCheckpointStore | undefined; // UI showLoginProgressSpinner(label: string): LoginProgressSpinnerHandle; @@ -342,6 +347,7 @@ const SESSION_REQUIRING_COMMANDS: ReadonlySet = new Set 'goal', 'init', 'plan', + 'rewind', 'swarm', 'undo', 'web', @@ -512,6 +518,9 @@ async function handleBuiltInSlashCommand( case 'undo': await handleUndoCommand(host, args); return; + case 'rewind': + await handleRewindCommand(host, args); + return; case 'web': await handleWebCommand(host); return; diff --git a/apps/kimi-code/src/tui/commands/index.ts b/apps/kimi-code/src/tui/commands/index.ts index 7449dba9bc..f51a03419f 100644 --- a/apps/kimi-code/src/tui/commands/index.ts +++ b/apps/kimi-code/src/tui/commands/index.ts @@ -30,6 +30,7 @@ export { handleGoalCommand, parseGoalCommand } from './goal'; export { goalArgumentCompletions } from './registry'; export { handleForkCommand, handleInitCommand, handleTitleCommand } from './session'; export { handleUndoCommand } from './undo'; +export { handleRewindCommand } from './rewind'; export { handleWebCommand } from './web'; export { promptApiKey, diff --git a/apps/kimi-code/src/tui/commands/registry.ts b/apps/kimi-code/src/tui/commands/registry.ts index 48b57aa3f8..2251b54882 100644 --- a/apps/kimi-code/src/tui/commands/registry.ts +++ b/apps/kimi-code/src/tui/commands/registry.ts @@ -354,6 +354,13 @@ export const BUILTIN_SLASH_COMMANDS = [ priority: 80, availability: 'idle-only', }, + { + name: 'rewind', + aliases: [], + description: 'Restore prompts and workspace files from local checkpoints', + priority: 80, + availability: 'idle-only', + }, { name: 'editor', aliases: [], diff --git a/apps/kimi-code/src/tui/commands/rewind.ts b/apps/kimi-code/src/tui/commands/rewind.ts new file mode 100644 index 0000000000..679a277b3e --- /dev/null +++ b/apps/kimi-code/src/tui/commands/rewind.ts @@ -0,0 +1,229 @@ +import { join } from 'node:path'; + +import { ChoicePickerComponent } from '../components/dialogs/choice-picker'; +import { UndoSelectorComponent } from '../components/dialogs/undo-selector'; +import { NO_ACTIVE_SESSION_MESSAGE } from '../constant/kimi-tui'; +import { formatErrorMessage } from '../utils/event-payload'; +import type { WorkspaceChange, WorkspaceRewindPlan } from '../workspace-checkpoints'; +import type { SlashCommandHost } from './dispatch'; +import { + createUndoChoices, + parseUndoCount, + resolveUndoAvailability, + undoByCount, +} from './undo'; + +export async function handleRewindCommand( + host: SlashCommandHost, + args: string = '', +): Promise { + if (host.state.appState.streamingPhase !== 'idle') { + host.showError('Cannot rewind while streaming — press Esc or Ctrl-C first.'); + return; + } + if (host.session === undefined) { + host.showError(NO_ACTIVE_SESSION_MESSAGE); + return; + } + const store = host.getWorkspaceCheckpointStore?.(); + if (store === undefined) { + host.showError('Workspace rewind is unavailable because this session has no local checkpoint store.'); + return; + } + + const availability = await resolveUndoAvailability(host); + const checkpointCount = await store.availableCount(); + const maxCount = Math.min(availability.maxCount, checkpointCount); + const trimmed = args.trim(); + if (trimmed.length === 0) { + const choices = createUndoChoices( + host.state.transcriptEntries, + host.state.transcriptContainer.children, + maxCount, + ); + if (choices.length === 0) { + host.showStatus(checkpointCount === 0 ? 'No workspace checkpoints to rewind.' : 'Nothing to rewind.'); + return; + } + host.mountEditorReplacement( + new UndoSelectorComponent({ + title: 'Select messages and files to rewind', + choices, + onSelect: (choice) => { + void rewindByCount(host, choice.count, true).then((rewound) => { + if (rewound) host.restoreInputText(choice.input); + }); + }, + onCancel: () => { + host.restoreEditor(); + }, + }), + ); + return; + } + + const count = parseUndoCount(trimmed); + if (count === undefined) { + host.showError('Usage: /rewind [count], where count is a positive integer.'); + return; + } + if (count > maxCount) { + host.showError( + `Cannot rewind ${formatPromptCount(count)}; only ${formatPromptCount(maxCount)} have both conversation history and workspace checkpoints.`, + ); + return; + } + await rewindByCount(host, count, false); +} + +async function rewindByCount( + host: SlashCommandHost, + count: number, + editorIsReplaced: boolean, +): Promise { + const store = host.getWorkspaceCheckpointStore?.(); + if (store === undefined) { + if (editorIsReplaced) host.restoreEditor(); + host.showError('Workspace rewind is unavailable for this session.'); + return false; + } + + let plan: WorkspaceRewindPlan; + try { + plan = await store.prepareRewind(count); + } catch (error) { + if (editorIsReplaced) host.restoreEditor(); + host.showError(`Cannot prepare rewind: ${formatErrorMessage(error)}`); + return false; + } + + const confirmed = await confirmRewind(host, plan); + if (!confirmed) { + await store.releasePreview().catch(() => undefined); + return false; + } + + try { + await store.apply(plan); + } catch (error) { + await store.releasePreview().catch(() => undefined); + host.showError(`Rewind aborted before conversation history changed: ${formatErrorMessage(error)}`); + return false; + } + + const conversationUndone = await undoByCount(host, count, { + preserveWorkspaceCheckpoints: true, + }); + if (!conversationUndone) { + try { + await store.rollback(plan); + } catch (error) { + host.showError( + `Conversation rewind failed and workspace rollback also failed: ${formatErrorMessage(error)}`, + ); + return false; + } + await store.releasePreview().catch(() => undefined); + host.showStatus( + 'Conversation rewind failed; workspace files were restored to their pre-rewind state.', + 'warning', + ); + return false; + } + + try { + await store.commit(plan); + } catch (error) { + // The user-visible rewind already succeeded. Invalidate stale metadata so + // a later command cannot apply checkpoints against the wrong turn suffix. + await store.invalidate().catch(() => undefined); + host.showStatus( + `Rewind completed, but its checkpoint metadata could not be finalized: ${formatErrorMessage(error)}`, + 'warning', + ); + return true; + } + + host.showStatus( + `Rewound ${formatPromptCount(count)} and ${formatFileCount(plan.changes.length)}.`, + 'success', + ); + return true; +} + +function confirmRewind(host: SlashCommandHost, plan: WorkspaceRewindPlan): Promise { + const summary = summarizeChanges(plan.changes); + return new Promise((resolveConfirmed) => { + let completed = false; + const finish = (confirmed: boolean): void => { + if (completed) return; + completed = true; + host.restoreEditor(); + resolveConfirmed(confirmed); + }; + host.mountEditorReplacement( + new ChoicePickerComponent({ + title: `Rewind ${formatPromptCount(plan.count)} and workspace files?`, + hint: 'Review the workspace changes below · Enter/Space select · Esc cancel', + notice: summary.notice, + noticeTone: plan.changes.length === 0 ? 'success' : 'warning', + currentValue: 'cancel', + options: [ + { + value: 'cancel', + label: 'Cancel', + description: 'Leave conversation history and workspace files unchanged.', + }, + { + value: 'rewind', + label: 'Rewind conversation and files', + tone: 'danger', + description: summary.description, + }, + ], + onSelect: (value) => { + finish(value === 'rewind'); + }, + onCancel: () => { + finish(false); + }, + }), + ); + }); +} + +function summarizeChanges(changes: readonly WorkspaceChange[]): { + readonly notice: string; + readonly description: string; +} { + const created = changes.filter((change) => change.kind === 'created').length; + const modified = changes.filter((change) => change.kind === 'modified').length; + const deleted = changes.filter((change) => change.kind === 'deleted').length; + if (changes.length === 0) { + return { + notice: 'No tracked workspace files changed; only conversation history will be rewound.', + description: 'Withdraw the selected prompts from the active context.', + }; + } + const preview = changes.slice(0, 8).map((change) => { + const action = change.kind === 'created' ? 'delete' : change.kind === 'deleted' ? 'restore' : 'restore'; + return `${action.padEnd(7)} ${displayWorkspacePath(change)}`; + }); + if (changes.length > preview.length) preview.push(`…and ${changes.length - preview.length} more`); + return { + notice: [`Workspace delta: ${created} created, ${modified} modified, ${deleted} deleted.`, ...preview].join('\n'), + description: `Restore ${formatFileCount(changes.length)} to the state before the selected prompts. Files ignored by .gitignore/.ignore, dependency trees, VCS metadata, and symlinks are outside the checkpoint.`, + }; +} + +function displayWorkspacePath(change: WorkspaceChange): string { + return JSON.stringify(join(change.root, change.path)); +} + +function formatPromptCount(count: number): string { + return `${count} ${count === 1 ? 'prompt' : 'prompts'}`; +} + +function formatFileCount(count: number): string { + return `${count} workspace ${count === 1 ? 'file' : 'files'}`; +} diff --git a/apps/kimi-code/src/tui/commands/undo.ts b/apps/kimi-code/src/tui/commands/undo.ts index 23d5a673e2..fa7ea65d73 100644 --- a/apps/kimi-code/src/tui/commands/undo.ts +++ b/apps/kimi-code/src/tui/commands/undo.ts @@ -30,7 +30,7 @@ import type { SlashCommandHost } from './dispatch'; // Undo command // --------------------------------------------------------------------------- -interface UndoAvailability { +export interface UndoAvailability { readonly maxCount: number; readonly stoppedAtCompaction: boolean; } @@ -77,7 +77,11 @@ export async function handleUndoCommand( await undoByCount(host, count); } -async function undoByCount(host: SlashCommandHost, count: number): Promise { +export async function undoByCount( + host: SlashCommandHost, + count: number, + options: { readonly preserveWorkspaceCheckpoints?: boolean } = {}, +): Promise { const session = host.session; if (session === undefined) { host.showError(NO_ACTIVE_SESSION_MESSAGE); @@ -122,6 +126,17 @@ async function undoByCount(host: SlashCommandHost, count: number): Promise { ); } -function parseUndoCount(args: string): number | undefined { +export function parseUndoCount(args: string): number | undefined { const value = args.trim(); if (value.length === 0) return 1; if (!/^[1-9]\d*$/.test(value)) return undefined; @@ -170,7 +185,7 @@ function parseUndoCount(args: string): number | undefined { return Number.isSafeInteger(count) ? count : undefined; } -async function resolveUndoAvailability( +export async function resolveUndoAvailability( host: SlashCommandHost, ): Promise { const local = undoAvailabilityFromTranscript( @@ -246,7 +261,7 @@ function isContextUndoAnchor(message: ContextMessage): boolean { return false; } -function createUndoChoices( +export function createUndoChoices( entries: readonly TranscriptEntry[], children: readonly Component[], maxCount: number, diff --git a/apps/kimi-code/src/tui/components/dialogs/undo-selector.ts b/apps/kimi-code/src/tui/components/dialogs/undo-selector.ts index 37320f92bd..08b5677d6f 100644 --- a/apps/kimi-code/src/tui/components/dialogs/undo-selector.ts +++ b/apps/kimi-code/src/tui/components/dialogs/undo-selector.ts @@ -22,6 +22,7 @@ export interface UndoChoice { } export interface UndoSelectorOptions { + readonly title?: string; readonly choices: readonly UndoChoice[]; readonly onSelect: (choice: UndoChoice) => void; readonly onCancel: () => void; @@ -70,7 +71,7 @@ export class UndoSelectorComponent extends Container implements Focusable { const lines: string[] = [ currentTheme.fg('primary', '─'.repeat(width)), - currentTheme.boldFg('primary', ' Select messages to undo'), + currentTheme.boldFg('primary', ` ${this.opts.title ?? 'Select messages to undo'}`), currentTheme.fg('textMuted', ' ' + hintParts.join(' · ')), '', ]; diff --git a/apps/kimi-code/src/tui/kimi-tui.ts b/apps/kimi-code/src/tui/kimi-tui.ts index 7c118e57a3..5d77bb3671 100644 --- a/apps/kimi-code/src/tui/kimi-tui.ts +++ b/apps/kimi-code/src/tui/kimi-tui.ts @@ -118,6 +118,7 @@ import { SessionReplayRenderer } from './controllers/session-replay'; import { StreamingUIController } from './controllers/streaming-ui'; import { TasksBrowserController } from './controllers/tasks-browser'; import { installRainbowDance } from './easter-eggs/dance'; +import { WorkspaceCheckpointStore } from './workspace-checkpoints'; import { adaptPanelResponse } from './reverse-rpc/approval/adapter'; import { ApprovalController } from './reverse-rpc/approval/controller'; import { createApprovalRequestHandler } from './reverse-rpc/approval/handler'; @@ -349,6 +350,9 @@ export class KimiTUI { private currentLoadingTip: { kind: LoadingTipKind; tip: string | undefined } | undefined = undefined; private lastHistoryContent: string | undefined; + private workspaceCheckpointCache: + | { readonly key: string; readonly store: WorkspaceCheckpointStore } + | undefined; // Live `!` shell output entries, keyed by commandId so concurrent commands // each update their own card and stale events are dropped. Mutated in place // as `shell.output` events arrive; removed when the command completes. @@ -1447,7 +1451,7 @@ export class KimiTUI { // the message. Steer instead: the engine buffers it into the running goal // turn, or launches a turn of its own if the loop just ended. if (this.state.appState.goal?.status === 'active') { - void session.steer(sdkInput).catch((error: unknown) => { + void this.runCheckpointedSessionRequest(session, () => session.steer(sdkInput)).catch((error: unknown) => { const message = formatErrorMessage(error); // Same reset as the prompt path: beginSessionRequest already moved the // TUI to the waiting phase, and no turn events may follow a failed @@ -1457,7 +1461,7 @@ export class KimiTUI { }); return; } - void session.prompt(sdkInput).catch((error: unknown) => { + void this.runCheckpointedSessionRequest(session, () => session.prompt(sdkInput)).catch((error: unknown) => { const message = formatErrorMessage(error); this.failSessionRequest(`Failed to send: ${message}`); }); @@ -1479,7 +1483,10 @@ export class KimiTUI { } if (!this.validateMediaCapabilities(rewrite)) return; this.beginSessionRequest(); - void session.activateSkill(skillName, rewrite.text).catch((error: unknown) => { + void this.runCheckpointedSessionRequest( + session, + () => session.activateSkill(skillName, rewrite.text), + ).catch((error: unknown) => { const message = formatErrorMessage(error); this.failSessionRequest(`Skill "${skillName}" failed: ${message}`); }); @@ -1503,14 +1510,78 @@ export class KimiTUI { } if (!this.validateMediaCapabilities(rewrite)) return; this.beginSessionRequest(); - void session - .activatePluginCommand(pluginId, commandName, rewrite.text) + void this.runCheckpointedSessionRequest( + session, + () => session.activatePluginCommand(pluginId, commandName, rewrite.text), + ) .catch((error: unknown) => { const message = formatErrorMessage(error); this.failSessionRequest(`Command "${pluginId}:${commandName}" failed: ${message}`); }); } + getWorkspaceCheckpointStore(): WorkspaceCheckpointStore | undefined { + const session = this.session; + return session === undefined ? undefined : this.workspaceCheckpointStoreFor(session); + } + + private workspaceCheckpointStoreFor(session: Session): WorkspaceCheckpointStore | undefined { + const summary = session.summary; + if ( + summary === undefined || + typeof summary.sessionDir !== 'string' || + summary.sessionDir.trim().length === 0 || + typeof summary.workDir !== 'string' || + summary.workDir.trim().length === 0 + ) { + return undefined; + } + const roots = [summary.workDir, ...(summary.additionalDirs ?? [])]; + const key = JSON.stringify([session.id, summary.sessionDir, roots]); + if (this.workspaceCheckpointCache?.key === key) return this.workspaceCheckpointCache.store; + const store = new WorkspaceCheckpointStore(summary.sessionDir, roots); + this.workspaceCheckpointCache = { key, store }; + return store; + } + + private async runCheckpointedSessionRequest( + session: Session, + request: () => Promise, + ): Promise { + const store = this.workspaceCheckpointStoreFor(session); + let checkpointId: string | undefined; + if (store !== undefined) { + try { + checkpointId = await store.captureBeforeTurn(); + } catch (error) { + // A missing checkpoint must invalidate older entries: otherwise the + // newest stored before-image would no longer align with `/undo 1`. + try { + await store.invalidate(); + } catch { + // The original capture error is the actionable one. + } + this.showStatus( + `Workspace rewind unavailable for this turn: ${formatErrorMessage(error)}`, + 'warning', + ); + } + } + try { + await request(); + } catch (error) { + if (store !== undefined && checkpointId !== undefined) { + try { + await store.discardCaptured(checkpointId); + } catch { + // The request failure remains primary; stale checkpoints fail closed + // during `/rewind` root/tail validation. + } + } + throw error; + } + } + private sendMessage(session: Session, input: string, options?: SendMessageOptions): void { if ( this.deferUserMessages || diff --git a/apps/kimi-code/src/tui/workspace-checkpoints.ts b/apps/kimi-code/src/tui/workspace-checkpoints.ts new file mode 100644 index 0000000000..1176127009 --- /dev/null +++ b/apps/kimi-code/src/tui/workspace-checkpoints.ts @@ -0,0 +1,726 @@ +import { createHash, randomUUID } from 'node:crypto'; +import { + chmod, + lstat, + mkdir, + readFile, + readdir, + realpath, + rename, + unlink, + writeFile, +} from 'node:fs/promises'; +import { dirname, isAbsolute, join, relative, resolve, sep } from 'node:path'; + +import createIgnore, { type Ignore } from 'ignore'; + +const STORE_VERSION = 1; +const DEFAULT_MAX_FILES = 50_000; +const DEFAULT_MAX_BYTES = 512 * 1024 * 1024; +const DEFAULT_MAX_FILE_BYTES = 64 * 1024 * 1024; +const DEFAULT_MAX_CHECKPOINTS = 20; +const ALWAYS_IGNORED_DIRECTORIES = new Set(['.git', '.hg', '.svn', 'node_modules']); + +interface FileRecord { + readonly hash: string; + readonly mode: number; + readonly size: number; +} + +interface RootSnapshot { + readonly path: string; + readonly files: Readonly>; +} + +interface WorkspaceSnapshot { + readonly roots: readonly RootSnapshot[]; +} + +interface CheckpointRecord { + readonly id: string; + readonly createdAt: string; + readonly roots: readonly string[]; +} + +interface CheckpointIndex { + readonly version: typeof STORE_VERSION; + readonly checkpoints: readonly CheckpointRecord[]; +} + +export type WorkspaceChangeKind = 'created' | 'modified' | 'deleted'; + +export interface WorkspaceChange { + readonly root: string; + readonly path: string; + readonly kind: WorkspaceChangeKind; +} + +export interface WorkspaceRewindPlan { + readonly count: number; + readonly checkpointIds: readonly string[]; + readonly changes: readonly WorkspaceChange[]; + /** Internal snapshots are deliberately opaque to command/UI consumers. */ + readonly target: WorkspaceSnapshot; + readonly current: WorkspaceSnapshot; +} + +export interface WorkspaceCheckpointLimits { + readonly maxFiles?: number; + readonly maxBytes?: number; + readonly maxFileBytes?: number; + readonly maxCheckpoints?: number; +} + +export class WorkspaceCheckpointError extends Error { + constructor(message: string) { + super(message); + this.name = 'WorkspaceCheckpointError'; + } +} + +/** + * Session-local, content-addressed workspace checkpoints for `/rewind`. + * + * A checkpoint is captured before a prompt is submitted. Restores never use + * Git: they compare the current workspace with that before-image, validate it + * again immediately before applying, and then recreate/delete regular files. + */ +export class WorkspaceCheckpointStore { + private readonly storeDir: string; + private readonly blobsDir: string; + private readonly checkpointsDir: string; + private readonly indexPath: string; + private rootsPromise: Promise | undefined; + private operation: Promise = Promise.resolve(); + private invalidated = false; + private readonly maxFiles: number; + private readonly maxBytes: number; + private readonly maxFileBytes: number; + private readonly maxCheckpoints: number; + + constructor( + sessionDir: string, + private readonly configuredRoots: readonly string[], + limits: WorkspaceCheckpointLimits = {}, + ) { + this.storeDir = join(sessionDir, 'workspace-checkpoints'); + this.blobsDir = join(this.storeDir, 'blobs'); + this.checkpointsDir = join(this.storeDir, 'checkpoints'); + this.indexPath = join(this.storeDir, 'index.json'); + this.maxFiles = limits.maxFiles ?? DEFAULT_MAX_FILES; + this.maxBytes = limits.maxBytes ?? DEFAULT_MAX_BYTES; + this.maxFileBytes = limits.maxFileBytes ?? DEFAULT_MAX_FILE_BYTES; + this.maxCheckpoints = limits.maxCheckpoints ?? DEFAULT_MAX_CHECKPOINTS; + } + + async captureBeforeTurn(): Promise { + return this.exclusive(async () => { + const snapshot = await this.captureSnapshot(); + const record: CheckpointRecord = { + id: randomUUID(), + createdAt: new Date().toISOString(), + roots: snapshot.roots.map((root) => root.path), + }; + const wasInvalidated = this.invalidated; + const index = wasInvalidated + ? { version: STORE_VERSION, checkpoints: [] } + : await this.readIndex(); + const previousRoots = index.checkpoints.at(-1)?.roots; + const compatible = previousRoots === undefined || sameStrings(previousRoots, record.roots); + await this.writeSnapshot(record.id, snapshot); + const checkpoints = [...(compatible ? index.checkpoints : []), record].slice( + -this.maxCheckpoints, + ); + await this.writeIndex({ + version: STORE_VERSION, + checkpoints, + }); + this.invalidated = false; + if (wasInvalidated || !compatible || checkpoints.length <= index.checkpoints.length) { + await this.pruneStore(); + } + return record.id; + }); + } + + async discardCaptured(checkpointId: string): Promise { + await this.exclusive(async () => { + const index = await this.readIndex(); + if (index.checkpoints.at(-1)?.id !== checkpointId) return; + await this.writeIndex({ version: STORE_VERSION, checkpoints: index.checkpoints.slice(0, -1) }); + await this.pruneStore(); + }); + } + + async discardLast(count: number): Promise { + await this.exclusive(async () => { + const index = await this.readIndex(); + const keep = Math.max(0, index.checkpoints.length - count); + await this.writeIndex({ version: STORE_VERSION, checkpoints: index.checkpoints.slice(0, keep) }); + await this.pruneStore(); + }); + } + + async invalidate(): Promise { + this.invalidated = true; + await this.exclusive(async () => { + await this.writeIndex({ version: STORE_VERSION, checkpoints: [] }); + await this.pruneStore(); + }); + } + + async availableCount(): Promise { + if (this.invalidated) return 0; + return this.exclusive(async () => (await this.readIndex()).checkpoints.length); + } + + async prepareRewind(count: number): Promise { + return this.exclusive(async () => { + if (this.invalidated) { + throw new WorkspaceCheckpointError('Workspace checkpoints were invalidated by an earlier failure.'); + } + if (!Number.isSafeInteger(count) || count < 1) { + throw new WorkspaceCheckpointError('Rewind count must be a positive integer.'); + } + const index = await this.readIndex(); + if (count > index.checkpoints.length) { + throw new WorkspaceCheckpointError( + `Only ${index.checkpoints.length} workspace checkpoint${index.checkpoints.length === 1 ? '' : 's'} available.`, + ); + } + const selected = index.checkpoints.slice(-count); + const target = await this.readSnapshot(selected[0]!.id); + await this.assertCurrentRoots(target); + const current = await this.captureSnapshot(); + return { + count, + checkpointIds: selected.map((checkpoint) => checkpoint.id), + changes: diffSnapshots(target, current), + target, + current, + }; + }); + } + + async apply(plan: WorkspaceRewindPlan): Promise { + await this.exclusive(async () => { + await this.validatePlan(plan); + try { + await this.applySnapshot(plan.target, plan.current, plan.changes); + } catch (error) { + try { + await this.applySnapshot(plan.current, plan.target, plan.changes); + } catch (rollbackError) { + throw new WorkspaceCheckpointError( + `Workspace restore failed (${errorMessage(error)}) and rollback also failed (${errorMessage(rollbackError)}).`, + ); + } + throw error; + } + }); + } + + async rollback(plan: WorkspaceRewindPlan): Promise { + await this.exclusive(async () => { + await this.applySnapshot(plan.current, plan.target, plan.changes); + }); + } + + async commit(plan: WorkspaceRewindPlan): Promise { + await this.exclusive(async () => { + const index = await this.readIndex(); + const tail = index.checkpoints.slice(-plan.count).map((checkpoint) => checkpoint.id); + if (!sameStrings(tail, plan.checkpointIds)) { + throw new WorkspaceCheckpointError('Workspace checkpoints changed while rewind was open.'); + } + await this.writeIndex({ + version: STORE_VERSION, + checkpoints: index.checkpoints.slice(0, -plan.count), + }); + await this.pruneStore(); + }); + } + + async releasePreview(): Promise { + await this.exclusive(() => this.pruneStore()); + } + + private async exclusive(task: () => Promise): Promise { + const previous = this.operation; + let release!: () => void; + this.operation = new Promise((resolveOperation) => { + release = resolveOperation; + }); + await previous; + try { + return await task(); + } finally { + release(); + } + } + + private async roots(): Promise { + this.rootsPromise ??= normalizeRoots(this.configuredRoots); + return this.rootsPromise; + } + + private async captureSnapshot(): Promise { + await mkdir(this.blobsDir, { recursive: true, mode: 0o700 }); + const excludedStoreDir = await realpath(this.storeDir); + const roots = await this.roots(); + const snapshots: RootSnapshot[] = []; + let fileCount = 0; + let byteCount = 0; + for (const root of roots) { + const matcher = await loadIgnoreMatcher(root); + const files: Record = {}; + const pending: Array<{ absolute: string; relative: string }> = [{ absolute: root, relative: '' }]; + while (pending.length > 0) { + const directory = pending.pop()!; + let entries; + try { + entries = await readdir(directory.absolute, { withFileTypes: true }); + } catch (error) { + throw checkpointFsError('scan', directory.absolute, error); + } + entries.sort((left, right) => left.name.localeCompare(right.name)); + for (let index = entries.length - 1; index >= 0; index -= 1) { + const entry = entries[index]!; + const relativePath = directory.relative === '' ? entry.name : `${directory.relative}/${entry.name}`; + const absolutePath = join(directory.absolute, entry.name); + if (entry.isSymbolicLink()) continue; + if (entry.isDirectory()) { + if (resolve(absolutePath) === excludedStoreDir) continue; + if (ALWAYS_IGNORED_DIRECTORIES.has(entry.name) || matcher.ignores(`${relativePath}/`)) { + continue; + } + pending.push({ absolute: absolutePath, relative: relativePath }); + continue; + } + if (!entry.isFile() || matcher.ignores(relativePath)) continue; + let contents: Buffer; + let stats; + try { + stats = await lstat(absolutePath); + if (!stats.isFile()) continue; + if (stats.size > this.maxFileBytes) { + throw new WorkspaceCheckpointError( + `Cannot checkpoint ${relativePath}: file exceeds ${formatBytes(this.maxFileBytes)}.`, + ); + } + contents = await readFile(absolutePath); + if (contents.byteLength > this.maxFileBytes) { + throw new WorkspaceCheckpointError( + `Cannot checkpoint ${relativePath}: file exceeds ${formatBytes(this.maxFileBytes)}.`, + ); + } + } catch (error) { + if (error instanceof WorkspaceCheckpointError) throw error; + if (isMissing(error)) continue; + throw checkpointFsError('read', absolutePath, error); + } + fileCount += 1; + byteCount += contents.byteLength; + if (fileCount > this.maxFiles) { + throw new WorkspaceCheckpointError(`Workspace exceeds the ${this.maxFiles} file checkpoint limit.`); + } + if (byteCount > this.maxBytes) { + throw new WorkspaceCheckpointError( + `Workspace exceeds the ${formatBytes(this.maxBytes)} checkpoint limit.`, + ); + } + const hash = createHash('sha256').update(contents).digest('hex'); + await this.writeBlob(hash, contents); + files[relativePath] = { hash, mode: stats.mode & 0o777, size: contents.byteLength }; + } + } + snapshots.push({ path: root, files }); + } + return { roots: snapshots }; + } + + private async writeBlob(hash: string, contents: Buffer): Promise { + const path = join(this.blobsDir, hash); + try { + await writeFile(path, contents, { flag: 'wx', mode: 0o600 }); + } catch (error) { + if (!isAlreadyExists(error)) throw checkpointFsError('store', path, error); + } + } + + private async validatePlan(plan: WorkspaceRewindPlan): Promise { + const currentRoots = rootMap(plan.current); + for (const change of plan.changes) { + const expected = currentRoots.get(change.root)?.files[change.path]; + const actual = await readCurrentRecord(change.root, change.path); + if (!sameFile(expected, actual)) { + throw new WorkspaceCheckpointError( + `Workspace changed after the rewind preview: ${join(change.root, change.path)}`, + ); + } + } + } + + private async applySnapshot( + desired: WorkspaceSnapshot, + source: WorkspaceSnapshot, + changes: readonly WorkspaceChange[], + ): Promise { + const desiredRoots = rootMap(desired); + const sourceRoots = rootMap(source); + for (const change of changes) { + const wanted = desiredRoots.get(change.root)?.files[change.path]; + const previous = sourceRoots.get(change.root)?.files[change.path]; + const targetPath = safeWorkspacePath(change.root, change.path); + await assertNoSymlinkParent(change.root, change.path); + if (wanted === undefined) { + if (previous !== undefined) { + try { + await unlink(targetPath); + } catch (error) { + if (!isMissing(error)) throw checkpointFsError('remove', targetPath, error); + } + } + continue; + } + const contents = await readFile(join(this.blobsDir, wanted.hash)); + const actualHash = createHash('sha256').update(contents).digest('hex'); + if (actualHash !== wanted.hash || contents.byteLength !== wanted.size) { + throw new WorkspaceCheckpointError(`Checkpoint content is corrupt for ${targetPath}.`); + } + await mkdir(dirname(targetPath), { recursive: true }); + const temporary = `${targetPath}.kimi-rewind-${randomUUID()}`; + try { + await writeFile(temporary, contents, { mode: wanted.mode }); + await chmod(temporary, wanted.mode); + await rename(temporary, targetPath); + } catch (error) { + try { + await unlink(temporary); + } catch { + // Best effort cleanup; preserve the original write failure. + } + throw checkpointFsError('restore', targetPath, error); + } + } + } + + private async assertCurrentRoots(snapshot: WorkspaceSnapshot): Promise { + const current = await this.roots(); + const captured = snapshot.roots.map((root) => root.path); + if (!sameStrings(current, captured)) { + throw new WorkspaceCheckpointError( + 'Workspace roots changed after this checkpoint; rewind is unavailable for that turn.', + ); + } + } + + private async readIndex(): Promise { + try { + const parsed: unknown = JSON.parse(await readFile(this.indexPath, 'utf8')); + if (!isCheckpointIndex(parsed)) throw new Error('unsupported checkpoint index'); + return parsed; + } catch (error) { + if (isMissing(error)) return { version: STORE_VERSION, checkpoints: [] }; + throw checkpointFsError('read', this.indexPath, error); + } + } + + private async writeIndex(index: CheckpointIndex): Promise { + await this.writeJson(this.indexPath, index); + } + + private async readSnapshot(id: string): Promise { + try { + const parsed: unknown = JSON.parse(await readFile(this.checkpointPath(id), 'utf8')); + if (!isWorkspaceSnapshot(parsed)) throw new Error('invalid workspace snapshot'); + return parsed; + } catch (error) { + throw checkpointFsError('read', this.checkpointPath(id), error); + } + } + + private async writeSnapshot(id: string, snapshot: WorkspaceSnapshot): Promise { + await this.writeJson(this.checkpointPath(id), snapshot); + } + + private checkpointPath(id: string): string { + return join(this.checkpointsDir, `${id}.json`); + } + + private async writeJson(path: string, value: unknown): Promise { + await mkdir(dirname(path), { recursive: true, mode: 0o700 }); + const temporary = `${path}.${randomUUID()}.tmp`; + try { + await writeFile(temporary, `${JSON.stringify(value)}\n`, { mode: 0o600 }); + await rename(temporary, path); + } catch (error) { + try { + await unlink(temporary); + } catch { + // Best effort cleanup; preserve the original write failure. + } + throw checkpointFsError('write', path, error); + } + } + + private async pruneStore(): Promise { + const index = await this.readIndex(); + const retained = new Set(); + for (const checkpoint of index.checkpoints) { + const snapshot = await this.readSnapshot(checkpoint.id); + for (const root of snapshot.roots) { + for (const file of Object.values(root.files)) retained.add(file.hash); + } + } + await this.pruneCheckpointFiles(new Set(index.checkpoints.map((checkpoint) => checkpoint.id))); + let entries; + try { + entries = await readdir(this.blobsDir, { withFileTypes: true }); + } catch (error) { + if (isMissing(error)) return; + throw checkpointFsError('scan', this.blobsDir, error); + } + await Promise.all( + entries + .filter((entry) => entry.isFile() && !retained.has(entry.name)) + .map(async (entry) => { + try { + await unlink(join(this.blobsDir, entry.name)); + } catch (error) { + if (!isMissing(error)) throw error; + } + }), + ); + } + + private async pruneCheckpointFiles(retained: ReadonlySet): Promise { + let entries; + try { + entries = await readdir(this.checkpointsDir, { withFileTypes: true }); + } catch (error) { + if (isMissing(error)) return; + throw checkpointFsError('scan', this.checkpointsDir, error); + } + await Promise.all( + entries + .filter((entry) => entry.isFile() && entry.name.endsWith('.json')) + .filter((entry) => !retained.has(entry.name.slice(0, -'.json'.length))) + .map(async (entry) => { + try { + await unlink(join(this.checkpointsDir, entry.name)); + } catch (error) { + if (!isMissing(error)) throw error; + } + }), + ); + } +} + +function diffSnapshots(target: WorkspaceSnapshot, current: WorkspaceSnapshot): WorkspaceChange[] { + const currentRoots = rootMap(current); + const changes: WorkspaceChange[] = []; + for (const targetRoot of target.roots) { + const currentRoot = currentRoots.get(targetRoot.path); + if (currentRoot === undefined) continue; + const paths = new Set([...Object.keys(targetRoot.files), ...Object.keys(currentRoot.files)]); + for (const path of [...paths].toSorted()) { + const before = targetRoot.files[path]; + const now = currentRoot.files[path]; + if (sameFile(before, now)) continue; + changes.push({ + root: targetRoot.path, + path, + kind: before === undefined ? 'created' : now === undefined ? 'deleted' : 'modified', + }); + } + } + return changes; +} + +async function normalizeRoots(configuredRoots: readonly string[]): Promise { + const unique = new Set(); + for (const configured of configuredRoots) { + if (!isAbsolute(configured)) { + throw new WorkspaceCheckpointError(`Workspace root must be absolute: ${configured}`); + } + let canonical: string; + try { + canonical = await realpath(resolve(configured)); + } catch (error) { + throw checkpointFsError('resolve', configured, error); + } + unique.add(canonical); + } + const sorted = [...unique].toSorted((left, right) => left.length - right.length || left.localeCompare(right)); + return sorted.filter( + (candidate, index) => !sorted.slice(0, index).some((parent) => isPathInside(parent, candidate)), + ); +} + +async function loadIgnoreMatcher(root: string): Promise { + const matcher = createIgnore(); + for (const filename of ['.gitignore', '.ignore']) { + try { + matcher.add(await readFile(join(root, filename), 'utf8')); + } catch (error) { + if (!isMissing(error)) throw checkpointFsError('read', join(root, filename), error); + } + } + return matcher; +} + +function rootMap(snapshot: WorkspaceSnapshot): Map { + return new Map(snapshot.roots.map((root) => [root.path, root])); +} + +async function readCurrentRecord(root: string, relativePath: string): Promise { + const absolutePath = safeWorkspacePath(root, relativePath); + try { + const stats = await lstat(absolutePath); + if (!stats.isFile()) return undefined; + const contents = await readFile(absolutePath); + return { + hash: createHash('sha256').update(contents).digest('hex'), + mode: stats.mode & 0o777, + size: contents.byteLength, + }; + } catch (error) { + if (isMissing(error)) return undefined; + throw checkpointFsError('read', absolutePath, error); + } +} + +function safeWorkspacePath(root: string, relativePath: string): string { + if (relativePath.length === 0 || isAbsolute(relativePath)) { + throw new WorkspaceCheckpointError(`Invalid checkpoint path: ${relativePath}`); + } + const absolutePath = resolve(root, relativePath); + if (!isPathInside(root, absolutePath)) { + throw new WorkspaceCheckpointError(`Checkpoint path escapes workspace: ${relativePath}`); + } + return absolutePath; +} + +function isPathInside(parent: string, candidate: string): boolean { + const pathFromParent = relative(resolve(parent), resolve(candidate)); + return ( + pathFromParent.length > 0 && + pathFromParent !== '..' && + !pathFromParent.startsWith(`..${sep}`) && + !isAbsolute(pathFromParent) + ); +} + +async function assertNoSymlinkParent(root: string, relativePath: string): Promise { + const parts = relativePath.split('/').slice(0, -1); + let cursor = root; + for (const part of parts) { + cursor = join(cursor, part); + try { + const stats = await lstat(cursor); + if (stats.isSymbolicLink()) { + throw new WorkspaceCheckpointError(`Refusing to restore through symlink: ${cursor}`); + } + if (!stats.isDirectory()) { + throw new WorkspaceCheckpointError(`Restore parent is not a directory: ${cursor}`); + } + } catch (error) { + if (isMissing(error)) return; + if (error instanceof WorkspaceCheckpointError) throw error; + throw checkpointFsError('inspect', cursor, error); + } + } +} + +function sameFile(left: FileRecord | undefined, right: FileRecord | undefined): boolean { + if (left === undefined || right === undefined) return left === right; + return left.hash === right.hash && left.mode === right.mode && left.size === right.size; +} + +function sameStrings(left: readonly string[], right: readonly string[]): boolean { + return left.length === right.length && left.every((value, index) => value === right[index]); +} + +function isCheckpointIndex(value: unknown): value is CheckpointIndex { + if (typeof value !== 'object' || value === null) return false; + const candidate = value as { version?: unknown; checkpoints?: unknown }; + return ( + candidate.version === STORE_VERSION && + Array.isArray(candidate.checkpoints) && + candidate.checkpoints.every(isCheckpointRecord) + ); +} + +function isCheckpointRecord(value: unknown): value is CheckpointRecord { + if (typeof value !== 'object' || value === null) return false; + const candidate = value as { id?: unknown; createdAt?: unknown; roots?: unknown }; + return ( + typeof candidate.id === 'string' && + /^[a-f0-9-]{36}$/.test(candidate.id) && + typeof candidate.createdAt === 'string' && + Array.isArray(candidate.roots) && + candidate.roots.every((root) => typeof root === 'string' && isAbsolute(root)) + ); +} + +function isWorkspaceSnapshot(value: unknown): value is WorkspaceSnapshot { + if (typeof value !== 'object' || value === null || !('roots' in value)) return false; + const roots = value.roots; + return Array.isArray(roots) && roots.every(isRootSnapshot); +} + +function isRootSnapshot(value: unknown): value is RootSnapshot { + if (typeof value !== 'object' || value === null) return false; + const candidate = value as { path?: unknown; files?: unknown }; + if (typeof candidate.path !== 'string' || !isAbsolute(candidate.path)) return false; + if (typeof candidate.files !== 'object' || candidate.files === null || Array.isArray(candidate.files)) { + return false; + } + return Object.entries(candidate.files).every(([path, record]) => { + if (path.length === 0 || isAbsolute(path) || path.split('/').includes('..')) return false; + return isFileRecord(record); + }); +} + +function isFileRecord(value: unknown): value is FileRecord { + if (typeof value !== 'object' || value === null) return false; + const candidate = value as { hash?: unknown; mode?: unknown; size?: unknown }; + return ( + typeof candidate.hash === 'string' && + /^[a-f0-9]{64}$/.test(candidate.hash) && + typeof candidate.mode === 'number' && + Number.isInteger(candidate.mode) && + candidate.mode >= 0 && + candidate.mode <= 0o777 && + typeof candidate.size === 'number' && + Number.isSafeInteger(candidate.size) && + candidate.size >= 0 + ); +} + +function checkpointFsError(action: string, path: string, cause: unknown): WorkspaceCheckpointError { + const message = errorMessage(cause); + return new WorkspaceCheckpointError(`Failed to ${action} workspace checkpoint at ${path}: ${message}`); +} + +function errorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} + +function isMissing(error: unknown): boolean { + return errorCode(error) === 'ENOENT'; +} + +function isAlreadyExists(error: unknown): boolean { + return errorCode(error) === 'EEXIST'; +} + +function errorCode(error: unknown): string | undefined { + if (typeof error !== 'object' || error === null || !('code' in error)) return undefined; + return typeof error.code === 'string' ? error.code : undefined; +} + +function formatBytes(bytes: number): string { + return `${Math.ceil(bytes / (1024 * 1024))} MiB`; +} diff --git a/apps/kimi-code/test/tui/commands/registry.test.ts b/apps/kimi-code/test/tui/commands/registry.test.ts index a1964b5cbb..a865ee13cc 100644 --- a/apps/kimi-code/test/tui/commands/registry.test.ts +++ b/apps/kimi-code/test/tui/commands/registry.test.ts @@ -39,6 +39,7 @@ describe('built-in slash command registry', () => { expect(findBuiltInSlashCommand('mcp')?.name).toBe('mcp'); expect(findBuiltInSlashCommand('status')?.name).toBe('status'); expect(findBuiltInSlashCommand('usage')?.aliases).not.toContain('status'); + expect((findBuiltInSlashCommand('rewind') as KimiSlashCommand).availability).toBe('idle-only'); expect(findBuiltInSlashCommand('unknown')).toBeUndefined(); }); diff --git a/apps/kimi-code/test/tui/commands/resolve.test.ts b/apps/kimi-code/test/tui/commands/resolve.test.ts index 614553bb42..43d498105c 100644 --- a/apps/kimi-code/test/tui/commands/resolve.test.ts +++ b/apps/kimi-code/test/tui/commands/resolve.test.ts @@ -89,6 +89,11 @@ describe('resolveSlashCommandInput', () => { commandName: 'undo', reason: 'streaming', }); + expect(resolve('/rewind', { isStreaming: true })).toEqual({ + kind: 'blocked', + commandName: 'rewind', + reason: 'streaming', + }); expect(resolve('/reload', { isStreaming: true })).toEqual({ kind: 'blocked', commandName: 'reload', diff --git a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts index acd25f34ad..a9d62cf5ba 100644 --- a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts +++ b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts @@ -1,5 +1,5 @@ import { AsyncLocalStorage } from 'node:async_hooks'; -import { mkdtemp, rm, writeFile } from 'node:fs/promises'; +import { mkdir, mkdtemp, readFile, rm, writeFile } from 'node:fs/promises'; import { tmpdir } from 'node:os'; import { join, resolve } from 'node:path'; @@ -38,6 +38,7 @@ import { WelcomeComponent } from '#/tui/components/chrome/welcome'; import { ModelSelectorComponent } from '#/tui/components/dialogs/model-selector'; import { TabbedModelSelectorComponent } from '#/tui/components/dialogs/tabbed-model-selector'; import { UndoSelectorComponent } from '#/tui/components/dialogs/undo-selector'; +import { ChoicePickerComponent } from '#/tui/components/dialogs/choice-picker'; import { PluginInstallTrustConfirmComponent, PluginMcpSelectorComponent, @@ -2095,6 +2096,46 @@ command = "vim" ]); }); + it('rewinds workspace files and conversation history from a pre-turn checkpoint', async () => { + const parent = await mkdtemp(join(tmpdir(), 'kimi-code-rewind-flow-')); + tempDirs.push(parent); + const workDir = join(parent, 'workspace'); + const sessionDir = join(parent, 'session'); + await mkdir(workDir); + const source = join(workDir, 'source.ts'); + await writeFile(source, 'before\n'); + const session = makeSession({ + summary: { title: null, sessionDir, workDir, additionalDirs: [] }, + }); + const startup = makeStartupInput(); + const { driver } = await makeDriver(session, {}, { ...startup, workDir }); + + driver.handleUserInput('change source.ts'); + await vi.waitFor(() => { + expect(session.prompt).toHaveBeenCalledWith('change source.ts'); + }); + await writeFile(source, 'after\n'); + driver.state.appState.streamingPhase = 'idle'; + + driver.handleUserInput('/rewind 1'); + await vi.waitFor(() => { + expect(driver.state.editorContainer.children[0]).toBeInstanceOf(ChoicePickerComponent); + }); + const confirmation = driver.state.editorContainer.children[0] as ChoicePickerComponent; + confirmation.handleInput('\u001B[B'); + confirmation.handleInput('\r'); + + await vi.waitFor(() => { + expect(session.undoHistory).toHaveBeenCalledWith(1); + }); + await vi.waitFor(async () => { + expect(await readFile(source, 'utf8')).toBe('before\n'); + }); + expect(driver.state.transcriptEntries).not.toContainEqual( + expect.objectContaining({ kind: 'user', content: 'change source.ts' }), + ); + }); + it('keeps the transcript intact when undo RPC fails', async () => { const session = makeSession({ undoHistory: vi.fn(async () => { diff --git a/apps/kimi-code/test/tui/workspace-checkpoints.test.ts b/apps/kimi-code/test/tui/workspace-checkpoints.test.ts new file mode 100644 index 0000000000..21123cd7f2 --- /dev/null +++ b/apps/kimi-code/test/tui/workspace-checkpoints.test.ts @@ -0,0 +1,210 @@ +import { chmod, mkdir, mkdtemp, readFile, realpath, rm, symlink, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { createHash } from 'node:crypto'; + +import { afterEach, describe, expect, it } from 'vitest'; + +import { + WorkspaceCheckpointError, + WorkspaceCheckpointStore, +} from '../../src/tui/workspace-checkpoints'; + +const temporaryDirectories: string[] = []; + +afterEach(async () => { + await Promise.all(temporaryDirectories.splice(0).map((path) => rm(path, { recursive: true, force: true }))); +}); + +async function fixture(): Promise<{ root: string; session: string; store: WorkspaceCheckpointStore }> { + const parent = await mkdtemp(join(tmpdir(), 'kimi-rewind-')); + temporaryDirectories.push(parent); + const root = join(parent, 'workspace'); + const session = join(parent, 'session'); + await mkdir(root); + const canonicalRoot = await realpath(root); + return { root: canonicalRoot, session, store: new WorkspaceCheckpointStore(session, [canonicalRoot]) }; +} + +describe('WorkspaceCheckpointStore', () => { + it('restores modified, created, and deleted files without Git', async () => { + const { root, store } = await fixture(); + await mkdir(join(root, 'src')); + await writeFile(join(root, 'src', 'modified.ts'), 'before\n'); + await writeFile(join(root, 'deleted.txt'), 'bring me back\n'); + + await store.captureBeforeTurn(); + await writeFile(join(root, 'src', 'modified.ts'), 'after\n'); + await writeFile(join(root, 'created.txt'), 'remove me\n'); + await rm(join(root, 'deleted.txt')); + + const plan = await store.prepareRewind(1); + expect(plan.changes).toEqual([ + { root, path: 'created.txt', kind: 'created' }, + { root, path: 'deleted.txt', kind: 'deleted' }, + { root, path: 'src/modified.ts', kind: 'modified' }, + ]); + + await store.apply(plan); + expect(await readFile(join(root, 'src', 'modified.ts'), 'utf8')).toBe('before\n'); + expect(await readFile(join(root, 'deleted.txt'), 'utf8')).toBe('bring me back\n'); + await expect(readFile(join(root, 'created.txt'))).rejects.toMatchObject({ code: 'ENOENT' }); + await store.commit(plan); + expect(await store.availableCount()).toBe(0); + }); + + it('rewinds multiple turns to the oldest selected before-image', async () => { + const { root, store } = await fixture(); + const file = join(root, 'counter.txt'); + await writeFile(file, 'zero'); + await store.captureBeforeTurn(); + await writeFile(file, 'one'); + await store.captureBeforeTurn(); + await writeFile(file, 'two'); + + const plan = await store.prepareRewind(2); + await store.apply(plan); + expect(await readFile(file, 'utf8')).toBe('zero'); + }); + + it('fails closed if a file changes after preview and leaves every file untouched', async () => { + const { root, store } = await fixture(); + const first = join(root, 'first.txt'); + const second = join(root, 'second.txt'); + await writeFile(first, 'before first'); + await writeFile(second, 'before second'); + await store.captureBeforeTurn(); + await writeFile(first, 'agent first'); + await writeFile(second, 'agent second'); + const plan = await store.prepareRewind(1); + + await writeFile(second, 'external edit'); + await expect(store.apply(plan)).rejects.toThrow(/changed after the rewind preview/); + expect(await readFile(first, 'utf8')).toBe('agent first'); + expect(await readFile(second, 'utf8')).toBe('external edit'); + }); + + it('can roll the workspace forward when conversation undo fails', async () => { + const { root, store } = await fixture(); + const file = join(root, 'file.txt'); + await writeFile(file, 'before'); + await store.captureBeforeTurn(); + await writeFile(file, 'after'); + const plan = await store.prepareRewind(1); + + await store.apply(plan); + await store.rollback(plan); + expect(await readFile(file, 'utf8')).toBe('after'); + expect(await store.availableCount()).toBe(1); + }); + + it('respects ignore files and never traverses symlinks or dependency trees', async () => { + const { root, session, store } = await fixture(); + const outside = join(session, 'outside.txt'); + await mkdir(session, { recursive: true }); + await writeFile(outside, 'outside before'); + await writeFile(join(root, '.gitignore'), 'ignored.txt\n'); + await writeFile(join(root, 'ignored.txt'), 'ignored before'); + await mkdir(join(root, 'node_modules')); + await writeFile(join(root, 'node_modules', 'dependency.js'), 'dependency before'); + await symlink(outside, join(root, 'linked.txt')); + await store.captureBeforeTurn(); + + await writeFile(join(root, 'ignored.txt'), 'ignored after'); + await writeFile(join(root, 'node_modules', 'dependency.js'), 'dependency after'); + await writeFile(outside, 'outside after'); + const plan = await store.prepareRewind(1); + expect(plan.changes).toEqual([]); + }); + + it('tracks executable-mode changes', async () => { + const { root, store } = await fixture(); + const script = join(root, 'script.sh'); + await writeFile(script, '#!/bin/sh\n'); + await chmod(script, 0o644); + await store.captureBeforeTurn(); + await chmod(script, 0o755); + + const plan = await store.prepareRewind(1); + expect(plan.changes).toEqual([{ root, path: 'script.sh', kind: 'modified' }]); + await store.apply(plan); + const restored = await import('node:fs/promises').then(({ stat }) => stat(script)); + expect(restored.mode & 0o777).toBe(0o644); + }); + + it('rejects oversized workspaces instead of creating a partial checkpoint', async () => { + const { root, session } = await fixture(); + await writeFile(join(root, 'large.txt'), '12345'); + const store = new WorkspaceCheckpointStore(session, [root], { maxBytes: 4 }); + await expect(store.captureBeforeTurn()).rejects.toBeInstanceOf(WorkspaceCheckpointError); + expect(await store.availableCount()).toBe(0); + }); + + it('excludes its own content store when a session directory is inside the workspace', async () => { + const parent = await mkdtemp(join(tmpdir(), 'kimi-rewind-nested-')); + temporaryDirectories.push(parent); + const root = join(parent, 'workspace'); + const session = join(root, '.kimi', 'sessions', 'current'); + await mkdir(root); + await writeFile(join(root, 'source.txt'), 'unchanged'); + const store = new WorkspaceCheckpointStore(session, [root]); + + await store.captureBeforeTurn(); + const plan = await store.prepareRewind(1); + expect(plan.changes).toEqual([]); + }); + + it('refuses to restore through a parent directory replaced by a symlink', async () => { + const { root, session, store } = await fixture(); + const outside = join(session, 'outside'); + await mkdir(join(root, 'src')); + await writeFile(join(root, 'src', 'file.txt'), 'before'); + await store.captureBeforeTurn(); + await rm(join(root, 'src'), { recursive: true }); + await mkdir(outside, { recursive: true }); + await symlink(outside, join(root, 'src')); + + const plan = await store.prepareRewind(1); + await expect(store.apply(plan)).rejects.toThrow(/symlink/); + await expect(readFile(join(outside, 'file.txt'))).rejects.toMatchObject({ code: 'ENOENT' }); + }); + + it('bounds retained history and keeps the newest checkpoints aligned', async () => { + const { root, session } = await fixture(); + const store = new WorkspaceCheckpointStore(session, [root], { maxCheckpoints: 2 }); + const file = join(root, 'state.txt'); + await writeFile(file, 'zero'); + await store.captureBeforeTurn(); + await writeFile(file, 'one'); + await store.captureBeforeTurn(); + await writeFile(file, 'two'); + await store.captureBeforeTurn(); + await writeFile(file, 'three'); + + expect(await store.availableCount()).toBe(2); + const plan = await store.prepareRewind(2); + await store.apply(plan); + expect(await readFile(file, 'utf8')).toBe('one'); + }); + + it('rolls back a partially applied restore when a later checkpoint blob is corrupt', async () => { + const { root, session, store } = await fixture(); + const first = join(root, 'a.txt'); + const second = join(root, 'b.txt'); + await writeFile(first, 'before a'); + await writeFile(second, 'before b'); + await store.captureBeforeTurn(); + await writeFile(first, 'after a'); + await writeFile(second, 'after b'); + const plan = await store.prepareRewind(1); + const secondBeforeHash = createHash('sha256').update('before b').digest('hex'); + await writeFile( + join(session, 'workspace-checkpoints', 'blobs', secondBeforeHash), + 'corrupt', + ); + + await expect(store.apply(plan)).rejects.toThrow(/corrupt/); + expect(await readFile(first, 'utf8')).toBe('after a'); + expect(await readFile(second, 'utf8')).toBe('after b'); + }); +}); diff --git a/docs/en/reference/slash-commands.md b/docs/en/reference/slash-commands.md index f4aa141714..3f34566f15 100644 --- a/docs/en/reference/slash-commands.md +++ b/docs/en/reference/slash-commands.md @@ -34,6 +34,7 @@ Some commands are only available in the idle state. Executing these commands whi | `/title []` | `/rename` | Without arguments, display the current session title; with an argument, set a new title (max 200 characters) | Yes | | `/compact []` | — | Compact the current conversation context to free up token usage; an optional custom instruction can hint to the model what to preserve | No | | `/undo []` | — | Undo recent prompts from the active context. Without a count, opens a selector; with a count, undoes that many prompts. Prompts before the last compaction cannot be undone. Undoing also rolls back the todo list and plan mode state produced by those prompts (code changes are not reverted) | No | +| `/rewind []` | — | Rewind recent prompts and tracked workspace files to their pre-prompt checkpoint. Without a count, opens a selector; with a count, prepares that many prompts for rewind. A file preview and explicit confirmation are required | No | | `/reload` | — | Reload the current session and apply the latest `config.toml` settings (providers, models, etc.) and `tui.toml` UI preferences, without restarting the CLI | No | | `/reload-tui` | — | Reload only the `tui.toml` UI preferences (theme, editor, notifications, etc.) without rebuilding the session | Yes | | `/init` | — | Analyze the current codebase and generate `AGENTS.md` | No | @@ -43,6 +44,14 @@ Some commands are only available in the idle state. Executing these commands whi | `/add-dir []` | — | Add an extra workspace directory to the current session. Run without a path (or with `list`) to list configured directories. When adding, choose whether to remember the directory for the project in `.kimi-code/local.toml` | No | | `/web` | — | Open the current session in the web UI: pick a running server to connect to, or start a new foreground server after the TUI exits. See [`kimi web`](./kimi-command.md#kimi-web) | Yes | +### Workspace rewind + +`/rewind` is the filesystem-restoring counterpart to `/undo`. Before each user-initiated prompt, skill activation, or plugin command, the TUI stores a pre-prompt file snapshot under the session directory; identical file contents are stored only once. The command compares that checkpoint with the current workspace and previews which files it will delete, restore, or replace. If any affected file changes after the preview opens, the operation aborts before conversation history is changed. + +The checkpoint covers the primary workspace and additional directories. It does not follow symbolic links, and excludes version-control metadata, `node_modules`, and paths matched by the workspace root's `.gitignore` or `.ignore`. A checkpoint is skipped entirely above 50,000 files, 512 MiB total content, or 64 MiB for one file; the prompt still runs, but older checkpoints are invalidated so that a later rewind cannot target the wrong turn. Each session retains the latest 20 checkpoints. Checkpoints begin with prompts submitted by a version that supports `/rewind`; earlier session history has no file snapshot. + +Workspace snapshots cannot tell whether a tracked change came from Kimi Code, a formatter, or the user. Rewinding therefore restores the entire tracked workspace delta since the selected prompt, not only tool-authored edits. Review the preview before confirming. The content store is local to the session, is excluded from session and debug exports, and does not require or modify Git history. + ## Modes & Run Control | Command | Alias | Description | Always available | diff --git a/docs/zh/reference/slash-commands.md b/docs/zh/reference/slash-commands.md index 4a8b24451f..5a79c356b3 100644 --- a/docs/zh/reference/slash-commands.md +++ b/docs/zh/reference/slash-commands.md @@ -34,6 +34,7 @@ | `/title []` | `/rename` | 不带参数时显示当前会话标题;带参数时设置为新标题(最长 200 字符) | 是 | | `/compact []` | — | 压缩当前对话上下文,释放 token 占用;可附带自定义指令,提示模型压缩时保留哪些信息 | 否 | | `/undo []` | — | 从当前上下文撤销最近的提示词。不带数量时打开选择器;带数量时撤销对应条数。最后一次上下文压缩之前的提示词不能撤销。撤销会一并回滚这些提示词产生的 todo 列表和计划模式状态(不回滚代码改动) | 否 | +| `/rewind []` | — | 同时将最近的提示词和已跟踪工作区文件恢复到提示词执行前的检查点。不带数量时打开选择器;带数量时准备回退对应条数。执行前必须查看文件预览并明确确认 | 否 | | `/init` | — | 分析当前代码库并生成 `AGENTS.md` | 否 | | `/export-md []` | `/export` | 将当前会话导出为 Markdown 文件 | 否 | | `/export-debug-zip` | — | 将当前会话导出为调试用 ZIP 压缩包(与 [`kimi export`](./kimi-command.md#kimi-export) 行为一致) | 否 | @@ -41,6 +42,14 @@ | `/add-dir []` | — | 为当前会话添加额外的工作目录。不带路径(或传入 `list`)运行时列出已配置的目录。添加时可选择是否将目录记入项目的 `.kimi-code/local.toml` | 否 | | `/web` | — | 在 web UI 中打开当前会话:选择一个运行中的实例进行连接,或在 TUI 退出后新开一个前台服务器。参见 [`kimi web`](./kimi-command.md#kimi-web) | 是 | +### 工作区回退 + +`/rewind` 是会恢复文件系统的 `/undo`。每次用户主动发送提示词、激活 Skill 或运行插件命令前,TUI 都会在会话目录中保存提示词执行前的文件快照;内容相同的文件只保存一份。执行命令时会比较该检查点与当前工作区,预览将要删除、还原或替换的文件。如果打开预览后任何受影响文件再次变化,操作会在修改对话历史前中止。 + +检查点覆盖主工作区和附加目录,但不会跟随符号链接,并排除版本控制元数据、`node_modules`,以及工作区根目录 `.gitignore` 或 `.ignore` 匹配的路径。文件数超过 50,000、内容总量超过 512 MiB,或单个文件超过 64 MiB 时,本轮检查点会完全跳过;提示词仍会执行,但旧检查点会失效,避免后续回退错误地对应到其他轮次。每个会话保留最近 20 个检查点。只有升级到支持 `/rewind` 的版本后新提交的提示词才有文件快照,旧会话历史不能恢复文件。 + +工作区快照无法判断一项改动来自 Kimi Code、格式化工具还是用户。因此,回退的是所选提示词之后全部已跟踪的工作区差异,而不仅是工具写入的内容;确认前应检查预览。内容存储只存在于本地会话目录,不会进入会话导出或调试导出,也不依赖或修改 Git 历史。 + ## 模式与运行控制 | 命令 | 别名 | 说明 | 随时可用 | diff --git a/packages/agent-core-v2/src/app/sessionExport/sessionExportService.ts b/packages/agent-core-v2/src/app/sessionExport/sessionExportService.ts index 561d3c5585..303157c3e9 100644 --- a/packages/agent-core-v2/src/app/sessionExport/sessionExportService.ts +++ b/packages/agent-core-v2/src/app/sessionExport/sessionExportService.ts @@ -47,6 +47,7 @@ const SESSION_LOG_REL = 'logs/kimi-code.log'; const GLOBAL_LOG_REL = 'logs/global/kimi-code.log'; const WEB_LOG_REL = 'logs/kimi-web.jsonl'; const DESKTOP_LOG_REL = 'logs/kimi-desktop.log'; +const NON_EXPORTABLE_SESSION_DIRECTORIES = ['workspace-checkpoints'] as const; export class SessionExportService implements ISessionExportService { declare readonly _serviceBrand: undefined; @@ -205,7 +206,9 @@ export async function exportSessionDirectory(input: { if (input.desktopLogPath !== undefined) { desktopSource = await openOptionalZipSource(input.desktopLogPath, input.signal); } - const sessionFiles = await collectFilesRecursive(sessionDir); + const sessionFiles = await collectFilesRecursive(sessionDir, { + excludeTopLevelDirectories: NON_EXPORTABLE_SESSION_DIRECTORIES, + }); if (sessionFiles.length === 0 && sessionLogSource === undefined) { throw new Error2( ErrorCodes.SESSION_EXPORT_NOT_FOUND, diff --git a/packages/agent-core-v2/src/app/sessionExport/zip.ts b/packages/agent-core-v2/src/app/sessionExport/zip.ts index 7dbc97ae7a..10039d9cca 100644 --- a/packages/agent-core-v2/src/app/sessionExport/zip.ts +++ b/packages/agent-core-v2/src/app/sessionExport/zip.ts @@ -23,12 +23,17 @@ import { } from './file-source'; import type { ExportSessionManifest } from './sessionExport'; -export async function collectFilesRecursive(root: string): Promise { +export async function collectFilesRecursive( + root: string, + options: { readonly excludeTopLevelDirectories?: readonly string[] } = {}, +): Promise { + const excluded = new Set(options.excludeTopLevelDirectories ?? []); try { const entries = await readdir(root, { recursive: true, withFileTypes: true }); return entries .filter((entry) => entry.isFile()) .map((entry) => join(entry.parentPath, entry.name)) + .filter((path) => !excluded.has(relative(root, path).split(/[\\/]/, 1)[0] ?? '')) .toSorted((a, b) => a.localeCompare(b)); } catch (error) { if (!isMissingPath(error)) throw error; diff --git a/packages/agent-core-v2/test/app/sessionExport/sessionExport.test.ts b/packages/agent-core-v2/test/app/sessionExport/sessionExport.test.ts index 34a49c143b..31a268c210 100644 --- a/packages/agent-core-v2/test/app/sessionExport/sessionExport.test.ts +++ b/packages/agent-core-v2/test/app/sessionExport/sessionExport.test.ts @@ -92,6 +92,8 @@ describe('sessionExport', () => { await mkdir(join(sessionDir, 'agents', 'main'), { recursive: true }); await mkdir(join(sessionDir, 'logs'), { recursive: true }); await writeFile(join(sessionDir, 'state.json'), '{}\n', 'utf-8'); + await mkdir(join(sessionDir, 'workspace-checkpoints', 'blobs'), { recursive: true }); + await writeFile(join(sessionDir, 'workspace-checkpoints', 'blobs', 'source-content'), 'private source'); await writeFile(join(sessionDir, 'logs', 'kimi-code.log'), '{"msg":"session"}\n', 'utf-8'); await writeFile( join(sessionDir, 'agents', 'main', 'wire.jsonl'), @@ -130,6 +132,7 @@ describe('sessionExport', () => { 'state.json', 'logs/global/kimi-code.log', ]); + expect(result.entries.some((entry) => entry.startsWith('workspace-checkpoints/'))).toBe(false); expect(result.manifest).toMatchObject({ sessionId: 'ses_demo', kimiCodeVersion: '1.0.0-test', diff --git a/packages/agent-core/src/session/export/session-export.ts b/packages/agent-core/src/session/export/session-export.ts index 056374cdfb..3368dd81c0 100644 --- a/packages/agent-core/src/session/export/session-export.ts +++ b/packages/agent-core/src/session/export/session-export.ts @@ -14,6 +14,7 @@ import type { ExportSessionPayload, ExportSessionResult, SessionSummary } from ' const SESSION_LOG_REL = 'logs/kimi-code.log'; const GLOBAL_LOG_REL = 'logs/global/kimi-code.log'; +const NON_EXPORTABLE_SESSION_DIRECTORIES = ['workspace-checkpoints'] as const; export async function exportSessionDirectory(input: { readonly request: ExportSessionPayload; @@ -22,7 +23,9 @@ export async function exportSessionDirectory(input: { readonly globalLogPath?: string | undefined; }): Promise { const sessionDir = input.summary.sessionDir; - const sessionFiles = await collectFilesRecursive(sessionDir); + const sessionFiles = await collectFilesRecursive(sessionDir, { + excludeTopLevelDirectories: NON_EXPORTABLE_SESSION_DIRECTORIES, + }); if (sessionFiles.length === 0) { throw new KimiError(ErrorCodes.SESSION_EXPORT_NOT_FOUND, `Session "${input.summary.id}" has no exportable directory at "${sessionDir}"`, { details: { sessionId: input.summary.id, sessionDir }, diff --git a/packages/agent-core/src/session/export/zip.ts b/packages/agent-core/src/session/export/zip.ts index 51507a28a9..0261dded70 100644 --- a/packages/agent-core/src/session/export/zip.ts +++ b/packages/agent-core/src/session/export/zip.ts @@ -7,12 +7,17 @@ import { pipeline } from 'node:stream/promises'; import type { ExportSessionManifest } from '#/rpc/core-api'; import { ZipFile } from 'yazl'; -export async function collectFilesRecursive(root: string): Promise { +export async function collectFilesRecursive( + root: string, + options: { readonly excludeTopLevelDirectories?: readonly string[] } = {}, +): Promise { + const excluded = new Set(options.excludeTopLevelDirectories ?? []); try { const entries = await readdir(root, { recursive: true, withFileTypes: true }); return entries .filter((entry) => entry.isFile()) .map((entry) => join(entry.parentPath, entry.name)) + .filter((path) => !excluded.has(relative(root, path).split(/[\\/]/, 1)[0] ?? '')) .toSorted((a, b) => a.localeCompare(b)); } catch { return []; diff --git a/packages/node-sdk/test/export-session.test.ts b/packages/node-sdk/test/export-session.test.ts index 358e1fac4a..a2ceaf119b 100644 --- a/packages/node-sdk/test/export-session.test.ts +++ b/packages/node-sdk/test/export-session.test.ts @@ -122,6 +122,8 @@ describe('exportSessionDirectory', () => { ); await writeFile(join(sessionDir, 'state.json'), JSON.stringify({ session_id: sid }), 'utf-8'); await writeFile(join(sessionDir, 'subagents', 'a.txt'), 'child', 'utf-8'); + await mkdir(join(sessionDir, 'workspace-checkpoints', 'blobs'), { recursive: true }); + await writeFile(join(sessionDir, 'workspace-checkpoints', 'blobs', 'source-content'), 'private source'); const outputPath = join(tmp, 'out.zip'); const result = await exportSessionDirectory({ @@ -158,6 +160,7 @@ describe('exportSessionDirectory', () => { expect(entries.get('state.json')?.toString('utf-8')).toContain(sid); expect(entries.get('subagents/a.txt')?.toString('utf-8')).toBe('child'); expect([...entries.keys()].some((name) => name.includes(tmp))).toBe(false); + expect([...entries.keys()].some((name) => name.startsWith('workspace-checkpoints/'))).toBe(false); const manifest = JSON.parse(entries.get('manifest.json')!.toString('utf-8')) as { sessionId: string; diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 65ce665033..672a68944e 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -125,6 +125,9 @@ importers: commander: specifier: ^13.1.0 version: 13.1.0 + ignore: + specifier: ^5.3.2 + version: 5.3.2 jimp: specifier: ^1.6.1 version: 1.6.1