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'), {

packages/ui/src/lib/worktrees/worktreeSourceBranchPreference.test.ts

@@ -0,0 +1,70 @@
+import { describe, expect, test } from 'bun:test';
+
+import {
+  resolveWorktreeSourceBranchPreference,
+  resolveWorktreeSourceBranchToPersist,
+} from './worktreeSourceBranchPreference';
+
+describe('resolveWorktreeSourceBranchPreference', () => {
+  test('keeps a valid saved source branch', () => {
+    expect(resolveWorktreeSourceBranchPreference({
+      branches: ['develop', 'main', 'remotes/origin/main'],
+      savedSourceBranch: 'develop',
+      rootBranch: 'main',
+    })).toEqual({
+      sourceBranch: 'develop',
+      shouldClearSavedSourceBranch: false,
+    });
+  });
+
+  test('falls back to the root branch and marks a stale saved branch for cleanup', () => {
+    expect(resolveWorktreeSourceBranchPreference({
+      branches: ['main', 'develop', 'remotes/origin/main'],
+      savedSourceBranch: 'feature/stale',
+      rootBranch: 'main',
+    })).toEqual({
+      sourceBranch: 'main',
+      shouldClearSavedSourceBranch: true,
+    });
+  });
+
+  test('uses a deterministic fallback when the root branch is unavailable', () => {
+    expect(resolveWorktreeSourceBranchPreference({
+      branches: ['feature/new', 'master'],
+      savedSourceBranch: 'feature/stale',
+      rootBranch: null,
+    })).toEqual({
+      sourceBranch: 'master',
+      shouldClearSavedSourceBranch: true,
+    });
+  });
+});
+
+describe('resolveWorktreeSourceBranchToPersist', () => {
+  test('persists the source branch for a new worktree without a linked PR', () => {
+    expect(resolveWorktreeSourceBranchToPersist({
+      mode: 'new-branch',
+      sourceBranch: 'develop',
+      linkedPr: false,
+      selectedBranch: '',
+    })).toBe('develop');
+  });
+
+  test('skips persisting a PR-linked source branch', () => {
+    expect(resolveWorktreeSourceBranchToPersist({
+      mode: 'new-branch',
+      sourceBranch: 'feature/pr',
+      linkedPr: true,
+      selectedBranch: '',
+    })).toBeNull();
+  });
+
+  test('persists the selected branch for existing-branch mode', () => {
+    expect(resolveWorktreeSourceBranchToPersist({
+      mode: 'existing-branch',
+      sourceBranch: 'feature/unused',
+      linkedPr: false,
+      selectedBranch: 'hotfix/fix-123',
+    })).toBe('hotfix/fix-123');
+  });
+});

packages/ui/src/lib/worktrees/worktreeSourceBranchPreference.ts

@@ -0,0 +1,64 @@
+export const LAST_WORKTREE_SOURCE_BRANCH_KEY = 'oc:lastWorktreeSourceBranch';
+
+export interface WorktreeSourceBranchPreferenceArgs {
+  branches: readonly string[];
+  savedSourceBranch: string | null;
+  rootBranch: string | null;
+}
+
+export interface WorktreeSourceBranchPreferenceResult {
+  sourceBranch: string;
+  shouldClearSavedSourceBranch: boolean;
+}
+
+export interface WorktreeSourceBranchPersistArgs {
+  mode: 'new-branch' | 'existing-branch';
+  sourceBranch: string;
+  linkedPr: boolean;
+  selectedBranch: string;
+}
+
+export const resolveWorktreeSourceBranchToPersist = ({
+  mode,
+  sourceBranch,
+  linkedPr,
+  selectedBranch,
+}: WorktreeSourceBranchPersistArgs): string | null => {
+  if (mode === 'existing-branch') {
+    return selectedBranch || null;
+  }
+
+  if (linkedPr) {
+    return null;
+  }
+
+  return sourceBranch || null;
+};
+
+export const resolveWorktreeSourceBranchPreference = ({
+  branches,
+  savedSourceBranch,
+  rootBranch,
+}: WorktreeSourceBranchPreferenceArgs): WorktreeSourceBranchPreferenceResult => {
+  const savedSourceBranchIsValid = Boolean(savedSourceBranch && branches.includes(savedSourceBranch));
+
+  if (savedSourceBranchIsValid && savedSourceBranch) {
+    return {
+      sourceBranch: savedSourceBranch,
+      shouldClearSavedSourceBranch: false,
+    };
+  }
+
+  const sourceBranch = rootBranch && branches.includes(rootBranch)
+    ? rootBranch
+    : branches.includes('main')
+      ? 'main'
+      : branches.includes('master')
+        ? 'master'
+        : branches[0] ?? '';
+
+  return {
+    sourceBranch,
+    shouldClearSavedSourceBranch: Boolean(savedSourceBranch),
+  };
+};