Keep New Workspace on the current project (#1806)

* fix(app): keep New Workspace on the active project

New Workspace preselection was falling back to stale remembered host state when multiple hosts existed. Pass durable active-project context from workspace entry points and centralize initial host selection so stale offline hosts cannot steal the form while startup restore remains last-workspace only.

* fix(app): address New Workspace review feedback

* fix(app): stabilize New Workspace host selection

* fix(app): preserve New Workspace host fallback

* test(app): loosen stale-host project assertion

* fix(app): prefer reachable new workspace hosts

* fix(app): avoid stale offline workspace preselect

* fix(app): prefer online project hosts

packages/app/src/screens/new-workspace-initial-context.test.ts

@@ -0,0 +1,283 @@
+import { describe, expect, it } from "vitest";
+import type { HostProjectListItem } from "@/projects/host-projects";
+import type { HostRuntimeConnectionStatus } from "@/runtime/host-runtime";
+import {
+  resolveNewWorkspaceAutomaticServerId,
+  resolveNewWorkspaceInitialServerId,
+} from "./new-workspace-initial-context";
+
+function projectFor(serverId: string, key = "project"): HostProjectListItem {
+  return {
+    projectKey: key,
+    projectName: key,
+    projectKind: "git",
+    iconWorkingDir: `/work/${key}`,
+    hosts: [{ serverId, iconWorkingDir: `/work/${key}`, canCreateWorktree: true }],
+    workspaceKeys: [],
+  };
+}
+
+function statuses(
+  entries: Record<string, HostRuntimeConnectionStatus>,
+): ReadonlyMap<string, HostRuntimeConnectionStatus> {
+  return new Map(Object.entries(entries));
+}
+
+function multiplicity(entries: Record<string, boolean> = {}): ReadonlyMap<string, boolean> {
+  return new Map(Object.entries(entries));
+}
+
+describe("resolveNewWorkspaceInitialServerId", () => {
+  it("prefers explicit route host context over online-host fallback", () => {
+    expect(
+      resolveNewWorkspaceInitialServerId({
+        allServerIds: ["offline", "online"],
+        routeServerId: "offline",
+        lastActiveProject: null,
+        projects: [projectFor("online")],
+        hostConnectionStatusByServerId: statuses({ offline: "offline", online: "online" }),
+        workspaceMultiplicityByServerId: multiplicity(),
+      }),
+    ).toBe("offline");
+  });
+
+  it("prefers the sole online host over a stale offline project", () => {
+    expect(
+      resolveNewWorkspaceInitialServerId({
+        allServerIds: ["offline", "online"],
+        routeServerId: null,
+        lastActiveProject: projectFor("offline"),
+        projects: [projectFor("offline")],
+        hostConnectionStatusByServerId: statuses({ offline: "offline", online: "online" }),
+        workspaceMultiplicityByServerId: multiplicity(),
+      }),
+    ).toBe("online");
+
+    expect(
+      resolveNewWorkspaceInitialServerId({
+        allServerIds: ["offline", "online"],
+        routeServerId: null,
+        lastActiveProject: null,
+        projects: [projectFor("online")],
+        hostConnectionStatusByServerId: statuses({ offline: "offline", online: "online" }),
+        workspaceMultiplicityByServerId: multiplicity(),
+      }),
+    ).toBe("online");
+  });
+
+  it("uses the last active project when there is no sole online host", () => {
+    expect(
+      resolveNewWorkspaceInitialServerId({
+        allServerIds: ["offline", "other"],
+        routeServerId: null,
+        lastActiveProject: projectFor("offline"),
+        projects: [projectFor("offline")],
+        hostConnectionStatusByServerId: statuses({ offline: "offline", other: "offline" }),
+        workspaceMultiplicityByServerId: multiplicity(),
+      }),
+    ).toBe("offline");
+  });
+
+  it("prefers a connecting project host over a stale offline last active project", () => {
+    expect(
+      resolveNewWorkspaceInitialServerId({
+        allServerIds: ["offline", "connecting"],
+        routeServerId: null,
+        lastActiveProject: projectFor("offline", "remembered"),
+        projects: [projectFor("offline", "remembered"), projectFor("connecting", "current")],
+        hostConnectionStatusByServerId: statuses({
+          offline: "offline",
+          connecting: "connecting",
+        }),
+        workspaceMultiplicityByServerId: multiplicity(),
+      }),
+    ).toBe("connecting");
+  });
+
+  it("prefers the online last active project over another hydrated online project", () => {
+    expect(
+      resolveNewWorkspaceInitialServerId({
+        allServerIds: ["host-a", "host-b"],
+        routeServerId: null,
+        lastActiveProject: projectFor("host-b", "remembered"),
+        projects: [projectFor("host-a")],
+        hostConnectionStatusByServerId: statuses({
+          "host-a": "online",
+          "host-b": "online",
+        }),
+        workspaceMultiplicityByServerId: multiplicity(),
+      }),
+    ).toBe("host-b");
+  });
+
+  it("falls back to the only online host even before projects have hydrated", () => {
+    expect(
+      resolveNewWorkspaceInitialServerId({
+        allServerIds: ["offline-a", "online", "offline-b"],
+        routeServerId: null,
+        lastActiveProject: null,
+        projects: [],
+        hostConnectionStatusByServerId: statuses({
+          "offline-a": "offline",
+          online: "online",
+          "offline-b": "offline",
+        }),
+        workspaceMultiplicityByServerId: multiplicity(),
+      }),
+    ).toBe("online");
+  });
+
+  it("prefers an online host over the only cached offline project", () => {
+    expect(
+      resolveNewWorkspaceInitialServerId({
+        allServerIds: ["offline", "online-a", "online-b"],
+        routeServerId: null,
+        lastActiveProject: projectFor("offline"),
+        projects: [projectFor("offline")],
+        hostConnectionStatusByServerId: statuses({
+          offline: "offline",
+          "online-a": "online",
+          "online-b": "online",
+        }),
+        workspaceMultiplicityByServerId: multiplicity(),
+      }),
+    ).toBe("online-a");
+  });
+
+  it("prefers the first online project host over an empty online host", () => {
+    expect(
+      resolveNewWorkspaceInitialServerId({
+        allServerIds: ["empty-online", "project-online-a", "project-online-b"],
+        routeServerId: null,
+        lastActiveProject: null,
+        projects: [projectFor("project-online-a", "a"), projectFor("project-online-b", "b")],
+        hostConnectionStatusByServerId: statuses({
+          "empty-online": "online",
+          "project-online-a": "online",
+          "project-online-b": "online",
+        }),
+        workspaceMultiplicityByServerId: multiplicity(),
+      }),
+    ).toBe("project-online-a");
+  });
+
+  it("uses the only host with selectable projects even before runtime status is online", () => {
+    expect(
+      resolveNewWorkspaceInitialServerId({
+        allServerIds: ["offline", "connected"],
+        routeServerId: null,
+        lastActiveProject: null,
+        projects: [projectFor("connected")],
+        hostConnectionStatusByServerId: statuses({
+          offline: "offline",
+          connected: "connecting",
+        }),
+        workspaceMultiplicityByServerId: multiplicity(),
+      }),
+    ).toBe("connected");
+  });
+});
+
+describe("resolveNewWorkspaceAutomaticServerId", () => {
+  it("keeps a usable automatic host stable when the computed default changes", () => {
+    expect(
+      resolveNewWorkspaceAutomaticServerId({
+        allServerIds: ["host-a", "host-b"],
+        routeServerId: null,
+        lastActiveProject: null,
+        projects: [projectFor("host-a"), projectFor("host-b")],
+        hostConnectionStatusByServerId: statuses({ "host-a": "online", "host-b": "online" }),
+        workspaceMultiplicityByServerId: multiplicity(),
+        currentServerId: "host-a",
+        nextServerId: "host-b",
+      }),
+    ).toBe("host-a");
+  });
+
+  it("switches to the remembered online host after it hydrates", () => {
+    expect(
+      resolveNewWorkspaceAutomaticServerId({
+        allServerIds: ["host-a", "host-b"],
+        routeServerId: null,
+        lastActiveProject: projectFor("host-b", "remembered"),
+        projects: [projectFor("host-a"), projectFor("host-b", "remembered")],
+        hostConnectionStatusByServerId: statuses({
+          "host-a": "online",
+          "host-b": "online",
+        }),
+        workspaceMultiplicityByServerId: multiplicity(),
+        currentServerId: "host-a",
+        nextServerId: "host-b",
+      }),
+    ).toBe("host-b");
+  });
+
+  it("switches from an offline automatic host to the online default", () => {
+    expect(
+      resolveNewWorkspaceAutomaticServerId({
+        allServerIds: ["offline", "online"],
+        routeServerId: null,
+        lastActiveProject: null,
+        projects: [projectFor("offline")],
+        hostConnectionStatusByServerId: statuses({ offline: "offline", online: "online" }),
+        workspaceMultiplicityByServerId: multiplicity(),
+        currentServerId: "offline",
+        nextServerId: "online",
+      }),
+    ).toBe("online");
+  });
+
+  it("switches from an offline automatic host to a connecting default with projects", () => {
+    expect(
+      resolveNewWorkspaceAutomaticServerId({
+        allServerIds: ["offline", "connecting"],
+        routeServerId: null,
+        lastActiveProject: projectFor("offline", "remembered"),
+        projects: [projectFor("offline", "remembered"), projectFor("connecting", "current")],
+        hostConnectionStatusByServerId: statuses({
+          offline: "offline",
+          connecting: "connecting",
+        }),
+        workspaceMultiplicityByServerId: multiplicity(),
+        currentServerId: "offline",
+        nextServerId: "connecting",
+      }),
+    ).toBe("connecting");
+  });
+
+  it("switches to the default when the current automatic host has no selectable projects", () => {
+    expect(
+      resolveNewWorkspaceAutomaticServerId({
+        allServerIds: ["empty", "with-project"],
+        routeServerId: null,
+        lastActiveProject: null,
+        projects: [projectFor("with-project")],
+        hostConnectionStatusByServerId: statuses({
+          empty: "connecting",
+          "with-project": "connecting",
+        }),
+        workspaceMultiplicityByServerId: multiplicity(),
+        currentServerId: "empty",
+        nextServerId: "with-project",
+      }),
+    ).toBe("with-project");
+  });
+
+  it("does not switch from an online host to an offline cached project", () => {
+    expect(
+      resolveNewWorkspaceAutomaticServerId({
+        allServerIds: ["online-empty", "offline-project"],
+        routeServerId: null,
+        lastActiveProject: null,
+        projects: [projectFor("offline-project")],
+        hostConnectionStatusByServerId: statuses({
+          "online-empty": "online",
+          "offline-project": "offline",
+        }),
+        workspaceMultiplicityByServerId: multiplicity(),
+        currentServerId: "online-empty",
+        nextServerId: "offline-project",
+      }),
+    ).toBe("online-empty");
+  });
+});

packages/app/src/screens/new-workspace-initial-context.ts

@@ -0,0 +1,254 @@
+import {
+  canCreateWorkspaceForHostProject,
+  type HostProjectListItem,
+} from "@/projects/host-projects";
+import type { HostRuntimeConnectionStatus } from "@/runtime/host-runtime";
+
+export interface NewWorkspaceInitialServerInput {
+  allServerIds: readonly string[];
+  routeServerId: string | null | undefined;
+  lastActiveProject: HostProjectListItem | null;
+  projects: readonly HostProjectListItem[];
+  hostConnectionStatusByServerId: ReadonlyMap<string, HostRuntimeConnectionStatus>;
+  workspaceMultiplicityByServerId: ReadonlyMap<string, boolean>;
+}
+
+function knownServerId(serverIds: ReadonlySet<string>, serverId: string | null | undefined) {
+  const normalized = serverId?.trim() ?? "";
+  return normalized && serverIds.has(normalized) ? normalized : null;
+}
+
+function supportsAllProjects(
+  workspaceMultiplicityByServerId: ReadonlyMap<string, boolean>,
+  serverId: string,
+) {
+  return workspaceMultiplicityByServerId.get(serverId) === true;
+}
+
+function getProjectForServer(input: {
+  candidate: HostProjectListItem;
+  projects: readonly HostProjectListItem[];
+  serverId: string;
+}) {
+  return (
+    input.projects.find(
+      (project) =>
+        project.projectKey === input.candidate.projectKey &&
+        project.hosts.some((host) => host.serverId === input.serverId),
+    ) ?? input.candidate
+  );
+}
+
+function canUseProjectForServer(input: {
+  project: HostProjectListItem;
+  projects: readonly HostProjectListItem[];
+  serverId: string;
+  workspaceMultiplicityByServerId: ReadonlyMap<string, boolean>;
+}) {
+  const project = getProjectForServer({
+    candidate: input.project,
+    projects: input.projects,
+    serverId: input.serverId,
+  });
+  return canCreateWorkspaceForHostProject({
+    project,
+    serverId: input.serverId,
+    allowAllProjects: supportsAllProjects(input.workspaceMultiplicityByServerId, input.serverId),
+  });
+}
+
+function findLastActiveProjectServerId(input: {
+  serverIds: ReadonlySet<string>;
+  lastActiveProject: HostProjectListItem | null;
+  projects: readonly HostProjectListItem[];
+  workspaceMultiplicityByServerId: ReadonlyMap<string, boolean>;
+}) {
+  if (!input.lastActiveProject) {
+    return null;
+  }
+
+  for (const host of input.lastActiveProject.hosts) {
+    if (!input.serverIds.has(host.serverId)) {
+      continue;
+    }
+    if (
+      canUseProjectForServer({
+        project: input.lastActiveProject,
+        projects: input.projects,
+        serverId: host.serverId,
+        workspaceMultiplicityByServerId: input.workspaceMultiplicityByServerId,
+      })
+    ) {
+      return host.serverId;
+    }
+  }
+
+  return null;
+}
+
+function hasSelectableProject(input: {
+  projects: readonly HostProjectListItem[];
+  serverId: string;
+  workspaceMultiplicityByServerId: ReadonlyMap<string, boolean>;
+}) {
+  return input.projects.some((project) =>
+    canUseProjectForServer({
+      project,
+      projects: input.projects,
+      serverId: input.serverId,
+      workspaceMultiplicityByServerId: input.workspaceMultiplicityByServerId,
+    }),
+  );
+}
+
+function isOnline(
+  statuses: ReadonlyMap<string, HostRuntimeConnectionStatus>,
+  serverId: string,
+): boolean {
+  return statuses.get(serverId) === "online";
+}
+
+function isKnownUnreachable(
+  statuses: ReadonlyMap<string, HostRuntimeConnectionStatus>,
+  serverId: string,
+): boolean {
+  const status = statuses.get(serverId);
+  return status === "offline" || status === "error";
+}
+
+export function resolveNewWorkspaceInitialServerId(input: NewWorkspaceInitialServerInput): string {
+  const serverIds = new Set(input.allServerIds);
+  const routeServerId = knownServerId(serverIds, input.routeServerId);
+  if (routeServerId) {
+    return routeServerId;
+  }
+
+  const onlineServerIds = input.allServerIds.filter((serverId) =>
+    isOnline(input.hostConnectionStatusByServerId, serverId),
+  );
+  const onlineServerIdsWithProjects = onlineServerIds.filter((serverId) =>
+    hasSelectableProject({
+      projects: input.projects,
+      serverId,
+      workspaceMultiplicityByServerId: input.workspaceMultiplicityByServerId,
+    }),
+  );
+  const serverIdsWithProjects = input.allServerIds.filter((serverId) =>
+    hasSelectableProject({
+      projects: input.projects,
+      serverId,
+      workspaceMultiplicityByServerId: input.workspaceMultiplicityByServerId,
+    }),
+  );
+
+  const lastActiveProjectServerId = findLastActiveProjectServerId({
+    serverIds,
+    lastActiveProject: input.lastActiveProject,
+    projects: input.projects,
+    workspaceMultiplicityByServerId: input.workspaceMultiplicityByServerId,
+  });
+  if (
+    lastActiveProjectServerId &&
+    isOnline(input.hostConnectionStatusByServerId, lastActiveProjectServerId)
+  ) {
+    return lastActiveProjectServerId;
+  }
+
+  if (onlineServerIdsWithProjects.length > 0) {
+    return onlineServerIdsWithProjects[0] ?? "";
+  }
+  if (onlineServerIds.length === 1) {
+    return onlineServerIds[0] ?? "";
+  }
+
+  if (onlineServerIds.length > 0) {
+    return onlineServerIds[0] ?? "";
+  }
+
+  const reachableServerIdsWithProjects = serverIdsWithProjects.filter(
+    (serverId) => !isKnownUnreachable(input.hostConnectionStatusByServerId, serverId),
+  );
+  if (reachableServerIdsWithProjects.length === 1) {
+    return reachableServerIdsWithProjects[0] ?? "";
+  }
+
+  if (lastActiveProjectServerId) {
+    return lastActiveProjectServerId;
+  }
+
+  if (serverIdsWithProjects.length === 1) {
+    return serverIdsWithProjects[0] ?? "";
+  }
+
+  return input.allServerIds[0] ?? "";
+}
+
+export function resolveNewWorkspaceAutomaticServerId(
+  input: NewWorkspaceInitialServerInput & {
+    currentServerId: string | null | undefined;
+    nextServerId: string | null | undefined;
+  },
+): string {
+  const serverIds = new Set(input.allServerIds);
+  const currentServerId = knownServerId(serverIds, input.currentServerId);
+  const nextServerId = knownServerId(serverIds, input.nextServerId) ?? input.allServerIds[0] ?? "";
+  if (!currentServerId || currentServerId === nextServerId) {
+    return nextServerId;
+  }
+
+  if (
+    isOnline(input.hostConnectionStatusByServerId, nextServerId) &&
+    !isOnline(input.hostConnectionStatusByServerId, currentServerId)
+  ) {
+    return nextServerId;
+  }
+
+  const routeServerId = knownServerId(serverIds, input.routeServerId);
+  if (routeServerId === nextServerId) {
+    return nextServerId;
+  }
+
+  const lastActiveProjectServerId = findLastActiveProjectServerId({
+    serverIds,
+    lastActiveProject: input.lastActiveProject,
+    projects: input.projects,
+    workspaceMultiplicityByServerId: input.workspaceMultiplicityByServerId,
+  });
+  const hasOnlineServer = input.allServerIds.some((serverId) =>
+    isOnline(input.hostConnectionStatusByServerId, serverId),
+  );
+  if (
+    lastActiveProjectServerId === nextServerId &&
+    (isOnline(input.hostConnectionStatusByServerId, nextServerId) || !hasOnlineServer)
+  ) {
+    return nextServerId;
+  }
+
+  const currentHasProject = hasSelectableProject({
+    projects: input.projects,
+    serverId: currentServerId,
+    workspaceMultiplicityByServerId: input.workspaceMultiplicityByServerId,
+  });
+  const nextHasProject = hasSelectableProject({
+    projects: input.projects,
+    serverId: nextServerId,
+    workspaceMultiplicityByServerId: input.workspaceMultiplicityByServerId,
+  });
+  if (
+    isKnownUnreachable(input.hostConnectionStatusByServerId, currentServerId) &&
+    nextHasProject &&
+    !isKnownUnreachable(input.hostConnectionStatusByServerId, nextServerId)
+  ) {
+    return nextServerId;
+  }
+  if (
+    !currentHasProject &&
+    nextHasProject &&
+    (!isOnline(input.hostConnectionStatusByServerId, currentServerId) ||
+      isOnline(input.hostConnectionStatusByServerId, nextServerId))
+  ) {
+    return nextServerId;
+  }
+
+  return currentServerId;
+}

packages/app/src/screens/new-workspace-screen.tsx

@@ -28,7 +28,14 @@ import { HEADER_INNER_HEIGHT, MAX_CONTENT_WIDTH, useIsCompactFormFactor } from "
 import { useToast } from "@/contexts/toast-context";
 import { useAgentInputDraft } from "@/composer/draft/input-draft";
 import { useGithubSearchQuery } from "@/git/use-github-search-query";
-import { useHostRuntimeClient, useHostRuntimeIsConnected, useHosts } from "@/runtime/host-runtime";
+import {
+  useHostRuntimeClient,
+  useHostRuntimeConnectionStatuses,
+  useHostRuntimeIsConnected,
+  useHosts,
+  type HostRuntimeConnectionStatus,
+} from "@/runtime/host-runtime";
+import { useHostFeature, useHostFeatureMap } from "@/runtime/host-features";
 import type { HostProfile } from "@/types/host-connection";
 import { navigateToWorkspace } from "@/stores/navigation-active-workspace-store";
 import { useLastWorkspaceSelection } from "@/stores/navigation-active-workspace-store";
@@ -57,7 +64,6 @@ import {
   resolveSelectedHostProject,
   useHostProjects,
   type HostProjectListItem,
-  type HostProjectRouteContext,
 } from "@/projects/host-projects";
 import { useProjectIconDataByProjectKey } from "@/projects/project-icons";
 import type { ComposerAttachment, UserComposerAttachment } from "@/attachments/types";
@@ -78,6 +84,10 @@ import {
   type PickerItem,
 } from "./new-workspace-picker-item";
 import { findCheckoutHintPrAttachment, syncPickerPrAttachment } from "./new-workspace-picker-state";
+import {
+  resolveNewWorkspaceAutomaticServerId,
+  resolveNewWorkspaceInitialServerId,
+} from "./new-workspace-initial-context";
 
 function resolveCheckoutRequest(
   selectedItem: PickerItem | null,
@@ -1160,29 +1170,95 @@ function submitWorkspaceDraft(input: SubmitDraftInput): void {
   useDraftStore.getState().clearDraftInput({ draftKey, lifecycle: "sent" });
 }
 
-function useNewWorkspaceHostSelector(initialServerId: string) {
-  const allHosts = useHosts();
-  const allServerIds = useMemo(() => allHosts.map((h) => h.serverId), [allHosts]);
-  const lastWorkspaceSelection = useLastWorkspaceSelection();
-  const normalizedInitialServerId = initialServerId.trim();
-  const routeInitialServerId = allServerIds.includes(normalizedInitialServerId)
-    ? normalizedInitialServerId
-    : null;
-  const fallbackServerId =
-    lastWorkspaceSelection && allServerIds.includes(lastWorkspaceSelection.serverId)
-      ? lastWorkspaceSelection.serverId
-      : (allServerIds[0] ?? "");
-  const [manualServerId, setManualServerId] = useState<string | null>(null);
+function useNewWorkspaceHostSelector(input: {
+  initialServerId: string;
+  allServerIds: string[];
+  projects: HostProjectListItem[];
+  lastActiveProject: HostProjectListItem | null;
+  hostConnectionStatusByServerId: ReadonlyMap<string, HostRuntimeConnectionStatus>;
+  workspaceMultiplicityByServerId: ReadonlyMap<string, boolean>;
+}) {
+  const routeServerId = input.initialServerId.trim();
+  const defaultServerId = useMemo(
+    () =>
+      resolveNewWorkspaceInitialServerId({
+        allServerIds: input.allServerIds,
+        routeServerId: input.initialServerId,
+        lastActiveProject: input.lastActiveProject,
+        projects: input.projects,
+        hostConnectionStatusByServerId: input.hostConnectionStatusByServerId,
+        workspaceMultiplicityByServerId: input.workspaceMultiplicityByServerId,
+      }),
+    [
+      input.allServerIds,
+      input.hostConnectionStatusByServerId,
+      input.initialServerId,
+      input.lastActiveProject,
+      input.projects,
+      input.workspaceMultiplicityByServerId,
+    ],
+  );
+  const [automaticSelection, setAutomaticSelection] = useState(() => ({
+    routeServerId,
+    serverId: defaultServerId,
+  }));
+  const [manualSelection, setManualSelection] = useState<{
+    routeServerId: string;
+    serverId: string;
+  } | null>(null);
   const [hostPickerOpen, setHostPickerOpen] = useState(false);
-  const selectedServerId =
-    manualServerId && allServerIds.includes(manualServerId)
-      ? manualServerId
-      : (routeInitialServerId ?? fallbackServerId);
 
-  const handleSelectHost = useCallback((id: string) => {
-    setManualServerId(id);
-    setHostPickerOpen(false);
-  }, []);
+  useEffect(() => {
+    setAutomaticSelection((current) => {
+      const nextServerId =
+        current.routeServerId === routeServerId
+          ? resolveNewWorkspaceAutomaticServerId({
+              allServerIds: input.allServerIds,
+              routeServerId: input.initialServerId,
+              lastActiveProject: input.lastActiveProject,
+              projects: input.projects,
+              hostConnectionStatusByServerId: input.hostConnectionStatusByServerId,
+              workspaceMultiplicityByServerId: input.workspaceMultiplicityByServerId,
+              currentServerId: current.serverId,
+              nextServerId: defaultServerId,
+            })
+          : defaultServerId;
+
+      if (current.routeServerId === routeServerId && current.serverId === nextServerId) {
+        return current;
+      }
+
+      return { routeServerId, serverId: nextServerId };
+    });
+  }, [
+    defaultServerId,
+    input.allServerIds,
+    input.hostConnectionStatusByServerId,
+    input.initialServerId,
+    input.lastActiveProject,
+    input.projects,
+    input.workspaceMultiplicityByServerId,
+    routeServerId,
+  ]);
+
+  const automaticServerId =
+    automaticSelection.routeServerId === routeServerId &&
+    input.allServerIds.includes(automaticSelection.serverId)
+      ? automaticSelection.serverId
+      : defaultServerId;
+  const selectedServerId =
+    manualSelection?.routeServerId === routeServerId &&
+    input.allServerIds.includes(manualSelection.serverId)
+      ? manualSelection.serverId
+      : automaticServerId;
+
+  const handleSelectHost = useCallback(
+    (id: string) => {
+      setManualSelection({ routeServerId, serverId: id });
+      setHostPickerOpen(false);
+    },
+    [routeServerId],
+  );
 
   const handleHostPickerOpenChange = useCallback((open: boolean) => {
     setHostPickerOpen(open);
@@ -1193,8 +1269,6 @@ function useNewWorkspaceHostSelector(initialServerId: string) {
   }, []);
 
   return {
-    allHosts,
-    allServerIds,
     selectedServerId,
     hostPickerOpen,
     handleSelectHost,
@@ -1203,9 +1277,93 @@ function useNewWorkspaceHostSelector(initialServerId: string) {
   };
 }
 
-interface NewWorkspaceProjectPickerInput extends HostProjectRouteContext {
+interface NewWorkspaceInitialContextState {
+  allHosts: HostProfile[];
   selectedServerId: string;
-  allServerIds: string[];
+  hostPickerOpen: boolean;
+  handleSelectHost: (id: string) => void;
+  handleHostPickerOpenChange: (open: boolean) => void;
+  openHostPicker: () => void;
+  projects: HostProjectListItem[];
+  routeProject: HostProjectListItem | null;
+  lastActiveProject: HostProjectListItem | null;
+  routeDisplayName: string;
+}
+
+function useNewWorkspaceInitialContext({
+  serverId,
+  sourceDirectory: sourceDirectoryProp,
+  projectId,
+  displayName: displayNameProp,
+}: NewWorkspaceScreenProps): NewWorkspaceInitialContextState {
+  const allHosts = useHosts();
+  const allServerIds = useMemo(() => allHosts.map((h) => h.serverId), [allHosts]);
+  const projects = useHostProjects(allServerIds);
+  const routeDisplayName = displayNameProp?.trim() ?? "";
+  const routeProject = useMemo(
+    () =>
+      hostProjectFromRoute({
+        serverId,
+        projectId,
+        displayName: routeDisplayName,
+        sourceDirectory: sourceDirectoryProp,
+      }),
+    [projectId, routeDisplayName, serverId, sourceDirectoryProp],
+  );
+  const lastWorkspaceSelection = useLastWorkspaceSelection();
+  const lastWorkspaceServerId = useMemo(
+    () =>
+      lastWorkspaceSelection && allServerIds.includes(lastWorkspaceSelection.serverId)
+        ? lastWorkspaceSelection.serverId
+        : null,
+    [allServerIds, lastWorkspaceSelection],
+  );
+  const lastWorkspaceId = lastWorkspaceServerId ? lastWorkspaceSelection!.workspaceId : null;
+  const lastWorkspace = useWorkspace(lastWorkspaceServerId, lastWorkspaceId);
+  const lastActiveProject = useMemo(
+    () =>
+      lastWorkspaceServerId
+        ? hostProjectFromWorkspace({ serverId: lastWorkspaceServerId, workspace: lastWorkspace })
+        : null,
+    [lastWorkspace, lastWorkspaceServerId],
+  );
+  const hostConnectionStatusByServerId = useHostRuntimeConnectionStatuses(allServerIds);
+  const workspaceMultiplicityByServerId = useHostFeatureMap(allServerIds, "workspaceMultiplicity");
+  const {
+    selectedServerId,
+    hostPickerOpen,
+    handleSelectHost,
+    handleHostPickerOpenChange,
+    openHostPicker,
+  } = useNewWorkspaceHostSelector({
+    initialServerId: serverId,
+    allServerIds,
+    projects,
+    lastActiveProject,
+    hostConnectionStatusByServerId,
+    workspaceMultiplicityByServerId,
+  });
+
+  return {
+    allHosts,
+    selectedServerId,
+    hostPickerOpen,
+    handleSelectHost,
+    handleHostPickerOpenChange,
+    openHostPicker,
+    projects,
+    routeProject,
+    lastActiveProject,
+    routeDisplayName,
+  };
+}
+
+interface NewWorkspaceProjectPickerInput {
+  selectedServerId: string;
+  projects: HostProjectListItem[];
+  routeProject: HostProjectListItem | null;
+  lastActiveProject: HostProjectListItem | null;
+  displayName?: string;
   allowAllProjects: boolean;
 }
 
@@ -1222,39 +1380,15 @@ interface NewWorkspaceProjectPickerState {
 }
 
 function useNewWorkspaceProjectPicker({
-  serverId,
   selectedServerId,
-  allServerIds,
-  sourceDirectory,
-  projectId,
+  projects,
+  routeProject,
+  lastActiveProject,
   displayName: displayNameProp,
   allowAllProjects,
 }: NewWorkspaceProjectPickerInput): NewWorkspaceProjectPickerState {
   const [manualProjectKey, setManualProjectKey] = useState<string | null>(null);
   const displayName = displayNameProp?.trim() ?? "";
-  const projects = useHostProjects(allServerIds);
-  const lastWorkspaceSelection = useLastWorkspaceSelection();
-  const lastWorkspaceServerId = useMemo(
-    () =>
-      lastWorkspaceSelection && allServerIds.includes(lastWorkspaceSelection.serverId)
-        ? lastWorkspaceSelection.serverId
-        : null,
-    [allServerIds, lastWorkspaceSelection],
-  );
-  const lastWorkspaceId = lastWorkspaceServerId ? lastWorkspaceSelection!.workspaceId : null;
-  const lastWorkspace = useWorkspace(lastWorkspaceServerId, lastWorkspaceId);
-
-  const routeProject = useMemo(
-    () => hostProjectFromRoute({ serverId, projectId, displayName, sourceDirectory }),
-    [displayName, projectId, serverId, sourceDirectory],
-  );
-  const lastActiveProject = useMemo(
-    () =>
-      lastWorkspaceServerId
-        ? hostProjectFromWorkspace({ serverId: lastWorkspaceServerId, workspace: lastWorkspace })
-        : null,
-    [lastWorkspace, lastWorkspaceServerId],
-  );
   const selectableProjects = useMemo(
     () =>
       filterWorkspaceProjectsForHost({ projects, serverId: selectedServerId, allowAllProjects }),
@@ -1556,18 +1690,23 @@ export function NewWorkspaceScreen({
   const mergeWorkspaces = useSessionStore((state) => state.mergeWorkspaces);
   const {
     allHosts,
-    allServerIds,
     selectedServerId,
     hostPickerOpen,
     handleSelectHost,
     handleHostPickerOpenChange,
     openHostPicker,
-  } = useNewWorkspaceHostSelector(serverId);
+    projects,
+    routeProject,
+    lastActiveProject,
+    routeDisplayName,
+  } = useNewWorkspaceInitialContext({
+    serverId,
+    sourceDirectory: sourceDirectoryProp,
+    projectId,
+    displayName: displayNameProp,
+  });
   // COMPAT(workspaceMultiplicity): added in v0.1.97, drop the gate when floor >= v0.1.97
-  const supportsWorkspaceMultiplicity = useSessionStore(
-    (state) =>
-      state.sessions[selectedServerId]?.serverInfo?.features?.workspaceMultiplicity === true,
-  );
+  const supportsWorkspaceMultiplicity = useHostFeature(selectedServerId, "workspaceMultiplicity");
   const [errorMessage, setErrorMessage] = useState<string | null>(null);
   const [createdWorkspace, setCreatedWorkspace] = useState<ReturnType<
     typeof normalizeWorkspaceDescriptor
@@ -1596,7 +1735,6 @@ export function NewWorkspaceScreen({
   const client = useHostRuntimeClient(selectedServerId);
   const isConnected = useHostRuntimeIsConnected(selectedServerId);
   const {
-    projects,
     selectedProject,
     selectedSourceDirectory,
     projectPickerOptions,
@@ -1605,12 +1743,11 @@ export function NewWorkspaceScreen({
     projectTriggerLabel,
     handleSelectProjectOption: selectProjectOption,
   } = useNewWorkspaceProjectPicker({
-    serverId,
     selectedServerId,
-    allServerIds,
-    sourceDirectory: sourceDirectoryProp,
-    projectId,
-    displayName: displayNameProp,
+    projects,
+    routeProject,
+    lastActiveProject,
+    displayName: routeDisplayName,
     allowAllProjects: supportsWorkspaceMultiplicity,
   });