Fix source control branch compare base (#5074)

Co-authored-by: Orca <help@stably.ai>

src/renderer/src/components/right-sidebar/SourceControl.compare-summary.test.ts

@@ -2,6 +2,7 @@ import { describe, expect, it, vi } from 'vitest'
 import {
   CompareSummary,
   CompareSummaryToolbarButton,
+  resolveSourceControlBaseRef,
   shouldShowCompareSummary
 } from './SourceControl'
 import type { GitBranchCompareSummary } from '../../../../shared/types'
@@ -75,6 +76,33 @@ const readySummary: GitBranchCompareSummary = {
 }
 
 describe('SourceControl compare summary', () => {
+  it('prefers the worktree creation base for branch compare', () => {
+    expect(
+      resolveSourceControlBaseRef({
+        worktreeBaseRef: 'refs/remotes/origin/main',
+        repoBaseRef: 'main',
+        defaultBaseRef: 'origin/main'
+      })
+    ).toBe('refs/remotes/origin/main')
+  })
+
+  it('falls back to repo and default base refs when worktree metadata is absent', () => {
+    expect(
+      resolveSourceControlBaseRef({
+        worktreeBaseRef: '  ',
+        repoBaseRef: ' origin/release ',
+        defaultBaseRef: 'origin/main'
+      })
+    ).toBe('origin/release')
+
+    expect(
+      resolveSourceControlBaseRef({
+        repoBaseRef: null,
+        defaultBaseRef: 'origin/main'
+      })
+    ).toBe('origin/main')
+  })
+
   it('wires toolbar actions without rendering the dead view-mode toggle', () => {
     const onChangeBaseRef = vi.fn()
     const onRetry = vi.fn()

src/renderer/src/components/right-sidebar/SourceControl.tsx

@@ -220,6 +220,19 @@ export type SourceControlActionError = {
   message: string
 }
 
+export function resolveSourceControlBaseRef(input: {
+  worktreeBaseRef?: string | null
+  repoBaseRef?: string | null
+  defaultBaseRef?: string | null
+}): string | null {
+  return (
+    input.worktreeBaseRef?.trim() ||
+    input.repoBaseRef?.trim() ||
+    input.defaultBaseRef?.trim() ||
+    null
+  )
+}
+
 const EMPTY_GIT_STATUS_ENTRIES: GitStatusEntry[] = []
 const EMPTY_BRANCH_CHANGE_ENTRIES: GitBranchChangeEntry[] = []
 
@@ -1153,7 +1166,15 @@ function SourceControlInner(): React.JSX.Element {
     }
   }, [activeRepo, isBranchVisible, isFolder])
 
-  const effectiveBaseRef = activeRepo?.worktreeBaseRef ?? defaultBaseRef
+  const normalizedWorktreeBaseRef = activeWorktree?.baseRef?.trim() || null
+  const normalizedRepoBaseRef = activeRepo?.worktreeBaseRef?.trim() || null
+  const effectiveBaseRef = resolveSourceControlBaseRef({
+    worktreeBaseRef: normalizedWorktreeBaseRef,
+    repoBaseRef: normalizedRepoBaseRef,
+    defaultBaseRef
+  })
+  const baseRefOwnedByWorktree = normalizedWorktreeBaseRef !== null
+  const pinnedBaseRef = normalizedWorktreeBaseRef ?? normalizedRepoBaseRef
   const hasUncommittedEntries = entries.length > 0
 
   const hostedReviewCreation =
@@ -4557,14 +4578,22 @@ function SourceControlInner(): React.JSX.Element {
           </DialogHeader>
           <BaseRefPicker
             repoId={activeRepo.id}
-            currentBaseRef={activeRepo.worktreeBaseRef}
+            currentBaseRef={pinnedBaseRef ?? undefined}
             onSelect={(ref) => {
-              void updateRepo(activeRepo.id, { worktreeBaseRef: ref })
+              if (baseRefOwnedByWorktree && activeWorktreeId) {
+                void updateWorktreeMeta(activeWorktreeId, { baseRef: ref })
+              } else {
+                void updateRepo(activeRepo.id, { worktreeBaseRef: ref })
+              }
               setBaseRefDialogOpen(false)
               window.setTimeout(() => void refreshBranchCompare(), 0)
             }}
             onUsePrimary={() => {
-              void updateRepo(activeRepo.id, { worktreeBaseRef: undefined })
+              if (baseRefOwnedByWorktree && activeWorktreeId) {
+                void updateWorktreeMeta(activeWorktreeId, { baseRef: undefined })
+              } else {
+                void updateRepo(activeRepo.id, { worktreeBaseRef: undefined })
+              }
               setBaseRefDialogOpen(false)
               window.setTimeout(() => void refreshBranchCompare(), 0)
             }}

src/renderer/src/store/slices/worktrees.test.ts

@@ -1259,6 +1259,31 @@ describe('createWorktree base status merge', () => {
     })
   })
 
+  it('merges create result metadata into a worktree inserted by the watcher race', async () => {
+    const store = createTestStore()
+    const watcherWorktree = makeWorktree({
+      id: 'repo1::/path/wt1',
+      repoId: 'repo1',
+      path: '/path/wt1'
+    })
+    const createdWorktree = makeWorktree({
+      ...watcherWorktree,
+      baseRef: 'refs/remotes/origin/main'
+    })
+    store.setState({
+      worktreesByRepo: { repo1: [watcherWorktree] }
+    } as Partial<AppState>)
+    mockApi.worktrees.create.mockResolvedValue({ worktree: createdWorktree })
+
+    await store.getState().createWorktree('repo1', 'feature', 'origin/main')
+
+    expect(store.getState().worktreesByRepo.repo1).toHaveLength(1)
+    expect(store.getState().worktreesByRepo.repo1[0]).toMatchObject({
+      id: watcherWorktree.id,
+      baseRef: 'refs/remotes/origin/main'
+    })
+  })
+
   it.each([
     {
       status: 'skipped_dirty_worktree',

src/renderer/src/store/slices/worktrees.ts

@@ -1149,10 +1149,17 @@ export const createWorktreeSlice: StateCreator<AppState, [], [], WorktreeSlice>
           set((s) => {
             const current = s.worktreesByRepo[repoId] ?? []
             const alreadyPresent = current.some((w) => w.id === result.worktree.id)
+            const nextWorktrees = alreadyPresent
+              ? current.map((worktree) =>
+                  worktree.id === result.worktree.id
+                    ? { ...worktree, ...result.worktree }
+                    : worktree
+                )
+              : [...current, result.worktree]
             return {
               worktreesByRepo: {
                 ...s.worktreesByRepo,
-                [repoId]: alreadyPresent ? current : [...current, result.worktree]
+                [repoId]: nextWorktrees
               },
               ...(result.initialBaseStatus
                 ? {