fix(mobile): a draft route's branch checkout runs once, even across a reconnect

apps/mobile/src/App.tsx

@@ -9,10 +9,10 @@ import { GestureHandlerRootView } from "react-native-gesture-handler";
 import { KeyboardProvider } from "react-native-keyboard-controller";
 import { SafeAreaProvider } from "react-native-safe-area-context";
 import {
-  createNavigationContainerRef,
   createStaticNavigation,
   getActionFromState,
   getStateFromPath,
+  type NavigationState,
 } from "@react-navigation/native";
 
 import { RegistryContext } from "@effect/atom-react";
@@ -52,6 +52,8 @@ import {
 } from "./features/workspace/workspace-selection-store";
 import { workspaceDetailLinkAction } from "./features/workspace/workspace-detail-selection";
 import { NativeLayoutMetricsProvider } from "./native/native-layout-metrics";
+import { attachRootNavigation, rootNavigationRef } from "./lib/rootNavigation";
+import { pruneNewTaskBranchCheckouts } from "./features/threads/NewTaskDraftRouteScreen";
 
 import "../global.css";
 
@@ -84,7 +86,15 @@ const appLinking = {
 };
 
 const Navigation = createStaticNavigation(RootStack);
-const navigationRef = createNavigationContainerRef();
+
+function handleNavigationStateChange(state: NavigationState | undefined): void {
+  recordWorkspaceSelections(state);
+  if (state) pruneNewTaskBranchCheckouts(state);
+}
+
+function handleNavigationReady(): void {
+  handleNavigationStateChange(rootNavigationRef.getRootState());
+}
 
 declare const module: { readonly hot?: FoundationHotModule } | undefined;
 
