fix(tui): move session when recovering missing directory (#51096)

Co-authored-by: vimtor <36263538+vimtor@users.noreply.github.com>

packages/tui/src/routes/session/location-missing.tsx

@@ -4,11 +4,56 @@ import { useTheme } from "../../context/theme"
 import { Locale } from "../../util/locale"
 import { abbreviateHome } from "../../util/path-format"
 import { SessionQuestion } from "./permission"
-import { usePromptMove } from "../../component/prompt/move"
+import { useDialog } from "../../ui/dialog"
+import { useClient } from "../../context/client"
+import { useToast } from "../../ui/toast"
+import { errorMessage } from "../../util/error"
+import { DialogWorkspaces, type WorkspaceSelection } from "../../component/dialog-workspaces"
+import { useData } from "../../context/data"
 
 export function SessionLocationMissing(props: { directory: string; projectID: string; sessionID: string }) {
-  const move = usePromptMove({ projectID: () => props.projectID, sessionID: () => props.sessionID })
-  return <SessionLocationUnavailable directory={props.directory} onMove={move.open} />
+  const dialog = useDialog()
+  const client = useClient()
+  const toast = useToast()
+  const data = useData()
+
+  function open() {
+    dialog.replace(() => (
+      <DialogWorkspaces
+        projectID={props.projectID}
+        location={{ directory: props.directory }}
+        current={{
+          type: "directory",
+          directory: props.directory,
+          subdirectory: !!data.session.get(props.sessionID)?.subpath,
+        }}
+        onSelect={(selection) => void select(selection)}
+      />
+    ))
+  }
+
+  async function select(selection: WorkspaceSelection) {
+    dialog.clear()
+    const directory =
+      selection.type === "directory"
+        ? selection.directory
+        : await client.api.worktree
+            .create({ projectID: props.projectID, name: selection.name })
+            .then((result) => {
+              if (!result.directory) throw new Error("No worktree directory returned")
+              return result.directory
+            })
+            .catch((error) => {
+              toast.show({ title: "Creating workspace failed", message: errorMessage(error), variant: "error" })
+              return undefined
+            })
+    if (!directory) return
+    await client.api.session.move({ sessionID: props.sessionID, directory }).catch((error) => {
+      toast.show({ title: "Failed to move session", message: errorMessage(error), variant: "error" })
+    })
+  }
+
+  return <SessionLocationUnavailable directory={props.directory} onMove={open} />
 }
 
 export function SessionLocationUnavailable(props: { directory: string; onMove: () => void }) {

packages/tui/test/cli/tui/prompt-move.test.tsx

@@ -12,6 +12,7 @@ import { RouteProvider, useRoute } from "../../../src/context/route"
 import { ThemeProvider } from "../../../src/context/theme"
 import { DialogProvider } from "../../../src/ui/dialog"
 import { ToastProvider, useToast } from "../../../src/ui/toast"
+import { SessionLocationMissing } from "../../../src/routes/session/location-missing"
 import { emptyThemeSource } from "../../fixture/fixture"
 import { createApi, createEventStream, createFetch, json } from "../../fixture/tui-client"
 import { TestTuiContexts } from "../../fixture/tui-environment"
@@ -175,6 +176,68 @@ test.each([false, true])("Ctrl+M moves only an existing session (home=%s)", asyn
   }
 })
 
