fix(app): preserve workspace choices across draft tabs (#45149)

Co-authored-by: Brendonovich <14191578+Brendonovich@users.noreply.github.com>

packages/app/src/new-session/composer-adapter.ts

@@ -133,7 +133,7 @@ export function createNewSessionComposerAdapter(props: {
 
   return {
     adapter,
-    project: createComposerProjectControls({ draftId: props.draftID }),
+    project: createComposerProjectControls({ draftId: props.draftID, worktree: props.worktree }),
     model,
     ready: prompt.ready,
   }

packages/app/src/new-session/project/controller.ts

@@ -4,10 +4,11 @@ import { useGlobal, useServerCtx } from "@/runtime/server/runtime"
 import { useServerSDK } from "@/runtime/server/client"
 import { serverName, ServerConnection, useServers } from "@/runtime/server/registry"
 import { useWorkspaceLocation } from "@/workspaces/location"
+import { workspaceSelectionDestination } from "@/workspaces/paths"
 import { useTabs } from "@/shell/tabs/tabs"
 import type { PromptProjectControls } from "./selector"
 
-export function createComposerProjectControls(props: { draftId: string }) {
+export function createComposerProjectControls(props: { draftId: string; worktree: () => string }) {
   const servers = useServers()
   const serverSDK = useServerSDK()
   const location = useWorkspaceLocation()
@@ -38,7 +39,7 @@ export function createComposerProjectControls(props: { draftId: string }) {
     tabs.updateDraft(props.draftId, {
       server: ServerConnection.key(connection),
       directory: worktree,
-      worktree: undefined,
+      worktree: workspaceSelectionDestination(props.worktree(), location().directory),
       branch: undefined,
     })
   }

packages/app/src/new-session/workspace/controller.ts

@@ -5,6 +5,8 @@ import { useWorkspaceLocation } from "@/workspaces/location"
 import { useServerSDK } from "@/runtime/server/client"
 import { useData } from "@/runtime/server/current"
 import { useSettings } from "@/settings/model"
+import { useTabs } from "@/shell/tabs/tabs"
+import { ServerConnection } from "@/runtime/server/registry"
 import { normalizeProjectInfo } from "@/runtime/server/global-sync/utils"
 import {
   isWorkspaceDirectory,
@@ -12,6 +14,7 @@ import {
   sameDirectory,
   workspaceDefaultSelection,
   workspaceDirectories,
+  workspaceSelectionDestination,
 } from "@/workspaces/paths"
 
 export function resolveNewSessionWorktree(input: {
@@ -57,6 +60,7 @@ export function createNewSessionWorkspaceController(input: {
   const serverSDK = useServerSDK()
   const data = useData()
   const settings = useSettings()
+  const tabs = useTabs()
   const [state, setState] = createStore({ search: "" })
   const searchBranches = debounce((search: string) => setState("search", search.trim()), 100)
   const currentProject = createMemo(() => {
@@ -121,7 +125,8 @@ export function createNewSessionWorkspaceController(input: {
   const remember = (worktree = value()) => {
     const project = currentProject()
     if (!project) return
-    const local = worktree === "main" || sameDirectory(worktree, project.worktree)
+    tabs.initializeDraftWorktrees(ServerConnection.key(serverSDK.server), sdk().directory, fallback())
+    const local = workspaceSelectionDestination(worktree, project.worktree) === "main"
     settings.workspaces.setLastUsed(serverSDK.scope, project.id, local ? "local" : "workspace")
   }
 
@@ -141,6 +146,7 @@ export function createNewSessionWorkspaceController(input: {
       set: (worktree: string) => {
         input.setSelectedBranch(undefined)
         input.setSelectedWorktree(normalizeNewSessionWorktree(worktree, sdk().directory, currentProject()?.worktree))
+        remember(worktree)
       },
       create: (branch: string) => {
         input.setSelectedBranch(branch)
@@ -160,7 +166,10 @@ export function createNewSessionWorkspaceController(input: {
         const loaded = branches.latest
         const list = loaded?.directory === projectRoot() ? loaded.data : []
         return [
-          ...new Set([...list, ...(current && current.toLowerCase().includes(state.search.toLowerCase()) ? [current] : [])]),
+          ...new Set([
+            ...list,
+            ...(current && current.toLowerCase().includes(state.search.toLowerCase()) ? [current] : []),
+          ]),
         ].slice(0, 50)
       },
       searchBranches,

packages/app/src/shell/tabs/tabs.tsx

@@ -245,6 +245,14 @@ export const { use: useTabs, provider: TabsProvider } = createSimpleContext({
           )
         })
       },
+      initializeDraftWorktrees(server: ServerConnection.Key, directory: string, worktree: string) {
+        setStore(
+          (tab) => tab.type === "draft" && tab.server === server && tab.directory === directory && !tab.worktree,
+          produce((tab) => {
+            if (tab.type === "draft") tab.worktree = worktree
+          }),
+        )
+      },
       promoteDraft(draftID: string, session: Omit<SessionTab, "type">) {
         // Keep the replacement and navigation atomic so /new-session never renders
         // after its backing draft tab has been removed from the store.

packages/app/src/workspaces/paths.test.ts

@@ -8,6 +8,7 @@ import {
   mergeWorkspaceSessionInventory,
   sessionsForWorkspace,
   workspaceInventory,
+  workspaceSelectionDestination,
 } from "./paths"
 
 describe("isWorkspaceDirectory", () => {
@@ -41,6 +42,19 @@ describe("isWorkspaceSelection", () => {
   })
 })
 
+describe("workspaceSelectionDestination", () => {
+  test("preserves local intent for the main selection and project root", () => {
+    expect(workspaceSelectionDestination("main", "/repo")).toBe("main")
+    expect(workspaceSelectionDestination("/repo/", "/repo")).toBe("main")
+    expect(workspaceSelectionDestination("c:\\repo\\", "C:\\repo")).toBe("main")
+  })
+
+  test("preserves workspace intent without carrying project-specific paths", () => {
+    expect(workspaceSelectionDestination("create", "/repo")).toBe("create")
+    expect(workspaceSelectionDestination("/workspaces/feature", "/repo")).toBe("create")
+  })
+})
+
 test("groups and filters workspace inventory by project", () => {
   const inventory = workspaceInventory([
     {

packages/app/src/workspaces/paths.ts

@@ -106,6 +106,10 @@ export function isWorkspaceSelection(project: WorkspaceProject | undefined, sele
   return isWorkspaceDirectory(project, selection)
 }
 
+export function workspaceSelectionDestination(selection: string, projectWorktree: string) {
+  return selection === "main" || sameDirectory(selection, projectWorktree) ? "main" : "create"
+}
+
 export function workspaceDefaultSelection(
   setting: WorkspaceDefaultDestination,
   lastUsed: WorkspaceLastUsed | undefined,