@@ -148,11 +158,11 @@ function AppContent() {
               <NativeLayoutMetricsProvider>
                 <IncomingShareProvider>
                   <Navigation
-                    ref={navigationRef}
+                    ref={attachRootNavigation}
                     linking={appLinking}
-                    onReady={() => recordWorkspaceSelections(navigationRef.getRootState())}
+                    onReady={handleNavigationReady}
                     theme={navigationTheme}
-                    onStateChange={recordWorkspaceSelections}
+                    onStateChange={handleNavigationStateChange}
                   />
                 </IncomingShareProvider>
                 <Dialog.Host />

apps/mobile/src/features/layout/AdaptiveWorkspaceLayout.tsx

@@ -1,6 +1,6 @@
 import { SafeAreaInsetsContext, useSafeAreaInsets } from "react-native-safe-area-context";
 import { AdaptiveWorkspaceContext, useAdaptiveWorkspaceLayout } from "./adaptive-workspace-context";
-import { createDeferredNavigation } from "./deferred-navigation";
+import { threadSelectionDeferral } from "./deferred-navigation";
 import { uuidv4 } from "../../lib/uuid";
 import type {
   EnvironmentProject,
@@ -491,10 +491,9 @@ function AdaptiveWorkspaceLayoutContent(
   };
   const handleOpenSettings = () => openWorkspaceTab("settings");
 
-  const deferredThreadSelection = useMemo(() => createDeferredNavigation(), []);
   /** The one way to show a page in the detail column; see `selectWorkspaceDetail`. */
   function handleSelectDetail(detail: WorkspaceDetail, sheet?: WorkspaceDetailSheet) {
-    deferredThreadSelection.run(() => {
+    threadSelectionDeferral.run(() => {
       if (Platform.OS === "android") Keyboard.dismiss();
       setFileInspectorPreferredVisible(false);
       if (layout.primarySidebarOverlaysContent) setPrimarySidebarVisible(false);
@@ -770,7 +769,6 @@ function AdaptiveWorkspaceLayoutContent(
     fileInspector,
     primarySidebarSearchQuery,
     startNewTask: handleStartNewTask,
-    deferThreadSelectionUntil: deferredThreadSelection.waitFor,
     selectDetail: handleSelectDetail,
     activateAuxiliaryPaneRole,
     registerWorkspaceInspector,

apps/mobile/src/features/layout/adaptive-workspace-context.ts

@@ -20,7 +20,6 @@ interface AdaptiveWorkspaceContextValue {
   readonly fileInspector: FileInspectorPaneLayout;
   readonly primarySidebarSearchQuery: string;
   readonly startNewTask: () => void;
-  readonly deferThreadSelectionUntil: (operation: Promise<unknown>) => void;
   /** Use `useSelectWorkspaceDetail`, which also works outside the layout. */
   readonly selectDetail?: (detail: WorkspaceDetail, sheet?: WorkspaceDetailSheet) => void;
   readonly activateAuxiliaryPaneRole: (role: WorkspaceAuxiliaryPaneRole) => () => void;
@@ -67,7 +66,6 @@ export const AdaptiveWorkspaceContext = createContext<AdaptiveWorkspaceContextVa
   fileInspector: compactFileInspector,
   primarySidebarSearchQuery: "",
   startNewTask: () => undefined,
-  deferThreadSelectionUntil: () => undefined,
   activateAuxiliaryPaneRole: () => () => undefined,
   registerWorkspaceInspector: () => () => undefined,
   setPrimarySidebarSearchQuery: () => undefined,

apps/mobile/src/features/layout/deferred-navigation.ts

@@ -20,3 +20,9 @@ export function createDeferredNavigation() {
     },
   };
 }
+
+/**
+ * The app's one thread-selection gate: a new-task branch checkout holds
+ * split-view thread selection until the draft route can be left.
+ */
+export const threadSelectionDeferral = createDeferredNavigation();

apps/mobile/src/features/threads/NewTaskDraftRouteScreen.test.tsx

@@ -0,0 +1,177 @@
+import type { ReactNode } from "react";
+import { act, useSyncExternalStore } from "react";
+import { createRoot } from "react-dom/client";
+import { AsyncResult } from "effect/reactivity";
+import { expect, it, vi } from "vite-plus/test";
+
+interface Snapshot {
+  readonly projects: ReadonlyArray<{
+    readonly environmentId: string;
+    readonly id: string;
+    readonly workspaceRoot: string;
+  }>;
+  readonly catalog: {
+    readonly isLoadingConnections: boolean;
+    readonly hasLoadedShellSnapshot: boolean;
+    readonly hasConnectingEnvironment: boolean;
+    readonly connectionError: string | null;
+  };
+}
+
+const harness = vi.hoisted(() => {
+  const listeners = new Set<() => void>();
+  const state: { snapshot: Snapshot | null } = { snapshot: null };
+  return {
+    state,
+    set(next: Snapshot) {
+      state.snapshot = next;
+      for (const listener of listeners) listener();
+    },
+    subscribe(listener: () => void) {
+      listeners.add(listener);
+      return () => listeners.delete(listener);
+    },
+    route: { key: "thread-route", name: "Thread", params: {} as object },
+    reset: () => undefined,
+    checkout: vi.fn(),
+    draftProjectRefs: [] as unknown[],
+    // Hook results stay stable across renders, as the real hooks' are. Only
+    // main's screen reads these (see the main-comparison mocks below).
+    layout: { deferThreadSelectionUntil: () => undefined },
+    navigation: { goBack: () => undefined },
+    preventContext: { preventedRoutes: {} },
+  };
+});
+
+function useSnapshot(): Snapshot {
+  return useSyncExternalStore(harness.subscribe, () => harness.state.snapshot!);
+}
+
+vi.mock("react-native", () => ({
+  View: ({ children }: { children: ReactNode }) => <div>{children}</div>,
+}));
+vi.mock("../../components/ui/AppText", () => ({
+  AppText: ({ children }: { children: ReactNode }) => <span>{children}</span>,
+}));
+vi.mock("../../components/ui/Dialog", () => ({ alert: vi.fn() }));
+vi.mock("../../native/StackHeader", () => ({ NativeStackScreenOptions: () => null }));
+vi.mock("./ThreadScreenHost", () => ({
+  ThreadScreenFallback: ({ children }: { children: ReactNode }) => <>{children}</>,
+}));
+vi.mock("./NewTaskDraftScreen", () => ({
+  NewTaskDraftScreen: (props: { initialProjectRef: { branch?: string | null } }) => {
+    harness.draftProjectRefs.push(props.initialProjectRef);
+    return <p>draft on {props.initialProjectRef.branch}</p>;
+  },
+}));
+vi.mock("./new-task-context-presentation", () => ({
+  resolveNewTaskBranchWorktreePath: () => null,
+}));
+vi.mock("./new-task-flow-provider", () => ({
+  useNewTaskFlow: () => harness,
+}));
+vi.mock("./checkout-new-task-branch", () => ({ checkoutNewTaskBranch: harness.checkout }));
+// Main-comparison mocks: this screen no longer uses them, but main's screen
+// did (deferThreadSelectionUntil, useAtomCommand, useNavigation). They stay so
+// this test can run against main's screen and fail there for the right reason
+// (a second checkout) instead of on an unmockable import.
+vi.mock("../layout/AdaptiveWorkspaceLayout", () => ({
+  useAdaptiveWorkspaceLayout: () => harness.layout,
+}));
+vi.mock("../../state/use-atom-command", () => ({ useAtomCommand: () => harness.checkout }));
+vi.mock("../../state/entities", () => ({ useProjects: () => useSnapshot().projects }));
+vi.mock("../../state/workspace", () => ({
+  useWorkspaceState: () => ({ state: useSnapshot().catalog }),
+}));
+vi.mock("../../state/vcs", () => ({ vcsEnvironment: { readRefs: {}, switchRef: {} } }));
+vi.mock("@react-navigation/native", () => ({
+  CommonActions: { goBack: () => ({ type: "GO_BACK" }) },
+  createNavigationContainerRef: () => ({ current: null }),
+  // Main-comparison only, like the mocks above.
+  useNavigation: () => harness.navigation,
+  useRoute: () => harness.route,
+  usePreventRemove: () => undefined,
+  usePreventRemoveContext: () => harness.preventContext,
+}));
+
+import { NewTaskDraftRouteScreen } from "./NewTaskDraftRouteScreen";
+
+const project = { environmentId: "env-1", id: "project-1", workspaceRoot: "/repo" };
+const connected: Snapshot["catalog"] = {
+  isLoadingConnections: false,
+  hasLoadedShellSnapshot: true,
+  hasConnectingEnvironment: false,
+  connectionError: null,
+};
+
+it("checks the branch out once when a reconnect drops and restores the project", async () => {
+  (globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true;
+  const firstCheckout = Promise.withResolvers<unknown>();
+  harness.checkout.mockReturnValueOnce(firstCheckout.promise);
+  harness.checkout.mockResolvedValue(
+    AsyncResult.success({ name: "feature", current: true, isDefault: false, worktreePath: null }),
+  );
+  harness.set({ projects: [project], catalog: connected });
+  const draft = { environmentId: "env-1", projectId: "project-1", branch: "feature" };
+  harness.route = { key: "thread-route", name: "Thread", params: { draft } };
+
+  const container = document.createElement("div");
+  const root = createRoot(container);
+  using _root = { [Symbol.dispose]: () => act(() => root.unmount()) };
+  act(() =>
+    root.render(
+      <NewTaskDraftRouteScreen
+        route={{ key: "thread-route", name: "Thread", params: draft } as never}
+      />,
+    ),
+  );
+  await vi.waitFor(() => expect(harness.checkout).toHaveBeenCalledTimes(1));
+
+  // The connection drops while the checkout runs, then the project comes back.
+  act(() => harness.set({ projects: [], catalog: { ...connected, isLoadingConnections: true } }));
+  act(() => harness.set({ projects: [project], catalog: connected }));
+  await act(async () => {
+    firstCheckout.resolve(
+      AsyncResult.success({ name: "feature", current: true, isDefault: false, worktreePath: null }),
+    );
+  });
+
+  await vi.waitFor(() =>
+    expect(
+      harness.checkout.mock.calls.length > 1 || container.textContent === "draft on feature",
+    ).toBe(true),
+  );
+  expect(harness.checkout).toHaveBeenCalledTimes(1);
+  expect(container.textContent).toBe("draft on feature");
+});
+
+it("hands the draft screen the same project ref across re-renders", async () => {
+  (globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true;
+  harness.checkout.mockReset();
+  harness.draftProjectRefs.length = 0;
+  harness.checkout.mockResolvedValue(
+    AsyncResult.success({ name: "other", current: true, isDefault: false, worktreePath: null }),
+  );
+  harness.set({ projects: [project], catalog: connected });
+  const draft = { environmentId: "env-1", projectId: "project-1", branch: "other" };
+  harness.route = { key: "stable-route", name: "Thread", params: { draft } };
+
+  const container = document.createElement("div");
+  const root = createRoot(container);
+  using _root = { [Symbol.dispose]: () => act(() => root.unmount()) };
+  act(() =>
+    root.render(
+      <NewTaskDraftRouteScreen
+        route={{ key: "stable-route", name: "Thread", params: draft } as never}
+      />,
+    ),
+  );
+  await vi.waitFor(() => expect(container.textContent).toBe("draft on other"));
+
+  // The draft screen re-applies its route project when this identity changes,
+  // and that write re-renders the route; an unstable ref loops forever.
+  act(() => harness.set({ projects: [{ ...project }], catalog: { ...connected } }));
+  const refs = new Set(harness.draftProjectRefs);
+  expect(harness.draftProjectRefs.length).toBeGreaterThan(1);
+  expect(refs.size).toBe(1);
+});