fix invalid code redo

Author: phantom5099

packages/codingcode/src/checkpoint/checkpoint-service.ts

@@ -170,51 +170,83 @@ function getCompletedTurnsFor(sg: ShadowGit, sessionId: string): number[] {
   return ids;
 }
 
+interface RestorePlan {
+  throughTurnId: number;
+  baseTurnId: number;
+  affectedTurns: number[];
+  baseline: string;
+}
+
+function getTurnRestorePlan(sg: ShadowGit, sessionId: string, turnId: number): RestorePlan | null {
+  const baseline = sg.findCommitByMessage(commitMsg(sessionId, turnId, 'baseline'));
+  const final = sg.findCommitByMessage(commitMsg(sessionId, turnId, 'final'));
+  if (!baseline || !final) return null;
+  return {
+    throughTurnId: turnId,
+    baseTurnId: turnId,
+    affectedTurns: [],
+    baseline,
+  };
+}
+
+function getRollbackToTurnPlan(
+  sg: ShadowGit,
+  sessionId: string,
+  throughTurnId: number
+): RestorePlan | null {
+  const completedTurns = getCompletedTurnsFor(sg, sessionId);
+  const affectedTurns = completedTurns.filter((id) => id >= throughTurnId);
+  if (affectedTurns.length === 0) return null;
+
+  const baseline = sg.findCommitByMessage(commitMsg(sessionId, throughTurnId, 'baseline'));
+  if (!baseline) return null;
+
+  return {
+    throughTurnId,
+    baseTurnId: throughTurnId,
+    affectedTurns,
+    baseline,
+  };
+}
+
 function revertFilesImpl(
   projectPath: string,
   sessionId: string,
-  turnId: number,
+  plan: RestorePlan,
   selectedFiles: string[],
   action: CodeRestoreEntry['action'],
   sg: ShadowGit
 ): CodeRollbackResult {
   if (selectedFiles.length === 0) {
     return {
       reverted: false,
-      throughTurnId: turnId,
-      baseTurnId: turnId,
-      affectedTurns: [],
+      throughTurnId: plan.throughTurnId,
+      baseTurnId: plan.baseTurnId,
+      affectedTurns: plan.affectedTurns,
       selectedFiles: [],
       restoreEntry: null,
     };
   }
 
-  const baseline = sg.findCommitByMessage(commitMsg(sessionId, turnId, 'baseline'));
-  if (!baseline)
-    return {
-      reverted: false,
-      throughTurnId: turnId,
-      baseTurnId: turnId,
-      affectedTurns: [],
-      selectedFiles: [],
-      restoreEntry: null,
-    };
-
   sg.lock();
   try {
     let safetyCommit: string;
     const existingEntry = readRestoreEntry(sg.gitDir, sessionId);
 
-    if (existingEntry && existingEntry.throughTurnId === turnId && existingEntry.safetyCommit) {
+    if (
+      existingEntry &&
+      existingEntry.throughTurnId === plan.throughTurnId &&
+      existingEntry.safetyCommit
+    ) {
       safetyCommit = existingEntry.safetyCommit;
     } else {
-      safetyCommit = sg.commit(commitMsg(sessionId, turnId, 'revert-safety'));
+      safetyCommit = sg.commit(commitMsg(sessionId, plan.throughTurnId, 'revert-safety'));
     }
 
-    sg.checkoutFiles(baseline, selectedFiles);
+    sg.checkoutFiles(plan.baseline, selectedFiles);
 
     const combinedFiles =
-      existingEntry && existingEntry.throughTurnId === turnId
+      existingEntry && existingEntry.throughTurnId === plan.throughTurnId
         ? [
             ...new Map(
               [...existingEntry.selectedFiles, ...selectedFiles].map((f) => [
@@ -227,14 +259,14 @@ function revertFilesImpl(
 
     const entry: CodeRestoreEntry = {
       id: createHash('sha256')
-        .update(`${sessionId}-${turnId}-${Date.now()}`)
+        .update(`${sessionId}-${plan.throughTurnId}-${Date.now()}`)
         .digest('hex')
         .slice(0, 12),
       sessionId,
       action,
-      throughTurnId: turnId,
-      baseTurnId: turnId,
-      affectedTurns: [],
+      throughTurnId: plan.throughTurnId,
+      baseTurnId: plan.baseTurnId,
+      affectedTurns: plan.affectedTurns,
       selectedFiles: combinedFiles,
       safetyCommit,
       timestamp: new Date().toISOString(),
@@ -243,9 +275,9 @@ function revertFilesImpl(
 
     return {
       reverted: true,
-      throughTurnId: turnId,
-      baseTurnId: turnId,
-      affectedTurns: [],
+      throughTurnId: plan.throughTurnId,
+      baseTurnId: plan.baseTurnId,
+      affectedTurns: plan.affectedTurns,
       selectedFiles: combinedFiles,
       restoreEntry: entry,
     };
@@ -440,7 +472,17 @@ export class CheckpointService extends Effect.Service<CheckpointService>()('Chec
         file: string
       ): CodeRollbackResult => {
         const sg = ensure(projectPath);
-        return revertFilesImpl(projectPath, sessionId, turnId, [file], 'checkpoint-file', sg);
+        const plan = getTurnRestorePlan(sg, sessionId, turnId);
+        if (!plan)
+          return {
+            reverted: false,
+            throughTurnId: turnId,
+            baseTurnId: turnId,
+            affectedTurns: [],
+            selectedFiles: [],
+            restoreEntry: null,
+          };
+        return revertFilesImpl(projectPath, sessionId, plan, [file], 'checkpoint-file', sg);
       },
 
       revertCheckpointFiles: (
@@ -450,7 +492,17 @@ export class CheckpointService extends Effect.Service<CheckpointService>()('Chec
         files: string[]
       ): CodeRollbackResult => {
         const sg = ensure(projectPath);
-        return revertFilesImpl(projectPath, sessionId, turnId, files, 'checkpoint-files', sg);
+        const plan = getTurnRestorePlan(sg, sessionId, turnId);
+        if (!plan)
+          return {
+            reverted: false,
+            throughTurnId: turnId,
+            baseTurnId: turnId,
+            affectedTurns: [],
+            selectedFiles: [],
+            restoreEntry: null,
+          };
+        return revertFilesImpl(projectPath, sessionId, plan, files, 'checkpoint-files', sg);
       },
 
       // ---- B3: revertCheckpointAgentFiles / revertCheckpointAllFiles ----
@@ -481,10 +533,20 @@ export class CheckpointService extends Effect.Service<CheckpointService>()('Chec
             selectedFiles: [],
             restoreEntry: null,
           };
+        const plan = getTurnRestorePlan(sg, sessionId, turnId);
+        if (!plan)
+          return {
+            reverted: false,
+            throughTurnId: turnId,
+            baseTurnId: null,
+            affectedTurns: [],
+            selectedFiles: [],
+            restoreEntry: null,
+          };
         return revertFilesImpl(
           projectPath,
           sessionId,
-          turnId,
+          plan,
           changes.agentModified,
           'checkpoint-agent',
           sg
@@ -518,7 +580,24 @@ export class CheckpointService extends Effect.Service<CheckpointService>()('Chec
             selectedFiles: [],
             restoreEntry: null,
           };
-        return revertFilesImpl(projectPath, sessionId, turnId, all, 'checkpoint-all', sg);
+        const plan = getTurnRestorePlan(sg, sessionId, turnId);
+        if (!plan)
+          return {
+            reverted: false,
+            throughTurnId: turnId,
+            baseTurnId: null,
+            affectedTurns: [],
+            selectedFiles: [],
+            restoreEntry: null,
+          };
+        return revertFilesImpl(
+          projectPath,
+          sessionId,
+          plan,
+          all,
+          'checkpoint-all',
+          sg
+        );
       },
 
       // ---- B4: previewRollbackDiff ----
@@ -529,20 +608,18 @@ export class CheckpointService extends Effect.Service<CheckpointService>()('Chec
         throughTurnId: number
       ): RollbackPreviewDiff => {
         const sg = ensure(projectPath);
-        const completedTurns = getCompletedTurnsFor(sg, sessionId);
-        const affectedTurns = completedTurns.filter((id) => id > throughTurnId);
-        if (affectedTurns.length === 0) {
+        const plan = getRollbackToTurnPlan(sg, sessionId, throughTurnId);
+        if (!plan) {
           return { throughTurnId, baseTurnId: null, affectedTurns: [], diff: '' };
         }
 
-        const baseTurnId = affectedTurns[0]! - 1;
-        const baseline = sg.findCommitByMessage(commitMsg(sessionId, baseTurnId, 'baseline'));
-        if (!baseline) {
-          return { throughTurnId, baseTurnId, affectedTurns, diff: '' };
-        }
-
-        const result = sg.git('diff', baseline);
-        return { throughTurnId, baseTurnId, affectedTurns, diff: result.stdout };
+        const result = sg.git('diff', plan.baseline);
+        return {
+          throughTurnId,
+          baseTurnId: plan.baseTurnId,
+          affectedTurns: plan.affectedTurns,
+          diff: result.stdout,
+        };
       },
 
       // ---- B5: rollbackCodeToTurn ----
@@ -553,33 +630,19 @@ export class CheckpointService extends Effect.Service<CheckpointService>()('Chec
         throughTurnId: number
       ): CodeRollbackResult => {
         const sg = ensure(projectPath);
-        const completedTurns = getCompletedTurnsFor(sg, sessionId);
-        const affectedTurns = completedTurns.filter((id) => id > throughTurnId);
-        if (affectedTurns.length === 0) {
-          return {
-            reverted: true,
-            throughTurnId,
-            baseTurnId: null,
-            affectedTurns: [],
-            selectedFiles: [],
-            restoreEntry: null,
-          };
-        }
-
-        const baseTurnId = affectedTurns[0]! - 1;
-        const baseline = sg.findCommitByMessage(commitMsg(sessionId, baseTurnId, 'baseline'));
-        if (!baseline) {
+        const plan = getRollbackToTurnPlan(sg, sessionId, throughTurnId);
+        if (!plan) {
           return {
             reverted: false,
             throughTurnId,
-            baseTurnId,
+            baseTurnId: null,
             affectedTurns: [],
             selectedFiles: [],
             restoreEntry: null,
           };
         }
 
-        const diffResult = sg.git('diff', '--name-only', baseline);
+        const diffResult = sg.git('diff', '--name-only', plan.baseline);
         const selectedFiles = diffResult.stdout
           .trim()
           .split('\n')
@@ -590,45 +653,14 @@ export class CheckpointService extends Effect.Service<CheckpointService>()('Chec
           return {
             reverted: true,
             throughTurnId,
-            baseTurnId,
-            affectedTurns,
+            baseTurnId: plan.baseTurnId,
+            affectedTurns: plan.affectedTurns,
             selectedFiles: [],
             restoreEntry: null,
           };
         }
 
-        sg.lock();
-        try {
-          const safetyCommit = sg.commit(commitMsg(sessionId, throughTurnId, 'rollback-safety'));
-          sg.checkoutFiles(baseline, selectedFiles);
-
-          const entry: CodeRestoreEntry = {
-            id: createHash('sha256')
-              .update(`${sessionId}-${Date.now()}`)
-              .digest('hex')
-              .slice(0, 12),
-            sessionId,
-            action: 'rollback-to-turn',
-            throughTurnId,
-            baseTurnId,
-            affectedTurns,
-            selectedFiles,
-            safetyCommit,
-            timestamp: new Date().toISOString(),
-          };
-          writeRestoreEntry(sg.gitDir, sessionId, entry);
-
-          return {
-            reverted: true,
-            throughTurnId,
-            baseTurnId,
-            affectedTurns,
-            selectedFiles,
-            restoreEntry: entry,
-          };
-        } finally {
-          sg.unlock();
-        }
+        return revertFilesImpl(projectPath, sessionId, plan, selectedFiles, 'rollback-to-turn', sg);
       },
 
       // ---- B6: undoLastCodeRollback ----

packages/codingcode/src/server/routes/sessions.ts

@@ -17,6 +17,19 @@ import { runWithLayer, errorResponse } from '../util.js';
 
 export const sessionsRouter = new Hono();
 
+function findUserMessageForTurn(sessionId: string, turnId: number): string {
+  const dir = resolveSessionDir(sessionId);
+  if (!dir) return '';
+  const jsonlPath = join(dir, `${sessionId}.jsonl`);
+  const rawEvents = readHistory(jsonlPath);
+  for (const ev of rawEvents) {
+    if (ev.type === 'user' && (ev as any).turnId === turnId) {
+      return (ev as any).content ?? '';
+    }
+  }
+  return '';
+}
+
 // Active session ApprovalService forks, keyed by sessionId.
 // messages.ts registers/unregisters; this file's PUT route updates them.
 export const activeApprovalForks = new Map<
@@ -365,21 +378,9 @@ sessionsRouter.post('/:id/rollback-context', async (c) => {
     Effect.gen(function* () {
       const svc = yield* SessionService;
       const state = yield* svc.create(cwd, 'unknown', sessionId);
+      const rolledBackMessage = findUserMessageForTurn(sessionId, body.throughTurnId);
       yield* svc.rollbackToTurn(state, body.throughTurnId, 'user rollback');
       const turns = readUIHistory(sessionId);
-      // Find user message of the rolled-back turn for input refill
-      let rolledBackMessage = '';
-      const dir = resolveSessionDir(sessionId);
-      if (dir) {
-        const jsonlPath = join(dir, `${sessionId}.jsonl`);
-        const rawEvents = readHistory(jsonlPath);
-        for (const ev of rawEvents) {
-          if (ev.type === 'user' && (ev as any).turnId === body.throughTurnId) {
-            rolledBackMessage = (ev as any).content ?? '';
-            break;
-          }
-        }
-      }
       return { ok: true, turns, rolledBackMessage, promptEstimate: state.promptEstimate };
     })
   );
@@ -402,9 +403,10 @@ sessionsRouter.post('/:id/rollback-both-to-turn', async (c) => {
       const svc = yield* SessionService;
       const codeResult = checkpoint.rollbackCodeToTurn(cwd, sessionId, body.throughTurnId);
       const state = yield* svc.create(cwd, 'unknown', sessionId);
+      const rolledBackMessage = findUserMessageForTurn(sessionId, body.throughTurnId);
       yield* svc.rollbackToTurn(state, body.throughTurnId, 'user rollback');
       const turns = readUIHistory(sessionId);
-      return { ok: true, turns, codeResult, promptEstimate: state.promptEstimate };
+      return { ok: true, turns, codeResult, rolledBackMessage, promptEstimate: state.promptEstimate };
     })
   );
   if (!result.ok) {