fix(ui): persist staged task state across app restarts

apps/frontend/src/main/ipc-handlers/task/worktree-handlers.ts

@@ -1,5 +1,5 @@
 import { ipcMain, BrowserWindow, shell, app } from 'electron';
-import { IPC_CHANNELS, AUTO_BUILD_PATHS, DEFAULT_APP_SETTINGS, DEFAULT_FEATURE_MODELS, DEFAULT_FEATURE_THINKING, MODEL_ID_MAP, THINKING_BUDGET_MAP } from '../../../shared/constants';
+import { IPC_CHANNELS, AUTO_BUILD_PATHS, DEFAULT_APP_SETTINGS, DEFAULT_FEATURE_MODELS, DEFAULT_FEATURE_THINKING, MODEL_ID_MAP, THINKING_BUDGET_MAP, getSpecsDir } from '../../../shared/constants';
 import type { IPCResult, WorktreeStatus, WorktreeDiff, WorktreeDiffFile, WorktreeMergeResult, WorktreeDiscardResult, WorktreeListResult, WorktreeListItem, SupportedIDE, SupportedTerminal, AppSettings } from '../../../shared/types';
 import path from 'path';
 import { existsSync, readdirSync, statSync, readFileSync } from 'fs';
@@ -2536,4 +2536,70 @@
       }
     }
   );
+
+  /**
+   * Clear the staged state for a task
+   * This allows the user to re-stage changes if needed
+   */
+  ipcMain.handle(
+    IPC_CHANNELS.TASK_CLEAR_STAGED_STATE,
+    async (_, taskId: string): Promise<IPCResult<{ cleared: boolean }>> => {
+      try {
+        const { task, project } = findTaskAndProject(taskId);
+        if (!task || !project) {
+          return { success: false, error: 'Task not found' };
+        }
+
+        const specsBaseDir = getSpecsDir(project.autoBuildPath);
+        const specDir = path.join(project.path, specsBaseDir, task.specId);
+        const planPath = path.join(specDir, AUTO_BUILD_PATHS.IMPLEMENTATION_PLAN);
+
+        if (!existsSync(planPath)) {
+          return { success: false, error: 'Implementation plan not found' };
+        }
+
+        // Read, update, and write the plan file
+        const { promises: fsPromises } = require('fs');
+        const planContent = await fsPromises.readFile(planPath, 'utf-8');
+        const plan = JSON.parse(planContent);
+
+        // Clear the staged state flags
+        delete plan.stagedInMainProject;
+        delete plan.stagedAt;
+        plan.updated_at = new Date().toISOString();
+
+        await fsPromises.writeFile(planPath, JSON.stringify(plan, null, 2));
+
+        // Also update worktree plan if it exists
+        const worktreePath = findTaskWorktree(project.path, task.specId);
+        if (worktreePath) {
+          const worktreePlanPath = path.join(worktreePath, specsBaseDir, task.specId, AUTO_BUILD_PATHS.IMPLEMENTATION_PLAN);
+          if (existsSync(worktreePlanPath)) {
+            try {
+              const worktreePlanContent = await fsPromises.readFile(worktreePlanPath, 'utf-8');
+              const worktreePlan = JSON.parse(worktreePlanContent);
+              delete worktreePlan.stagedInMainProject;
+              delete worktreePlan.stagedAt;
+              worktreePlan.updated_at = new Date().toISOString();
+              await fsPromises.writeFile(worktreePlanPath, JSON.stringify(worktreePlan, null, 2));
+            } catch (e) {
+              // Non-fatal - worktree plan update is best-effort
+              console.warn('[CLEAR_STAGED_STATE] Failed to update worktree plan:', e);
+            }
+          }
+        }
+
+        // Invalidate tasks cache to force reload
+        projectStore.invalidateTasksCache(project.id);
+
+        return { success: true, data: { cleared: true } };
+      } catch (error) {
+        console.error('Failed to clear staged state:', error);
+        return {
+          success: false,
+          error: error instanceof Error ? error.message : 'Failed to clear staged state'
+        };
+      }
+    }
+  );
 }

apps/frontend/src/preload/api/task-api.ts

