Fix optimistic workspace creation on Android

packages/app/src/utils/workspace-navigation.test.ts

@@ -1,12 +1,31 @@
 import { beforeEach, describe, expect, it, vi } from "vitest";
 
-vi.mock("expo-router", () => ({
-  router: {
+const { platformState, routerMock } = vi.hoisted(() => ({
+  platformState: {
+    isNative: false,
+    isWeb: true,
+  },
+  routerMock: {
+    back: vi.fn(),
+    canGoBack: vi.fn(() => false),
     navigate: vi.fn(),
     replace: vi.fn(),
   },
 }));
 
+vi.mock("expo-router", () => ({
+  router: routerMock,
+}));
+
+vi.mock("@/constants/platform", () => ({
+  get isNative() {
+    return platformState.isNative;
+  },
+  get isWeb() {
+    return platformState.isWeb;
+  },
+}));
+
 vi.mock("@react-native-async-storage/async-storage", () => {
   const storage = new Map<string, string>();
   return {
@@ -23,7 +42,11 @@ vi.mock("@react-native-async-storage/async-storage", () => {
 });
 
 import { useWorkspaceLayoutStore } from "@/stores/workspace-layout-store";
-import { prepareWorkspaceTab } from "@/utils/workspace-navigation";
+import {
+  getNavigationActiveWorkspaceSelection,
+  syncNavigationActiveWorkspace,
+} from "@/stores/navigation-active-workspace-store";
+import { navigateToPreparedWorkspaceTab, prepareWorkspaceTab } from "@/utils/workspace-navigation";
 
 const SERVER_ID = "server-1";
 const WORKSPACE_ID = "/repo/worktree";
@@ -31,6 +54,15 @@ const AGENT_ID = "agent-1";
 
 describe("prepareWorkspaceTab", () => {
   beforeEach(() => {
+    vi.useRealTimers();
+    platformState.isNative = false;
+    platformState.isWeb = true;
+    routerMock.back.mockReset();
+    routerMock.canGoBack.mockReset();
+    routerMock.canGoBack.mockReturnValue(false);
+    routerMock.navigate.mockReset();
+    routerMock.replace.mockReset();
+    syncNavigationActiveWorkspace({ current: null });
     useWorkspaceLayoutStore.setState({
       layoutByWorkspace: {},
       splitSizesByWorkspace: {},
@@ -50,4 +82,37 @@ describe("prepareWorkspaceTab", () => {
     const key = "server-1:/repo/worktree";
     expect(useWorkspaceLayoutStore.getState().getWorkspaceTabs(key)).toHaveLength(1);
   });
+
+  it("pops back to the retained workspace shell for native replace navigation", () => {
+    vi.useFakeTimers();
+    platformState.isNative = true;
+    platformState.isWeb = false;
+    routerMock.canGoBack.mockReturnValue(true);
+
+    syncNavigationActiveWorkspace({
+      current: {
+        getCurrentRoute: () => ({
+          path: "/h/server-1/workspace/source-workspace",
+        }),
+      },
+    });
+
+    const route = navigateToPreparedWorkspaceTab({
+      serverId: SERVER_ID,
+      workspaceId: WORKSPACE_ID,
+      target: { kind: "agent", agentId: AGENT_ID },
+      navigationMethod: "replace",
+    });
+
+    expect(route).toBe("/h/server-1/workspace/b64_L3JlcG8vd29ya3RyZWU");
+    expect(routerMock.back).toHaveBeenCalledOnce();
+    expect(routerMock.replace).not.toHaveBeenCalled();
+
+    vi.runAllTimers();
+
+    expect(getNavigationActiveWorkspaceSelection()).toEqual({
+      serverId: SERVER_ID,
+      workspaceId: WORKSPACE_ID,
+    });
+  });
 });

packages/app/src/utils/workspace-navigation.ts

@@ -1,4 +1,10 @@
 import { router } from "expo-router";
+import { isNative } from "@/constants/platform";
+import {
+  activateNavigationWorkspaceSelection,
+  getLastNavigationWorkspaceRouteSelection,
+  overrideNextNavigationWorkspaceRouteSelection,
+} from "@/stores/navigation-active-workspace-store";
 import { useWorkspaceLayoutStore } from "@/stores/workspace-layout-store";
 import { generateDraftId } from "@/stores/draft-keys";
 import {
@@ -45,6 +51,20 @@ export function prepareWorkspaceTab(input: PrepareWorkspaceTabInput) {
 export function navigateToPreparedWorkspaceTab(input: NavigateToPreparedWorkspaceTabInput): string {
   const route = prepareWorkspaceTab(input);
   if (input.navigationMethod === "replace") {
+    const canReturnToWorkspaceShell =
+      isNative && getLastNavigationWorkspaceRouteSelection() !== null && router.canGoBack();
+    if (canReturnToWorkspaceShell) {
+      const nextSelection = {
+        serverId: input.serverId,
+        workspaceId: input.workspaceId,
+      };
+      overrideNextNavigationWorkspaceRouteSelection(nextSelection);
+      router.back();
+      setTimeout(() => {
+        activateNavigationWorkspaceSelection(nextSelection);
+      }, 0);
+      return route;
+    }
     router.replace(route as any);
   } else {
     router.navigate(route as any);