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/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),
+  };
+};