@@ -51,6 +51,7 @@ export interface TaskAPI {
   mergeWorktree: (taskId: string, options?: { noCommit?: boolean }) => Promise<IPCResult<import('../../shared/types').WorktreeMergeResult>>;
   mergeWorktreePreview: (taskId: string) => Promise<IPCResult<import('../../shared/types').WorktreeMergeResult>>;
   discardWorktree: (taskId: string) => Promise<IPCResult<import('../../shared/types').WorktreeDiscardResult>>;
+  clearStagedState: (taskId: string) => Promise<IPCResult<{ cleared: boolean }>>;
   listWorktrees: (projectId: string) => Promise<IPCResult<import('../../shared/types').WorktreeListResult>>;
   worktreeOpenInIDE: (worktreePath: string, ide: SupportedIDE, customPath?: string) => Promise<IPCResult<{ opened: boolean }>>;
   worktreeOpenInTerminal: (worktreePath: string, terminal: SupportedTerminal, customPath?: string) => Promise<IPCResult<{ opened: boolean }>>;
@@ -142,6 +143,9 @@ export const createTaskAPI = (): TaskAPI => ({
   discardWorktree: (taskId: string): Promise<IPCResult<import('../../shared/types').WorktreeDiscardResult>> =>
     ipcRenderer.invoke(IPC_CHANNELS.TASK_WORKTREE_DISCARD, taskId),
 
+  clearStagedState: (taskId: string): Promise<IPCResult<{ cleared: boolean }>> =>
+    ipcRenderer.invoke(IPC_CHANNELS.TASK_CLEAR_STAGED_STATE, taskId),
+
   listWorktrees: (projectId: string): Promise<IPCResult<import('../../shared/types').WorktreeListResult>> =>
     ipcRenderer.invoke(IPC_CHANNELS.TASK_LIST_WORKTREES, projectId),
 

apps/frontend/src/renderer/components/task-detail/TaskDetailModal.tsx

@@ -442,6 +442,7 @@ function TaskDetailModalContent({ open, task, onOpenChange, onSwitchToTerminals,
                             onClose={handleClose}
                             onSwitchToTerminals={onSwitchToTerminals}
                             onOpenInbuiltTerminal={onOpenInbuiltTerminal}
+                            onReviewAgain={state.handleReviewAgain}
                           />
                         </>
                       )}

apps/frontend/src/renderer/components/task-detail/TaskReview.tsx

@@ -42,6 +42,7 @@ interface TaskReviewProps {
   onClose?: () => void;
   onSwitchToTerminals?: () => void;
   onOpenInbuiltTerminal?: (id: string, cwd: string) => void;
+  onReviewAgain?: () => void;
 }
 
 /**
@@ -83,7 +84,8 @@ export function TaskReview({
   onLoadMergePreview,
   onClose,
   onSwitchToTerminals,
-  onOpenInbuiltTerminal
+  onOpenInbuiltTerminal,
+  onReviewAgain
 }: TaskReviewProps) {
   return (
     <div className="space-y-4">
@@ -98,10 +100,24 @@ export function TaskReview({
         />
       )}
 
-      {/* Workspace Status - hide if staging was successful (worktree is deleted after staging) */}
+      {/* Workspace Status - priority: loading > fresh staging success > already staged (persisted) > worktree exists > no workspace */}
       {isLoadingWorktree ? (
         <LoadingMessage />
-      ) : worktreeStatus?.exists && !stagedSuccess ? (
+      ) : stagedSuccess ? (
+        /* Fresh staging just completed - StagedSuccessMessage is rendered above */
+        null
+      ) : task.stagedInMainProject ? (
+        /* Task was previously staged (persisted state) - show even if worktree still exists */
+        <StagedInProjectMessage
+          task={task}
+          projectPath={stagedProjectPath}
+          hasWorktree={worktreeStatus?.exists || false}
+          worktreeStatus={worktreeStatus}
+          onClose={onClose}
+          onReviewAgain={onReviewAgain}
+        />
+      ) : worktreeStatus?.exists ? (
+        /* Worktree exists but not yet staged - show staging UI */
         <WorkspaceStatus
           worktreeStatus={worktreeStatus}
           workspaceError={workspaceError}
@@ -120,13 +136,6 @@ export function TaskReview({
           onSwitchToTerminals={onSwitchToTerminals}
           onOpenInbuiltTerminal={onOpenInbuiltTerminal}
         />
-      ) : task.stagedInMainProject && !stagedSuccess ? (
-        <StagedInProjectMessage
-          task={task}
-          projectPath={stagedProjectPath}
-          hasWorktree={worktreeStatus?.exists || false}
-          onClose={onClose}
-        />
       ) : (
         <NoWorkspaceMessage task={task} onClose={onClose} />
       )}

apps/frontend/src/renderer/components/task-detail/hooks/useTaskDetail.ts

@@ -238,6 +238,37 @@ export function useTaskDetail({ task }: UseTaskDetailOptions) {
     }
   }, [task.id]);
 
+  // Handle "Review Again" - clears staged state and reloads worktree info
+  const handleReviewAgain = useCallback(async () => {
+    // Clear staged success state if it was set in this session
+    setStagedSuccess(null);
+    setStagedProjectPath(undefined);
+    setSuggestedCommitMessage(undefined);
+
+    // Reset merge preview to force re-check
+    setMergePreview(null);
+    hasLoadedPreviewRef.current = null;
+
+    // Reload worktree status
+    setIsLoadingWorktree(true);
+    try {
+      const [statusResult, diffResult] = await Promise.all([
+        window.electronAPI.getWorktreeStatus(task.id),
+        window.electronAPI.getWorktreeDiff(task.id)
+      ]);
+      if (statusResult.success && statusResult.data) {
+        setWorktreeStatus(statusResult.data);
+      }
+      if (diffResult.success && diffResult.data) {
+        setWorktreeDiff(diffResult.data);
+      }
+    } catch (err) {
+      console.error('Failed to reload worktree info:', err);
+    } finally {
+      setIsLoadingWorktree(false);
+    }
+  }, [task.id]);
+
   // NOTE: Merge preview is NO LONGER auto-loaded on modal open.
   // User must click "Check for Conflicts" button to trigger the expensive preview operation.
   // This improves modal open performance significantly (avoids 1-30+ second Python subprocess).
@@ -318,5 +349,6 @@ export function useTaskDetail({ task }: UseTaskDetailOptions) {
     handleLogsScroll,
     togglePhase,
     loadMergePreview,
+    handleReviewAgain,
   };
 }

apps/frontend/src/renderer/components/task-detail/task-review/WorkspaceMessages.tsx

@@ -1,8 +1,8 @@
-import { AlertCircle, GitMerge, Loader2, Trash2, Check } from 'lucide-react';
+import { AlertCircle, GitMerge, Loader2, Trash2, Check, RotateCcw } from 'lucide-react';
 import { useState } from 'react';
 import { Button } from '../../ui/button';
 import { persistTaskStatus } from '../../../stores/task-store';
-import type { Task } from '../../../../shared/types';
+import type { Task, WorktreeStatus } from '../../../../shared/types';
 
 interface LoadingMessageProps {
   message?: string;
@@ -88,14 +88,18 @@
   task: Task;
   projectPath?: string;
   hasWorktree?: boolean;
+  worktreeStatus?: WorktreeStatus | null;
   onClose?: () => void;
+  onReviewAgain?: () => void;
 }
 
 /**
  * Displays message when changes have already been staged in the main project
  */
-export function StagedInProjectMessage({ task, projectPath, hasWorktree = false, onClose }: StagedInProjectMessageProps) {
+export function StagedInProjectMessage({ task, projectPath, hasWorktree = false, worktreeStatus, onClose, onReviewAgain }: StagedInProjectMessageProps) {
   const [isDeleting, setIsDeleting] = useState(false);
+  const [isMarkingDone, setIsMarkingDone] = useState(false);
+  const [isResetting, setIsResetting] = useState(false);
   const [error, setError] = useState<string | null>(null);
 
   const handleDeleteWorktreeAndMarkDone = async () => {
@@ -124,6 +128,46 @@
     }
   };
 
+  const handleMarkDoneOnly = async () => {
+    setIsMarkingDone(true);
+    setError(null);
+
+    try {
+      await persistTaskStatus(task.id, 'done');
+      onClose?.();
+    } catch (err) {
+      console.error('Error marking task as done:', err);
+      setError(err instanceof Error ? err.message : 'Failed to mark as done');
+    } finally {
+      setIsMarkingDone(false);
+    }
+  };
+
+  const handleReviewAgain = async () => {
+    if (!onReviewAgain) return;
+
+    setIsResetting(true);
+    setError(null);
+
+    try {
+      // Clear the staged flag via IPC
+      const result = await window.electronAPI.clearStagedState(task.id);
+
+      if (!result.success) {
+        setError(result.error || 'Failed to reset staged state');
+        return;
+      }
+
+      // Trigger re-render by calling parent callback
+      onReviewAgain();
+    } catch (err) {
+      console.error('Error resetting staged state:', err);
+      setError(err instanceof Error ? err.message : 'Failed to reset staged state');
+    } finally {
+      setIsResetting(false);
+    }
+  };
+
   return (
     <div className="rounded-xl border border-success/30 bg-success/10 p-4">
       <h3 className="font-medium text-sm text-foreground mb-2 flex items-center gap-2">
@@ -143,12 +187,13 @@
       </div>
 
       {/* Action buttons */}
-      {hasWorktree && (
-        <div className="flex flex-col gap-2">
-          <div className="flex gap-2">
+      <div className="flex flex-col gap-2">
+        <div className="flex gap-2">
+          {/* Primary action: Mark Done or Delete Worktree & Mark Done */}
+          {hasWorktree ? (
             <Button
               onClick={handleDeleteWorktreeAndMarkDone}
-              disabled={isDeleting}
+              disabled={isDeleting || isMarkingDone || isResetting}
               size="sm"
               variant="default"
               className="flex-1"
@@ -165,15 +210,88 @@
                 </>
               )}
             </Button>
-          </div>
-          {error && (
-            <p className="text-xs text-destructive">{error}</p>
+          ) : (
+            <Button
+              onClick={handleMarkDoneOnly}
+              disabled={isDeleting || isMarkingDone || isResetting}
+              size="sm"
+              variant="default"
+              className="flex-1"
+            >
+              {isMarkingDone ? (
+                <>
+                  <Loader2 className="h-4 w-4 mr-2 animate-spin" />
+                  Marking done...
+                </>
+              ) : (
+                <>
+                  <Check className="h-4 w-4 mr-2" />
+                  Mark as Done
+                </>
+              )}
+            </Button>
+          )}
+        </div>
+
+        {/* Secondary actions row */}
+        <div className="flex gap-2">
+          {/* Mark Done Only (when worktree exists) - allows keeping worktree */}
+          {hasWorktree && (
+            <Button
+              onClick={handleMarkDoneOnly}
+              disabled={isDeleting || isMarkingDone || isResetting}
+              size="sm"
+              variant="outline"
+              className="flex-1"
+            >
+              {isMarkingDone ? (
+                <>
+                  <Loader2 className="h-4 w-4 mr-2 animate-spin" />
+                  Marking done...
+                </>
+              ) : (
+                <>
+                  <Check className="h-4 w-4 mr-2" />
+                  Mark Done Only
+                </>
+              )}
+            </Button>
+          )}
+
+          {/* Review Again button - only show if worktree exists and callback provided */}
+          {hasWorktree && onReviewAgain && (
+            <Button
+              onClick={handleReviewAgain}
+              disabled={isDeleting || isMarkingDone || isResetting}
+              size="sm"
+              variant="outline"
+              className="flex-1"
+            >
+              {isResetting ? (
+                <>
+                  <Loader2 className="h-4 w-4 mr-2 animate-spin" />
+                  Resetting...
+                </>
+              ) : (
+                <>
+                  <RotateCcw className="h-4 w-4 mr-2" />
+                  Review Again
+                </>
+              )}
+            </Button>
           )}
+        </div>
+
+        {error && (
+          <p className="text-xs text-destructive">{error}</p>
+        )}
+
+        {hasWorktree && (
           <p className="text-xs text-muted-foreground">
-            This will delete the isolated workspace and mark the task as complete.
+            "Delete Worktree & Mark Done" cleans up the isolated workspace. "Mark Done Only" keeps it for reference.
           </p>
-        </div>
-      )}
+        )}
+      </div>
     </div>
   );
 }