feat(task-input): cache cloud default branch for instant "start on trunk" (#3496)

Author: pauldambra

packages/ui/src/features/git-interaction/components/BranchSelector.test.tsx

@@ -75,6 +75,133 @@ describe("BranchSelector cloud mode", () => {
     expect(screen.getByRole("combobox", { name: "Branch" })).toBeEnabled();
   });
 
+  it("seeds the default branch as a pickable item with a loading row while the cloud list loads", async () => {
+    const user = userEvent.setup();
+    renderInTheme(
+      <BranchSelector
+        repoPath="owner/repo"
+        currentBranch={null}
+        defaultBranch="main"
+        workspaceMode="cloud"
+        cloudBranches={[]}
+        cloudBranchesLoading={true}
+        cloudSearchQuery=""
+        selectedBranch="main"
+        onBranchSelect={vi.fn()}
+        onCloudSearchChange={vi.fn()}
+      />,
+    );
+
+    await user.click(screen.getByRole("combobox", { name: "Branch" }));
+
+    expect(await screen.findByRole("option", { name: "main" })).toBeVisible();
+    expect(screen.getByText("Loading branches…")).toBeVisible();
+    // The seeded row makes the list non-empty, so the empty-state stays gated
+    // off (Base UI only reveals it when the content group is data-empty) — it
+    // keeps the `hidden` class rather than flashing "No branches found." above
+    // the trunk row.
+    expect(screen.getByText("No branches found.")).toHaveClass("hidden");
+  });
+
+  it("does not seed the default branch once the user is searching", async () => {
+    const user = userEvent.setup();
+    renderInTheme(
+      <BranchSelector
+        repoPath="owner/repo"
+        currentBranch={null}
+        defaultBranch="main"
+        workspaceMode="cloud"
+        cloudBranches={[]}
+        cloudBranchesLoading={true}
+        cloudSearchQuery="feat"
+        onBranchSelect={vi.fn()}
+        onCloudSearchChange={vi.fn()}
+      />,
+    );
+
+    await user.click(screen.getByRole("combobox", { name: "Branch" }));
+
+    expect(screen.queryByRole("option", { name: "main" })).toBeNull();
+  });
+
+  it("re-selects the default when a stale cached default is replaced by the live one", () => {
+    const onBranchSelect = vi.fn();
+    const { rerender } = renderInTheme(
+      <BranchSelector
+        repoPath="owner/repo"
+        currentBranch={null}
+        defaultBranch="master"
+        workspaceMode="cloud"
+        cloudBranches={[]}
+        cloudBranchesLoading={true}
+        cloudSearchQuery=""
+        selectedBranch={null}
+        onBranchSelect={onBranchSelect}
+        onCloudSearchChange={vi.fn()}
+      />,
+    );
+
+    // Auto-selected the (stale) cached default.
+    expect(onBranchSelect).toHaveBeenLastCalledWith("master");
+
+    // Parent commits that selection; then the live default arrives differing.
+    rerender(
+      <Theme>
+        <BranchSelector
+          repoPath="owner/repo"
+          currentBranch={null}
+          defaultBranch="main"
+          workspaceMode="cloud"
+          cloudBranches={["main"]}
+          cloudBranchesLoading={false}
+          cloudSearchQuery=""
+          selectedBranch="master"
+          onBranchSelect={onBranchSelect}
+          onCloudSearchChange={vi.fn()}
+        />
+      </Theme>,
+    );
+
+    expect(onBranchSelect).toHaveBeenLastCalledWith("main");
+  });
+
+  it("does not override a branch the user picked when the default later changes", () => {
+    const onBranchSelect = vi.fn();
+    const { rerender } = renderInTheme(
+      <BranchSelector
+        repoPath="owner/repo"
+        currentBranch={null}
+        defaultBranch="master"
+        workspaceMode="cloud"
+        cloudBranches={["master", "feature-x"]}
+        cloudBranchesLoading={false}
+        cloudSearchQuery=""
+        selectedBranch="feature-x"
+        onBranchSelect={onBranchSelect}
+        onCloudSearchChange={vi.fn()}
+      />,
+    );
+
+    rerender(
+      <Theme>
+        <BranchSelector
+          repoPath="owner/repo"
+          currentBranch={null}
+          defaultBranch="main"
+          workspaceMode="cloud"
+          cloudBranches={["main", "feature-x"]}
+          cloudBranchesLoading={false}
+          cloudSearchQuery=""
+          selectedBranch="feature-x"
+          onBranchSelect={onBranchSelect}
+          onCloudSearchChange={vi.fn()}
+        />
+      </Theme>,
+    );
+
+    expect(onBranchSelect).not.toHaveBeenCalled();
+  });
+
   it("surfaces the 'Use input as branch name' action when the typed value is new", async () => {
     const user = userEvent.setup();
     renderInTheme(

packages/ui/src/features/git-interaction/components/BranchSelector.tsx

@@ -40,6 +40,10 @@ import { getSuggestedBranchName } from "../utils/getSuggestedBranchName";
 
 const COMBOBOX_LIMIT = 50;
 
+// Shared so the two "still loading branches" render sites (the empty-list
+// spinner and the seeded-default row) can never drift out of sync on a copy edit.
+const LOADING_BRANCHES_LABEL = "Loading branches…";
+
 // Sentinel value for the "Create new branch" action. Rendered as a real
 // ComboboxItem in the list footer so it's reachable by keyboard, not a
 // plain button the combobox's roving focus skips over.
@@ -140,8 +144,19 @@ export function BranchSelector({
   const isSelectionOnly = workspaceMode === "worktree" || isCloudMode;
   const displayedBranch = isSelectionOnly ? selectedBranch : currentBranch;
 
+  // The branch we auto-selected, so we can tell our own pick apart from one the
+  // user made. Lets us correct a stale default (e.g. a cached "trunk" that the
+  // live list later contradicts) without ever clobbering a deliberate choice.
+  const autoSelectedBranchRef = useRef<string | null>(null);
   useEffect(() => {
-    if (isSelectionOnly && defaultBranch && !selectedBranch && onBranchSelect) {
+    if (!isSelectionOnly || !defaultBranch || !onBranchSelect) return;
+    // Adopt the default when nothing is selected yet, or when the default has
+    // changed out from under a value we ourselves auto-selected — but leave a
+    // user's own selection alone.
+    const selectionIsOurs =
+      !selectedBranch || selectedBranch === autoSelectedBranchRef.current;
+    if (selectionIsOurs && selectedBranch !== defaultBranch) {
+      autoSelectedBranchRef.current = defaultBranch;
       onBranchSelect(defaultBranch);
     }
   }, [isSelectionOnly, defaultBranch, selectedBranch, onBranchSelect]);
@@ -175,12 +190,28 @@ export function BranchSelector({
     return byBranch;
   }, [repoCheckouts]);
 
-  const branches = isCloudMode ? (cloudBranches ?? []) : localBranches;
+  const liveBranches = isCloudMode ? (cloudBranches ?? []) : localBranches;
   const effectiveLoading = loading || (isCloudMode && cloudBranchesLoading);
   const branchListLoading = isCloudMode
     ? !!cloudBranchesLoading
     : localBranchesLoading;
 
+  // On a cold start the live cloud branch list is still empty while the (slow)
+  // remote fetch runs. Surface the known default ("trunk") branch as a real
+  // list item straight away — with a loading row rendered below it — so the
+  // common "start on trunk" case is pickable with zero wait. Only when there's
+  // no active search: once the user types, the results should be purely what
+  // the remote returns.
+  const seededDefaultBranch =
+    isCloudMode &&
+    branchListLoading &&
+    liveBranches.length === 0 &&
+    !!defaultBranch &&
+    !(cloudSearchQuery ?? "").trim()
+      ? defaultBranch
+      : null;
+  const branches = seededDefaultBranch ? [seededDefaultBranch] : liveBranches;
+
   const checkoutMutation = useMutation({
     ...trpc.git.checkoutBranch.mutationOptions(),
     onSuccess: () => {
@@ -456,7 +487,7 @@ export function BranchSelector({
         ) : null}
 
         {branchListLoading && branches.length === 0 ? (
-          <LoadingRow label="Loading branches…" />
+          <LoadingRow label={LOADING_BRANCHES_LABEL} />
         ) : (
           <ComboboxEmpty>No branches found.</ComboboxEmpty>
         )}
@@ -519,6 +550,15 @@ export function BranchSelector({
           }}
         </ComboboxList>
 
+        {/*
+          Cold start: the default ("trunk") branch is seeded as the only list
+          item while the remote list loads. A loading row directly below it
+          makes clear the rest of the branches are still on the way.
+        */}
+        {seededDefaultBranch ? (
+          <LoadingRow label={LOADING_BRANCHES_LABEL} />
+        ) : null}
+
         {isCloudMode && cloudBranchesHasMore ? (
           <ComboboxListFooter>
             <div className="px-2 pb-2">

packages/ui/src/features/settings/settingsStore.test.ts

@@ -58,6 +58,7 @@ describe("feature settingsStore cloud selections", () => {
       allowBypassPermissions: false,
       lastUsedCloudRepository: null,
       cachedCloudRepositoryMap: {},
+      cachedCloudDefaultBranchMap: {},
     });
   });
 
@@ -141,6 +142,63 @@ describe("feature settingsStore cloud selections", () => {
     });
   });
 
+  it("caches and persists the cloud default branch per repo", async () => {
+    useSettingsStore
+      .getState()
+      .setCachedCloudDefaultBranch("posthog/posthog", "master");
+    useSettingsStore
+      .getState()
+      .setCachedCloudDefaultBranch("posthog/code", "main");
+
+    expect(useSettingsStore.getState().cachedCloudDefaultBranchMap).toEqual({
+      "posthog/posthog": "master",
+      "posthog/code": "main",
+    });
+
+    await waitForPersistedWrite();
+
+    const lastCall = setItem.mock.calls[setItem.mock.calls.length - 1];
+    const persisted = JSON.parse(lastCall[1]);
+
+    expect(persisted.state.cachedCloudDefaultBranchMap).toEqual({
+      "posthog/posthog": "master",
+      "posthog/code": "main",
+    });
+  });
+
+  it("keeps the same map reference when the default branch is unchanged", () => {
+    useSettingsStore
+      .getState()
+      .setCachedCloudDefaultBranch("posthog/code", "main");
+    const first = useSettingsStore.getState().cachedCloudDefaultBranchMap;
+
+    useSettingsStore
+      .getState()
+      .setCachedCloudDefaultBranch("posthog/code", "main");
+    const second = useSettingsStore.getState().cachedCloudDefaultBranchMap;
+
+    expect(second).toBe(first);
+  });
+
+  it("rehydrates the cached cloud default branch map", async () => {
+    getItem.mockResolvedValue(
+      JSON.stringify({
+        state: {
+          cachedCloudDefaultBranchMap: { "posthog/code": "main" },
+        },
+        version: 0,
+      }),
+    );
+
+    useSettingsStore.setState({ cachedCloudDefaultBranchMap: {} });
+
+    await useSettingsStore.persist.rehydrate();
+
+    expect(useSettingsStore.getState().cachedCloudDefaultBranchMap).toEqual({
+      "posthog/code": "main",
+    });
+  });
+
   it("rehydrates the unsafe mode toggle", async () => {
     getItem.mockResolvedValue(
       JSON.stringify({

packages/ui/src/features/settings/settingsStore.ts

@@ -107,6 +107,10 @@ interface SettingsStore {
   lastUsedReasoningEffort: string | null;
   lastUsedCloudRepository: string | null;
   cachedCloudRepositoryMap: Record<string, UserRepositoryIntegrationRef>;
+  // Last-known default ("trunk") branch per cloud repo, keyed by lowercased
+  // "owner/repo". Persisted so a cold start can pre-select trunk in the branch
+  // picker immediately, before the (slow) live branch list resolves.
+  cachedCloudDefaultBranchMap: Record<string, string>;
   lastUsedEnvironments: Record<string, string>;
   defaultInitialTaskMode: DefaultInitialTaskMode;
   lastUsedInitialTaskMode: ExecutionMode;
@@ -126,6 +130,7 @@ interface SettingsStore {
   setCachedCloudRepositoryMap: (
     map: Record<string, UserRepositoryIntegrationRef>,
   ) => void;
+  setCachedCloudDefaultBranch: (repo: string, branch: string) => void;
   setLastUsedEnvironment: (
     repoPath: string,
     environmentId: string | null,
@@ -289,6 +294,7 @@ export const useSettingsStore = create<SettingsStore>()(
       lastUsedReasoningEffort: null,
       lastUsedCloudRepository: null,
       cachedCloudRepositoryMap: {},
+      cachedCloudDefaultBranchMap: {},
       lastUsedEnvironments: {},
       defaultInitialTaskMode: "plan",
       lastUsedInitialTaskMode: "plan",
@@ -308,6 +314,16 @@ export const useSettingsStore = create<SettingsStore>()(
         set({ lastUsedCloudRepository: repo }),
       setCachedCloudRepositoryMap: (map) =>
         set({ cachedCloudRepositoryMap: map }),
+      setCachedCloudDefaultBranch: (repo, branch) =>
+        set((state) => {
+          if (state.cachedCloudDefaultBranchMap[repo] === branch) return {};
+          return {
+            cachedCloudDefaultBranchMap: {
+              ...state.cachedCloudDefaultBranchMap,
+              [repo]: branch,
+            },
+          };
+        }),
       setLastUsedEnvironment: (repoPath, environmentId) =>
         set((state) => {
           const next = { ...state.lastUsedEnvironments };
@@ -498,6 +514,7 @@ export const useSettingsStore = create<SettingsStore>()(
         lastUsedReasoningEffort: state.lastUsedReasoningEffort,
         lastUsedCloudRepository: state.lastUsedCloudRepository,
         cachedCloudRepositoryMap: state.cachedCloudRepositoryMap,
+        cachedCloudDefaultBranchMap: state.cachedCloudDefaultBranchMap,
         lastUsedEnvironments: state.lastUsedEnvironments,
         defaultInitialTaskMode: state.defaultInitialTaskMode,
         lastUsedInitialTaskMode: state.lastUsedInitialTaskMode,

packages/ui/src/features/task-detail/components/TaskInput.tsx

@@ -189,6 +189,8 @@ export function TaskInput({
     setLastUsedAdapter,
     lastUsedCloudRepository,
     setLastUsedCloudRepository,
+    cachedCloudDefaultBranchMap,
+    setCachedCloudDefaultBranch,
     allowBypassPermissions,
     setLastUsedEnvironment,
     getLastUsedEnvironment,
@@ -430,7 +432,33 @@ export function TaskInput({
     cloudBranchSearchQuery,
   );
   const cloudBranches = cloudBranchData?.branches;
-  const cloudDefaultBranch = cloudBranchData?.defaultBranch ?? null;
+  const liveCloudDefaultBranch = cloudBranchData?.defaultBranch ?? null;
+  // Serve the persisted default branch until the live list resolves, so the
+  // majority "start on trunk" case pre-selects trunk with zero wait on a cold
+  // start. The cached value is best-effort: if it's stale (a default branch
+  // renamed since it was cached), `cloudDefaultBranch` switches to the live
+  // value on arrival and BranchSelector re-selects it — as long as the user
+  // hasn't picked a branch of their own in the meantime.
+  const cloudDefaultBranch =
+    liveCloudDefaultBranch ??
+    (selectedCloudRepository
+      ? (cachedCloudDefaultBranchMap[selectedCloudRepository] ?? null)
+      : null);
+
+  // Persist the freshly loaded default branch so the next cold start can
+  // pre-select trunk immediately.
+  useEffect(() => {
+    if (selectedCloudRepository && liveCloudDefaultBranch) {
+      setCachedCloudDefaultBranch(
+        selectedCloudRepository,
+        liveCloudDefaultBranch,
+      );
+    }
+  }, [
+    selectedCloudRepository,
+    liveCloudDefaultBranch,
+    setCachedCloudDefaultBranch,
+  ]);
 
   const {
     branchOpen,