fix(tui): cache plan path before ExitPlanMode arrives (#1337)

Co-authored-by: Letta Code <noreply@letta.com>

src/cli/App.tsx

@@ -1067,7 +1067,18 @@ export default function App({
     permissionMode.getMode(),
   );
   const uiPermissionModeRef = useRef<PermissionMode>(uiPermissionMode);
-  const setUiPermissionMode = useCallback((mode: PermissionMode) => {
+
+  // Store the last plan file path for post-approval rendering
+  // (needed because plan mode is exited before rendering the result)
+  const lastPlanFilePathRef = useRef<string | null>(null);
+  const cacheLastPlanFilePath = useCallback((planFilePath: string | null) => {
+    if (planFilePath) {
+      lastPlanFilePathRef.current = planFilePath;
+    }
+  }, []);
+
+  const setUiPermissionMode = useCallback(
+    (mode: PermissionMode) => {
       uiPermissionModeRef.current = mode;
       _setUiPermissionMode(mode);
 
@@ -1082,10 +1093,13 @@ export default function App({
         if (mode === "plan" && !permissionMode.getPlanFilePath()) {
           const planPath = generatePlanFilePath();
           permissionMode.setPlanFilePath(planPath);
+          cacheLastPlanFilePath(planPath);
         }
         permissionMode.setMode(mode);
       }
-  }, []);
+    },
+    [cacheLastPlanFilePath],
+  );
 
   const statusLineTriggerVersionRef = useRef(0);
   const [statusLineTriggerVersion, setStatusLineTriggerVersion] = useState(0);
@@ -2531,10 +2545,6 @@ export default function App({
     new Map(),
   );
 
-  // Store the last plan file path for post-approval rendering
-  // (needed because plan mode is exited before rendering the result)
-  const lastPlanFilePathRef = useRef<string | null>(null);
-
   // Track which approval tool call IDs have had their previews eagerly committed
   // This prevents double-committing when the approval changes
   const eagerCommittedPreviewsRef = useRef<Set<string>>(new Set());
@@ -9234,6 +9244,7 @@ export default function App({
           // Generate plan file path and enter plan mode
           const planPath = generatePlanFilePath();
           permissionMode.setPlanFilePath(planPath);
+          cacheLastPlanFilePath(planPath);
           permissionMode.setMode("plan");
           setUiPermissionMode("plan");
 
@@ -12009,12 +12020,13 @@ ${SYSTEM_REMINDER_CLOSE}
       if (mode === "plan") {
         const planPath = generatePlanFilePath();
         permissionMode.setPlanFilePath(planPath);
+        cacheLastPlanFilePath(planPath);
       }
       // permissionMode.setMode() is called in InputRich.tsx before this callback
       setUiPermissionMode(mode);
       triggerStatusLineRefresh();
     },
-    [triggerStatusLineRefresh, setUiPermissionMode],
+    [triggerStatusLineRefresh, setUiPermissionMode, cacheLastPlanFilePath],
   );
 
   // Reasoning tier cycling (Tab hotkey in InputRich.tsx)
@@ -12327,12 +12339,18 @@ ${SYSTEM_REMINDER_CLOSE}
         lastPlanFilePathRef.current = planFilePath;
       }
 
-      // Exit plan mode
+      // Exit plan mode — if user already cycled out (e.g., Shift+Tab to
+      // acceptEdits/yolo), keep their chosen mode instead of downgrading.
+      const currentMode = permissionMode.getMode();
+      if (currentMode === "plan") {
         const restoreMode = acceptEdits
           ? "acceptEdits"
           : (permissionMode.getModeBeforePlan() ?? "default");
         permissionMode.setMode(restoreMode);
         setUiPermissionMode(restoreMode);
+      } else {
+        setUiPermissionMode(currentMode);
+      }
 
       try {
         // Execute ExitPlanMode tool to get the result
@@ -12434,8 +12452,13 @@ ${SYSTEM_REMINDER_CLOSE}
 
       if (mode !== "plan") {
         if (hasUsablePlan) {
-          // User likely cycled out of plan mode (e.g., Shift+Tab to acceptEdits/yolo)
-          // Keep approval flow alive and let ExitPlanMode proceed using fallback plan path.
+          if (mode === "bypassPermissions") {
+            // User cycled to YOLO mode — auto-approve ExitPlanMode
+            // so they don't need to manually click through the approval.
+            handlePlanApprove();
+            return;
+          }
+          // Other modes: keep approval flow alive and let user manually approve.
           return;
         }
 
@@ -12491,6 +12514,7 @@ ${SYSTEM_REMINDER_CLOSE}
   }, [
     pendingApprovals,
     approvalResults.length,
+    handlePlanApprove,
     handlePlanKeepPlanning,
     refreshDerived,
     queueApprovalResults,
@@ -12573,6 +12597,7 @@ ${SYSTEM_REMINDER_CLOSE}
     // Toggle plan mode on and store plan file path
     permissionMode.setMode("plan");
     permissionMode.setPlanFilePath(planFilePath);
+    cacheLastPlanFilePath(planFilePath);
     setUiPermissionMode("plan");
 
     // Get the tool return message from the implementation
@@ -12629,6 +12654,7 @@ If using apply_patch, use this exact relative patch path: ${applyPatchRelativePa
     sendAllResults,
     refreshDerived,
     setUiPermissionMode,
+    cacheLastPlanFilePath,
   ]);
 
   const handleEnterPlanModeReject = useCallback(async () => {

src/tests/cli/permission-mode-retry-wiring.test.ts

@@ -11,11 +11,10 @@ describe("permission mode retry wiring", () => {
   test("setUiPermissionMode syncs singleton mode immediately", () => {
     const source = readAppSource();
 
-    const start = source.indexOf(
-      "const setUiPermissionMode = useCallback((mode: PermissionMode) => {",
-    );
+    const start = source.indexOf("const setUiPermissionMode = useCallback(");
     const end = source.indexOf(
       "const statusLineTriggerVersionRef = useRef(0);",
+      start,
     );
     expect(start).toBeGreaterThan(-1);
     expect(end).toBeGreaterThan(start);
@@ -26,9 +25,53 @@ describe("permission mode retry wiring", () => {
       'if (mode === "plan" && !permissionMode.getPlanFilePath())',
     );
     expect(segment).toContain("permissionMode.setPlanFilePath(planPath);");
+    expect(segment).toContain("cacheLastPlanFilePath(planPath);");
     expect(segment).toContain("permissionMode.setMode(mode);");
   });
 
+  test("caches the plan path at every plan-mode entry point", () => {
+    const source = readAppSource();
+
+    expect(source).toContain(
+      "const cacheLastPlanFilePath = useCallback((planFilePath: string | null) => {",
+    );
+
+    const slashPlanStart = source.indexOf('if (trimmed === "/plan") {');
+    const slashPlanEnd = source.indexOf(
+      "return { submitted: true };",
+      slashPlanStart,
+    );
+    expect(slashPlanStart).toBeGreaterThan(-1);
+    expect(slashPlanEnd).toBeGreaterThan(slashPlanStart);
+    expect(source.slice(slashPlanStart, slashPlanEnd)).toContain(
+      "cacheLastPlanFilePath(planPath);",
+    );
+
+    const modeChangeStart = source.indexOf(
+      "const handlePermissionModeChange = useCallback(",
+    );
+    const modeChangeEnd = source.indexOf(
+      "// Reasoning tier cycling (Tab hotkey in InputRich.tsx)",
+    );
+    expect(modeChangeStart).toBeGreaterThan(-1);
+    expect(modeChangeEnd).toBeGreaterThan(modeChangeStart);
+    expect(source.slice(modeChangeStart, modeChangeEnd)).toContain(
+      "cacheLastPlanFilePath(planPath);",
+    );
+
+    const enterPlanStart = source.indexOf(
+      "const handleEnterPlanModeApprove = useCallback(async () => {",
+    );
+    const enterPlanEnd = source.indexOf(
+      "const handleEnterPlanModeReject = useCallback(async () => {",
+    );
+    expect(enterPlanStart).toBeGreaterThan(-1);
+    expect(enterPlanEnd).toBeGreaterThan(enterPlanStart);
+    expect(source.slice(enterPlanStart, enterPlanEnd)).toContain(
+      "cacheLastPlanFilePath(planFilePath);",
+    );
+  });
+
   test("pins submission permission mode and defines a restore helper", () => {
     const source = readAppSource();