fix(worktree): restore last source branch reliably (#2030)

* fix(worktree): restore last source branch reliably

* chore: retrigger review

---------

Co-authored-by: bashrusakh <bashrusakh@users.noreply.github.com>

packages/ui/src/components/multirun/BranchSelector.tsx

@@ -12,11 +12,12 @@ import {
 import { useGitStore, useGitBranches, useGitLoadingBranches, useGitLoadingStatus, useIsGitRepo } from '@/stores/useGitStore';
 import { useRuntimeAPIs } from '@/hooks/useRuntimeAPIs';
 import { getRootBranch } from '@/lib/worktrees/worktreeStatus';
+import {
+  LAST_WORKTREE_SOURCE_BRANCH_KEY,
+  resolveWorktreeSourceBranchPreference,
+} from '@/lib/worktrees/worktreeSourceBranchPreference';
 import { useI18n } from '@/lib/i18n';
 
-/** localStorage key matching NewWorktreeDialog */
-const LAST_SOURCE_BRANCH_KEY = 'oc:lastWorktreeSourceBranch';
-
 export interface BranchSelectorProps {
   /** Current directory to check for git repository */
   directory: string | null;
@@ -117,24 +118,35 @@ export const BranchSelector: React.FC<BranchSelectorProps> = ({
   // Resolve default source branch (same priority as NewWorktreeDialog)
   React.useEffect(() => {
     if (disabled || isLoading || allBranches.length === 0) return;
-    // If current value is valid, keep it
     if (value && allBranches.includes(value)) return;
 
+    const currentValue = value;
+    let cancelled = false;
+
     const resolve = async () => {
       try {
         const rootBranch = directory ? await getRootBranch(directory).catch(() => null) : null;
-        const saved = localStorage.getItem(LAST_SOURCE_BRANCH_KEY);
+        if (cancelled) return;
 
-        if (saved && allBranches.includes(saved)) {
-          onChange(saved);
-        } else if (rootBranch && allBranches.includes(rootBranch)) {
-          onChange(rootBranch);
-        } else if (allBranches.includes('main')) {
-          onChange('main');
-        } else if (allBranches.includes('master')) {
-          onChange('master');
-        } else if (allBranches[0]) {
-          onChange(allBranches[0]);
+        const saved = localStorage.getItem(LAST_WORKTREE_SOURCE_BRANCH_KEY);
+
+        const {
+          sourceBranch,
+          shouldClearSavedSourceBranch,
+        } = resolveWorktreeSourceBranchPreference({
+          branches: allBranches,
+          savedSourceBranch: saved,
+          rootBranch,
+        });
+
+        if (shouldClearSavedSourceBranch) {
+          localStorage.removeItem(LAST_WORKTREE_SOURCE_BRANCH_KEY);
+        }
+
+        if (cancelled || (currentValue && allBranches.includes(currentValue))) return;
+
+        if (sourceBranch) {
+          onChange(sourceBranch);
         }
       } catch {
         // ignore
@@ -142,6 +154,9 @@ export const BranchSelector: React.FC<BranchSelectorProps> = ({
     };
 
     void resolve();
+    return () => {
+      cancelled = true;
+    };
   }, [allBranches, directory, disabled, isLoading, onChange, value]);
 
   const isDisabled = disabled || !isGitRepository || isLoading;

packages/ui/src/components/session/NewWorktreeDialog.tsx

@@ -40,6 +40,11 @@ import { generateBranchSlug } from '@/lib/git/branchNameGenerator';
 import { renderMagicPrompt } from '@/lib/magicPrompts';
 import { parseModelIdentifier } from '@/lib/modelIdentifier';
 import { rankBranchesForQuery } from '@/lib/worktrees/branchSearch';
+import {
+  LAST_WORKTREE_SOURCE_BRANCH_KEY,
+  resolveWorktreeSourceBranchPreference,
+  resolveWorktreeSourceBranchToPersist,
+} from '@/lib/worktrees/worktreeSourceBranchPreference';
 import { useRuntimeAPIs } from '@/hooks/useRuntimeAPIs';
 import { useGitBranches, useGitStore, useGitLoadingBranches } from '@/stores/useGitStore';
 import { GitHubIntegrationDialog } from './GitHubIntegrationDialog';
@@ -105,8 +110,6 @@ const slugifyWorktreeName = (value: string): string => {
     .slice(0, 80);
 };
 
-const LAST_SOURCE_BRANCH_KEY = 'oc:lastWorktreeSourceBranch';
-
 const sanitizeRemoteName = (value: string): string => {
   const normalized = String(value || '')
     .trim()
@@ -586,25 +589,35 @@ export function NewWorktreeDialog({
   // Get current state based on mode
   const currentState = mode === 'new-branch' ? newBranchState : existingBranchState;
 
-  // Set default source branch when branches become available
+  // Set default source branch when the dialog opens and branches become available
   React.useEffect(() => {
-    if (!branches?.all || !projectDirectory) return;
-    if (newBranchState.sourceBranch) return; // Already set
-    
+    if (!open || !branches?.all || !projectDirectory) return;
+    if (newBranchState.sourceBranch) return;
+
+    const currentSourceBranch = newBranchState.sourceBranch;
+    let cancelled = false;
+
     const loadDefaultSourceBranch = async () => {
       try {
         const rootBranch = await getRootBranch(projectDirectory).catch(() => null);
-        const savedSourceBranch = localStorage.getItem(LAST_SOURCE_BRANCH_KEY);
-        const defaultSourceBranch = savedSourceBranch && branches.all?.includes(savedSourceBranch)
-          ? savedSourceBranch
-          : rootBranch && branches.all?.includes(rootBranch)
-            ? rootBranch
-            : branches.all?.includes('main')
-              ? 'main'
-              : branches.all?.includes('master')
-                ? 'master'
-                : branches.all?.[0] || '';
-        
+        if (cancelled) return;
+
+        const savedSourceBranch = localStorage.getItem(LAST_WORKTREE_SOURCE_BRANCH_KEY);
+        const {
+          sourceBranch: defaultSourceBranch,
+          shouldClearSavedSourceBranch,
+        } = resolveWorktreeSourceBranchPreference({
+          branches: branches.all,
+          savedSourceBranch,
+          rootBranch,
+        });
+
+        if (shouldClearSavedSourceBranch) {
+          localStorage.removeItem(LAST_WORKTREE_SOURCE_BRANCH_KEY);
+        }
+
+        if (cancelled || currentSourceBranch) return;
+
         if (defaultSourceBranch) {
           setNewBranchState(prev => ({
             ...prev,
@@ -615,13 +628,16 @@ export function NewWorktreeDialog({
         // ignore
       }
     };
-    
+
     void loadDefaultSourceBranch();
-  }, [branches, projectDirectory, newBranchState.sourceBranch]);
+    return () => {
+      cancelled = true;
+    };
+  }, [open, branches?.all, projectDirectory, newBranchState.sourceBranch]);
 
   // Reset state on each open. Resetting on close would empty the form during
   // the close animation, causing visible flicker.
-  React.useEffect(() => {
+  React.useLayoutEffect(() => {
     if (!open) return;
 
     setMode('new-branch');
@@ -890,9 +906,16 @@ export function NewWorktreeDialog({
         setIsCreating(false);
       }
       
-      // Save source branch preference (only if not from PR)
-      if (newBranchState.sourceBranch && mode === 'new-branch' && !newBranchState.linkedPr) {
-        localStorage.setItem(LAST_SOURCE_BRANCH_KEY, newBranchState.sourceBranch);
+      // Save the last source-branch choice for the next open.
+      const lastSourceBranch = resolveWorktreeSourceBranchToPersist({
+        mode,
+        sourceBranch: newBranchState.sourceBranch,
+        linkedPr: !!newBranchState.linkedPr,
+        selectedBranch: existingBranchState.selectedBranch,
+      });
+
+      if (lastSourceBranch) {
+        localStorage.setItem(LAST_WORKTREE_SOURCE_BRANCH_KEY, lastSourceBranch);
       }
       
       toast.success(t('session.newWorktree.toast.worktreeCreated'), {