Rewind the conversation on /undo, staged so a redo can lift it (part 3 of #944)

/undo reverted files while the model kept its memory of the turn, so the agent
would look at the disk, not find the file it believed it had just written, and
write it again - the undo undid itself.

A turn now records where it started in both memories, and an undo stages a cut at
that point: nothing is deleted, the model's history is slice-copied on its way
into createRunConfig (the one place a run is built), and /redo lifts the mark. A
turn that runs afterwards absorbs the cut and clears it.

The picker half - hiding the reverted tail on screen, prefilling the composer,
and choosing code-only vs code-and-conversation - is the next PR.

Author: KazenDev

cli/src/app.tsx

@@ -213,7 +213,11 @@ export const App = ({
       closeUndoHistory()
       try {
         const message =
-          (await undoToRecord(getCurrentChatId(), projectRoot, recordId)) ??
+          (await undoToRecord(getCurrentChatId(), projectRoot, recordId, {
+            // Both halves by default: reverting the files while the model keeps
+            // remembering the turn is what made it write them back again.
+            conversation: true,
+          })) ??
           'Could not undo — the snapshot store is unavailable.'
         useChatStore.getState().setMessages((prev) => [
           ...prev,

cli/src/hooks/use-send-message.ts

@@ -8,7 +8,14 @@ import {
   setCurrentChatId,
 } from '../project-files'
 import { beginUndoTurn, isLatestUndoTurn } from '../state/undo-guards'
-import { recordUndoEntry } from '../state/undo-store'
+import { applyRewind, lastUserMessageIndex } from '../state/undo-rewind'
+import {
+  clearRewindBoundary,
+  getRewindBoundary,
+  recordUndoEntry,
+} from '../state/undo-store'
+
+import type { RewindAnchor } from '../state/undo-rewind'
 import { createStreamController } from './stream-state'
 import { useChatStore } from '../state/chat-store'
 import {
@@ -589,6 +596,9 @@ export const useSendMessage = ({
       // can run after a newer turn has already started (Esc releases the chain
       // lock first), and by then this turn's diff is no longer only its own.
       let undoTurnStamp: number | null = null
+      // Where the conversation stood when this turn started, recorded with the
+      // entry so an undo can rewind the model as well as the worktree.
+      let undoTurnAnchor: RewindAnchor | null = null
 
       // Checkpoint the turn to disk immediately so that killing the process
       // (closed terminal, crash) can't lose the user's prompt, then keep the
@@ -666,6 +676,11 @@ export const useSendMessage = ({
                 priorByok.revision === selectedByok.revision,
             )
           : !byok
+        // A staged conversation rewind is applied here, at the one place a run
+        // is built: the model's history is what this argument becomes. Nothing
+        // was deleted to stage it, so /redo can still lift it, and a chat that
+        // never rewound passes the very same state through.
+        const rewindBoundary = getRewindBoundary(runChatId)
         const runConfig = createRunConfig({
           logger,
           agent: resolvedAgent,
@@ -674,7 +689,10 @@ export const useSendMessage = ({
           // A persisted run has a non-secret source pin. Never resume its
           // transcript after switching to Freebuff or another BYOK revision.
           previousRunState: canResumePreviousRun
-            ? previousRunStateRef.current
+            ? applyRewind(
+                previousRunStateRef.current,
+                rewindBoundary?.historyLength ?? null,
+              )
             : null,
           agentDefinitions,
           eventHandlerState,
@@ -758,6 +776,19 @@ export const useSendMessage = ({
             undoSnapshotHash = await trackSnapshot(getProjectRoot())
             if (undoSnapshotHash) {
               undoTurnStamp = beginUndoTurn(runChatId)
+              // Read here, as the turn starts and before its own run can append
+              // anything, so "the last prompt" is this turn's.
+              const transcriptIndex = lastUserMessageIndex(
+                useChatStore.getState().messages,
+              )
+              if (transcriptIndex >= 0) {
+                undoTurnAnchor = {
+                  historyLength:
+                    previousRunStateRef.current?.sessionState?.mainAgentState
+                      .messageHistory.length ?? 0,
+                  transcriptIndex,
+                }
+              }
             }
           }
         } catch (error) {
@@ -796,6 +827,11 @@ export const useSendMessage = ({
           // and JSON.stringify of the (unbounded) transcript through proxy
           // traps is several times slower.
           saveChatState(runState, useChatStore.getState().messages, runChatDir)
+
+          // This run was built from the cut history, so the cut is baked into
+          // the state just persisted and the staged mark has done its job. An
+          // interrupted or failed run keeps it: nothing absorbed the cut.
+          if (rewindBoundary) clearRewindBoundary(runChatId)
         }
         handleRunCompletion({
           runState,
@@ -910,6 +946,7 @@ export const useSendMessage = ({
                   hashBefore: undoSnapshotHash,
                   files: undoFiles,
                   message: content,
+                  ...(undoTurnAnchor ? { anchor: undoTurnAnchor } : {}),
                 })
               }
             } catch (error) {

cli/src/state/__tests__/undo-rewind.test.ts

@@ -0,0 +1,108 @@
+import { describe, expect, test } from 'bun:test'
+
+import { anchorOf, applyRewind, lastUserMessageIndex } from '../undo-rewind'
+
+import type { ChatMessage } from '../../types/chat'
+import type { Message } from '@codebuff/common/types/messages/codebuff-message'
+import type { RunState } from '@codebuff/sdk'
+
+const message = (variant: 'user' | 'ai', content = 'x'): ChatMessage =>
+  ({
+    id: `${variant}-${content}`,
+    variant,
+    content,
+    timestamp: '2026-01-01T00:00:00.000Z',
+  }) as ChatMessage
+
+/** A message-shaped stand-in: only its identity matters to these helpers. */
+const fakeMessage = (id: string): Message => ({ id }) as unknown as Message
+
+/** Only the part of a run state these helpers read is real. */
+const runStateWithHistory = (ids: string[]): RunState =>
+  ({
+    output: { type: 'allMessages', value: ids },
+    traceSessionId: 'trace-1',
+    sessionState: {
+      mainAgentState: { messageHistory: ids.map(fakeMessage) },
+    },
+  }) as unknown as RunState
+
+describe('lastUserMessageIndex', () => {
+  test('finds the prompt a turn started from, not the newest ai reply', () => {
+    const messages = [
+      message('user', 'one'),
+      message('ai', 'reply'),
+      message('user', 'two'),
+      message('ai', 'reply'),
+      message('ai', 'still streaming'),
+    ]
+    expect(lastUserMessageIndex(messages)).toBe(2)
+  })
+
+  test('is -1 when the transcript holds no prompt yet', () => {
+    expect(lastUserMessageIndex([])).toBe(-1)
+    expect(lastUserMessageIndex([message('ai', 'hello')])).toBe(-1)
+  })
+})
+
+describe('anchorOf', () => {
+  test('reads an anchor a record carries', () => {
+    expect(anchorOf({ anchor: { historyLength: 4, transcriptIndex: 2 } })).toEqual(
+      { historyLength: 4, transcriptIndex: 2 },
+    )
+  })
+
+  test('an entry recorded before the anchor existed has none', () => {
+    expect(anchorOf({})).toBeNull()
+  })
+
+  test('refuses a nonsense anchor instead of cutting at a guess', () => {
+    expect(anchorOf({ anchor: { historyLength: -1, transcriptIndex: 0 } })).toBeNull()
+    expect(anchorOf({ anchor: { historyLength: 1.5, transcriptIndex: 0 } })).toBeNull()
+    expect(anchorOf({ anchor: { historyLength: 3, transcriptIndex: -2 } })).toBeNull()
+    expect(
+      anchorOf({
+        anchor: { historyLength: 3, transcriptIndex: Number.NaN },
+      }),
+    ).toBeNull()
+  })
+})
+
+describe('applyRewind', () => {
+  const original = runStateWithHistory(['m1', 'm2', 'm3', 'm4'])
+
+  test('cuts the model history back to where the turn started', () => {
+    const rewound = applyRewind(original, 2)
+    expect(rewound?.sessionState?.mainAgentState.messageHistory).toEqual([
+      fakeMessage('m1'),
+      fakeMessage('m2'),
+    ])
+  })
+
+  test('leaves the caller the full history it still persists', () => {
+    // The full state lives in a ref and is written to disk; a cut that mutated
+    // it would make /redo impossible even though nothing was ever deleted.
+    applyRewind(original, 1)
+    expect(original.sessionState?.mainAgentState.messageHistory).toHaveLength(4)
+  })
+
+  test('is a no-op, same reference, when there is nothing to cut', () => {
+    expect(applyRewind(original, null)).toBe(original)
+    expect(applyRewind(original, 4)).toBe(original)
+    expect(applyRewind(original, 9)).toBe(original)
+    expect(applyRewind(null, 2)).toBeNull()
+  })
+
+  test('an unstarted chat cuts to an empty history rather than throwing', () => {
+    // historyLength 0 is the first turn of a chat: everything the model had is
+    // that turn, so the cut leaves it with nothing to remember.
+    expect(
+      applyRewind(original, 0)?.sessionState?.mainAgentState.messageHistory,
+    ).toEqual([])
+  })
+
+  test('leaves a state with no history alone', () => {
+    const bare = { output: { type: 'allMessages', value: [] } } as unknown as RunState
+    expect(applyRewind(bare, 0)).toBe(bare)
+  })
+})

cli/src/state/__tests__/undo-store.test.ts

@@ -22,6 +22,8 @@ import {
   trackSnapshot,
 } from '../../utils/undo-snapshot'
 import {
+  clearRewindBoundary,
+  getRewindBoundary,
   loadUndoState,
   peekRedo,
   peekUndo,
@@ -30,10 +32,14 @@ import {
   pushRedo,
   pushUndo,
   recordUndoEntry,
+  redoToRecord,
+  setRewindBoundary,
   sweepAnchors,
   undoToRecord,
 } from '../undo-store'
 
+import type { RewindAnchor } from '../undo-rewind'
+
 const CHAT_ID = 'undo-store-test-chat'
 
 let projectDir: string
@@ -262,6 +268,128 @@ describe('snapshot anchors', () => {
   })
 })
 
+describe('the conversation rewind', () => {
+  /**
+   * The journal as persisted. Read raw on purpose: a staged cut is a promise
+   * about what is on disk, and a test that only asked the getter could pass
+   * while the write silently dropped it.
+   */
+  const readJournal = (): {
+    rewind?: {
+      recordId: string
+      historyLength: number
+      transcriptIndex: number
+      createdAt: string
+    }
+  } =>
+    JSON.parse(
+      readFileSync(
+        path.join(getProjectDataDir(), 'chats', CHAT_ID, 'undo.json'),
+        'utf8',
+      ),
+    )
+
+  /** A turn that rewrites a file, recorded with the anchor it captured. */
+  const recordTurn = async (name: string, anchor?: RewindAnchor) => {
+    const file = path.join(projectDir, name)
+    writeFileSync(file, 'the user wrote this\n')
+    const hash = await trackSnapshot(projectDir)
+    writeFileSync(file, 'the agent changed it\n')
+    recordUndoEntry(CHAT_ID, {
+      hashBefore: hash!,
+      files: [name],
+      message: 'a turn worth rewinding',
+      ...(anchor ? { anchor } : {}),
+    })
+    return { file, id: peekUndo(CHAT_ID)!.id }
+  }
+
+  test('an undo takes the turn out of the model conversation too', async () => {
+    const { file, id } = await recordTurn('rewritten.txt', {
+      historyLength: 3,
+      transcriptIndex: 1,
+    })
+
+    const message = await undoToRecord(CHAT_ID, projectDir, id, {
+      conversation: true,
+    })
+
+    // Both halves: the file is back, and the model no longer remembers writing
+    // it — which is what stops it from writing it a second time.
+    expect(readFileSync(file, 'utf8')).toBe('the user wrote this\n')
+    expect(message).toContain('↶')
+    expect(readJournal().rewind).toEqual({
+      recordId: id,
+      historyLength: 3,
+      transcriptIndex: 1,
+      createdAt: expect.any(String),
+    })
+  })
+
+  test('a file-only undo stages no conversation cut', async () => {
+    const { id } = await recordTurn('files-only.txt', {
+      historyLength: 2,
+      transcriptIndex: 0,
+    })
+
+    const message = await undoToRecord(CHAT_ID, projectDir, id)
+
+    expect(message).not.toContain('↶')
+    expect(readJournal().rewind).toBeUndefined()
+    expect(getRewindBoundary(CHAT_ID)).toBeNull()
+  })
+
+  test('an entry with no anchor reverts its files and does not cut', async () => {
+    const { file, id } = await recordTurn('unanchored.txt')
+
+    const message = await undoToRecord(CHAT_ID, projectDir, id, {
+      conversation: true,
+    })
+
+    // A guess is worse than no cut: the file half still did its job.
+    expect(readFileSync(file, 'utf8')).toBe('the user wrote this\n')
+    expect(message).not.toContain('↶')
+    expect(getRewindBoundary(CHAT_ID)).toBeNull()
+  })
+
+  test('a turn recorded while a cut is staged does not lift it', async () => {
+    setRewindBoundary(CHAT_ID, {
+      recordId: 'r1',
+      historyLength: 4,
+      transcriptIndex: 2,
+      createdAt: new Date().toISOString(),
+    })
+
+    recordUndoEntry(CHAT_ID, {
+      hashBefore: 'later',
+      files: ['later.txt'],
+      message: 'a newer turn',
+      anchor: { historyLength: 6, transcriptIndex: 4 },
+    })
+
+    expect(getRewindBoundary(CHAT_ID)?.historyLength).toBe(4)
+  })
+
+  test('a redo lifts the staged cut', async () => {
+    const { id } = await recordTurn('redone.txt', {
+      historyLength: 1,
+      transcriptIndex: 0,
+    })
+    await undoToRecord(CHAT_ID, projectDir, id, { conversation: true })
+    expect(getRewindBoundary(CHAT_ID)).not.toBeNull()
+
+    await redoToRecord(CHAT_ID, projectDir, peekRedo(CHAT_ID)!.id)
+
+    expect(getRewindBoundary(CHAT_ID)).toBeNull()
+  })
+
+  test('lifting a cut that is not there is safe', () => {
+    expect(getRewindBoundary(CHAT_ID)).toBeNull()
+    clearRewindBoundary(CHAT_ID)
+    expect(getRewindBoundary(CHAT_ID)).toBeNull()
+  })
+})
+
 describe('corrupt file handling', () => {
   test('loads an empty state for a nonexistent chat', () => {
     expect(loadUndoState('never-existed').undoStack).toEqual([])