+test("choosing a directory recovers the session when its location is unavailable", async () => {
+  const fixture = await renderMove({ directory: clone, unavailable: "location", showMissingLocation: true })
+  try {
+    await fixture.app.waitForFrame((frame) => frame.includes("Session location unavailable"))
+    fixture.app.mockInput.pressEnter()
+    await fixture.app.waitForFrame((frame) => frame.includes("Worktrees") && frame.includes(linked))
+    await fixture.app.waitFor(() => fixture.app.renderer.currentFocusedEditor instanceof InputRenderable)
+    await fixture.app.mockInput.typeText("linked")
+    await fixture.app.waitForFrame((frame) => frame.includes(linked) && !frame.includes(main))
+    fixture.app.mockInput.pressEnter()
+    await fixture.app.waitFor(() => fixture.moves.length === 1)
+
+    expect(fixture.moves).toEqual([{ directory: linked }])
+    expect(fixture.route.data).toEqual({ type: "session", sessionID: "ses_clone" })
+    expect(fixture.requests).toEqual([])
+  } finally {
+    fixture.app.renderer.destroy()
+  }
+})
+
+test("creating a worktree recovers the session without reading its removed location", async () => {
+  const fixture = await renderMove({ directory: clone, unavailable: "location", showMissingLocation: true })
+  try {
+    await fixture.app.waitForFrame((frame) => frame.includes("Session location unavailable"))
+    fixture.app.mockInput.pressEnter()
+    await fixture.app.waitForFrame((frame) => frame.includes("Worktrees") && frame.includes(linked))
+    fixture.app.mockInput.pressKey("a", { ctrl: true })
+    await fixture.app.waitForFrame((frame) => frame.includes("Name worktree"))
+    await fixture.app.waitFor(() => fixture.app.renderer.currentFocusedEditor instanceof InputRenderable)
+    await fixture.app.mockInput.typeText("fresh")
+    fixture.app.mockInput.pressEnter()
+    await fixture.app.waitFor(() => fixture.moves.length === 1)
+
+    expect(fixture.requests).toEqual([{ payload: { projectID: "proj_test", name: "fresh" }, directory: null }])
+    expect(fixture.moves).toEqual([{ directory: created }])
+    expect(fixture.route.data).toEqual({ type: "session", sessionID: "ses_clone" })
+    expect(fixture.reads.locations).not.toContain(clone)
+  } finally {
+    fixture.app.renderer.destroy()
+  }
+})
+
+test("failed recovery does not navigate away from the session", async () => {
+  const fixture = await renderMove({ directory: clone, unavailable: "location", showMissingLocation: true, moveFails: true })
+  try {
+    await fixture.app.waitForFrame((frame) => frame.includes("Session location unavailable"))
+    fixture.app.mockInput.pressEnter()
+    await fixture.app.waitForFrame((frame) => frame.includes("Worktrees") && frame.includes(linked))
+    await fixture.app.waitFor(() => fixture.app.renderer.currentFocusedEditor instanceof InputRenderable)
+    await fixture.app.mockInput.typeText("linked")
+    await fixture.app.waitForFrame((frame) => frame.includes(linked) && !frame.includes(main))
+    fixture.app.mockInput.pressEnter()
+    await fixture.app.waitFor(() => fixture.toast.currentToast !== null)
+
+    expect(fixture.moves).toEqual([{ directory: linked }])
+    expect(fixture.route.data).toEqual({ type: "session", sessionID: "ses_clone" })
+    expect(fixture.toast.currentToast).toMatchObject({ title: "Failed to move session", variant: "error" })
+  } finally {
+    fixture.app.renderer.destroy()
+  }
+})
+
 test.each([
   { name: "session", unavailable: "session" as const },
   { name: "location", unavailable: "location" as const },
@@ -223,6 +286,8 @@ async function renderMove(input: {
   launch?: string
   launchProjectID?: string
   unavailable?: "session" | "location"
+  showMissingLocation?: boolean
+  moveFails?: boolean
 }) {
   const launch = input.launch ?? (input.home ? input.directory : main)
   const requests: unknown[] = []
@@ -294,6 +359,7 @@ async function renderMove(input: {
     }
     if (url.pathname === "/api/session/ses_clone/move") {
       moves.push(await request.json())
+      if (input.moveFails) return json({ message: "Destination unavailable" }, { status: 503 })
       return new Response(null, { status: 204 })
     }
     return undefined
@@ -313,7 +379,9 @@ async function renderMove(input: {
       projectID: () => (input.home ? data.location.info()?.project.id : "proj_test"),
       sessionID: () => (input.home ? undefined : "ses_clone"),
     })
-    return null
+    return input.showMissingLocation ? (
+      <SessionLocationMissing directory={input.directory} projectID="proj_test" sessionID="ses_clone" />
+    ) : null
   }
 
   const app = await testRender(