diff --git a/packages/ui/src/components/multirun/BranchSelector.tsx b/packages/ui/src/components/multirun/BranchSelector.tsx index 8b86ac66..946e5432 100644 --- a/packages/ui/src/components/multirun/BranchSelector.tsx +++ b/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 = ({ // 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 = ({ }; void resolve(); + return () => { + cancelled = true; + }; }, [allBranches, directory, disabled, isLoading, onChange, value]); const isDisabled = disabled || !isGitRepository || isLoading; diff --git a/packages/ui/src/components/session/NewWorktreeDialog.tsx b/packages/ui/src/components/session/NewWorktreeDialog.tsx index cda33b7b..68291c76 100644 --- a/packages/ui/src/components/session/NewWorktreeDialog.tsx +++ b/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'), { diff --git a/packages/ui/src/lib/worktrees/worktreeSourceBranchPreference.test.ts b/packages/ui/src/lib/worktrees/worktreeSourceBranchPreference.test.ts new file mode 100644 index 00000000..bc41f07c --- /dev/null +++ b/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'); + }); +}); diff --git a/packages/ui/src/lib/worktrees/worktreeSourceBranchPreference.ts b/packages/ui/src/lib/worktrees/worktreeSourceBranchPreference.ts new file mode 100644 index 00000000..a8435221 --- /dev/null +++ b/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), + }; +};