feat(tui): optimistic session creation on first prompt (#43687)

packages/tui/src/component/prompt/index.tsx

@@ -19,6 +19,8 @@ import { useClipboard } from "../../context/clipboard"
 import { Spinner } from "../spinner"
 import { useClient } from "../../context/client"
 import { useRoute } from "../../context/route"
+import { usePromptRef } from "../../context/prompt"
+import { useSessionTabs } from "../../context/session-tabs"
 import { useEvent } from "../../context/event"
 import { editorSelectionKey, useEditorContext, type EditorSelection } from "../../context/editor"
 import { normalizePromptContent, openEditor } from "../../editor"
@@ -201,6 +203,8 @@ export function Prompt(props: PromptProps) {
   const client = useClient()
   const editor = useEditorContext()
   const route = useRoute()
+  const promptRef = usePromptRef()
+  const sessionTabs = useSessionTabs()
   const data = useData()
   const directoryRecents = useDirectoryRecents()
   const keymapCommands = Keymap.useCommands()
@@ -688,16 +692,20 @@ export function Prompt(props: PromptProps) {
       input.gotoBufferEnd()
     },
     reset() {
-      input.clear()
-      input.extmarks.clear()
-      setStore("prompt", emptyPrompt())
-      setStore("extmarkToPart", new Map())
+      resetComposer()
     },
     submit() {
       void submit()
     },
   }
 
+  function resetComposer() {
+    input.extmarks.clear()
+    setStore("prompt", emptyPrompt())
+    setStore("extmarkToPart", new Map())
+    input.clear()
+  }
+
   // Captured once: the session route is keyed by sessionID, so this Prompt
   // instance belongs to exactly one tab. Reading props.sessionID lazily would
   // observe the *next* route during onCleanup and stash under the wrong tab.
@@ -873,10 +881,7 @@ export function Prompt(props: PromptProps) {
         run: () => {
           if (!store.prompt.text) return
           stash.push({ prompt: store.prompt })
-          input.extmarks.clear()
-          input.clear()
-          setStore("prompt", emptyPrompt())
-          setStore("extmarkToPart", new Map())
+          resetComposer()
           dialog.clear()
         },
       },
@@ -1168,102 +1173,139 @@ export function Prompt(props: PromptProps) {
       return false
     }
 
+    // Snapshot the composer and clear it synchronously, before the first await.
+    // Everything below reads the snapshot: text typed while a request is in
+    // flight lands in the already-empty composer and survives, and prompt
+    // history records exactly what was submitted instead of the live store
+    // (which may have absorbed mid-flight typing). Failure paths restore the
+    // snapshot unless the user has started typing something new.
+    const currentMode = store.mode
+    const entry = { ...store.prompt, mode: currentMode }
+    history.append(entry)
+    resetComposer()
+    props.onSubmit?.()
+    const restoreEntry = () => {
+      if (disposed || input.isDestroyed || input.plainText !== "") return
+      input.setText(entry.text)
+      setStore("prompt", entry)
+      setStore("mode", entry.mode ?? "normal")
+      restoreExtmarksFromPrompt(entry)
+      input.cursorOffset = entry.text.length
+    }
+
     const variant = selection.variant
     let sessionID = props.sessionID
     let session = sessionID ? data.session.get(sessionID) : undefined
     let finishMoveProgress = false
+    // New-session sends wait for creation and environment setup.
+    let newSession: { gate: Promise<unknown>; recover: (error: unknown) => void } | undefined
     if (sessionID == null) {
       const directory = await move.getDirectory()
-      if (move.pending() && !directory) return false
+      if (move.pending() && !directory) {
+        restoreEntry()
+        return false
+      }
       finishMoveProgress = Boolean(move.progress())
       // The location context is where the next session is created: seeded by the home
       // route (launch cwd, inherited session location, or picked project) and updated
       // by /cd before a session exists.
       const location = currentLocation.ref ?? data.location.default()
 
-      const created = await client.api.session
-        .create({
-          location: directory ? { directory } : location,
-          agent: agent.id,
-          model: {
-            providerID: selection.providerID,
-            id: selection.modelID,
-            variant,
-          },
-        })
-        .catch(() => undefined)
-
-      if (!created) {
-        if (finishMoveProgress) move.finishSubmit()
-        toast.show({
-          message: "Creating a session failed. Open console for more details.",
-          variant: "error",
-        })
-
-        return true
-      }
-
+      // Optimistic create: the data layer mints the ID client-side and admits
+      // a local session record synchronously, so the navigation below happens
+      // immediately — enter feels sent even while the create round-trip is in
+      // flight. Sends against the new session gate on the request.
+      const created = data.session.create({
+        location: directory ? { directory } : location,
+        agent: agent.id,
+        model: {
+          providerID: selection.providerID,
+          id: selection.modelID,
+          variant,
+        },
+      })
       sessionID = created.id
-      session = created
-      if (created.location.workspaceID === undefined && terminalEnvironment.variables !== undefined) {
-        const error = await client.api.session
-          .environment({ sessionID, variables: terminalEnvironment.variables })
-          .then(
-            () => undefined,
-            (error) => error,
-          )
-        if (error) {
-          if (finishMoveProgress) move.finishSubmit()
-          toast.show({ title: "Failed to set session environment", message: errorMessage(error), variant: "error" })
-          return true
-        }
+      session = data.session.get(created.id)
+      newSession = {
+        gate: created.request.then(async (info) => {
+          if (info.location.workspaceID === undefined && terminalEnvironment.variables !== undefined) {
+            await client.api.session.environment({ sessionID: created.id, variables: terminalEnvironment.variables })
+          }
+        }),
+        recover: (error) => {
+          toast.show({
+            title: data.session.get(created.id) ? "Failed to set up session" : "Creating a session failed",
+            message: errorMessage(error),
+            variant: "error",
+          })
+          const active =
+            route.data.type === "session" && route.data.sessionID === created.id ? promptRef.current : undefined
+          const current = active?.current
+          const draft = current?.text
+            ? { prompt: { ...unwrap(current) }, cursor: current.text.length }
+            : (takeDraft(created.id) ?? { prompt: entry, cursor: entry.text.length })
+          saveDraft(undefined, draft)
+          active?.reset()
+          if (sessionTabs.enabled()) {
+            sessionTabs.close(created.id)
+          } else if (route.data.type === "session" && route.data.sessionID === created.id) {
+            route.navigate({ type: "home" })
+          }
+        },
       }
     }
 
-    // Capture mode before it gets reset
-    const currentMode = store.mode
-    if (store.mode === "shell") {
+    const target = sessionID
+    const dispatch = (send: () => Promise<unknown>) => {
+      const setup = newSession
+      if (setup) void setup.gate.then(send).catch(setup.recover)
+      else void send()
+    }
+    if (currentMode === "shell") {
       move.startSubmit()
-      void client.api.session.shell({
-        sessionID,
-        command: inputText,
-      })
+      dispatch(() => client.api.session.shell({ sessionID: target, command: inputText }))
       setStore("mode", "normal")
     } else if (slashHead && isCommand) {
       move.startSubmit()
       const model = { providerID: selection.providerID, id: selection.modelID, variant }
-      const cancelCommit = local.model.trackSessionCommit(sessionID, model)
+      const cancelCommit = local.model.trackSessionCommit(target, model)
 
-      void client.api.session
-        .command({
-          sessionID,
+      const send = () =>
+        client.api.session.command({
+          sessionID: target,
           command: slashHead.name,
           arguments: slashHead.arguments,
           agent: agent.id,
           model,
-          files: store.prompt.files,
-          agents: store.prompt.agents,
-          skills: store.prompt.skills?.length ? store.prompt.skills : undefined,
+          files: entry.files,
+          agents: entry.agents,
+          skills: entry.skills?.length ? entry.skills : undefined,
           delivery,
         })
-        .catch((error) => {
-          cancelCommit()
-          toast.show({ title: "Failed to run command", message: errorMessage(error), variant: "error" })
-        })
+      const setup = newSession
+      void (setup ? setup.gate.then(send) : send()).catch((error) => {
+        cancelCommit()
+        if (setup) return setup.recover(error)
+        toast.show({ title: "Failed to run command", message: errorMessage(error), variant: "error" })
+        restoreEntry()
+      })
     } else if (isSkill) {
       move.startSubmit()
-      void client.api.session.skill({
-        sessionID,
-        skill: slashHead.name,
-      })
+      dispatch(() => client.api.session.skill({ sessionID: target, skill: slashHead.name }))
     } else {
       move.startSubmit()
-      if (!session) {
-        await data.session.sync(sessionID)
-        session = data.session.get(sessionID)
-      }
-      if (session?.agent !== agent.id) {
-        await client.api.session.switchAgent({ sessionID, agent: agent.id })
+      try {
+        if (!session) {
+          await data.session.sync(target)
+          session = data.session.get(target)
+        }
+        if (session?.agent !== agent.id) {
+          await client.api.session.switchAgent({ sessionID: target, agent: agent.id })
+        }
+      } catch (error) {
+        toast.show({ title: "Failed to prepare session", message: errorMessage(error), variant: "error" })
+        restoreEntry()
+        return true
       }
       if (
         session?.model?.providerID !== selection.providerID ||
@@ -1271,83 +1313,94 @@ export function Prompt(props: PromptProps) {
         (session.model.variant ?? "default") !== (variant ?? "default")
       ) {
         const model = { providerID: selection.providerID, id: selection.modelID, variant }
-        const cancelCommit = local.model.trackSessionCommit(sessionID, model)
-        await client.api.session.switchModel({ sessionID, model }).catch((error) => {
+        const cancelCommit = local.model.trackSessionCommit(target, model)
+        const switchError = await client.api.session.switchModel({ sessionID: target, model }).then(
+          () => undefined,
+          (error) => error,
+        )
+        if (switchError) {
           cancelCommit()
-          throw error
-        })
+          toast.show({ title: "Failed to switch model", message: errorMessage(switchError), variant: "error" })
+          restoreEntry()
+          return true
+        }
       }
       if (session?.revert) {
-        const error = await client.api.session.revert.commit({ sessionID }).then(
+        const error = await client.api.session.revert.commit({ sessionID: target }).then(
           () => undefined,
           (error) => error,
         )
         if (error) {
           toast.show({ title: "Failed to commit revert", message: errorMessage(error), variant: "error" })
+          restoreEntry()
           return false
         }
       }
       if (pendingEditorSelection) {
         // Keep editor context hidden while admitting it before the corresponding user prompt.
-        const error = await client.api.session
-          .synthetic({
-            sessionID,
+        const send = () =>
+          client.api.session.synthetic({
+            sessionID: target,
             text: formatEditorContext(pendingEditorSelection),
             resume: false,
           })
-          .then(
+        if (newSession) {
+          // Fold into the setup gate so the context still admits before the
+          // user prompt once the session exists.
+          newSession.gate = newSession.gate.then(send)
+        } else {
+          const error = await send().then(
             () => undefined,
             (error) => error,
           )
-        if (error) {
-          toast.show({ title: "Failed to send editor context", message: errorMessage(error), variant: "error" })
-          return false
+          if (error) {
+            toast.show({ title: "Failed to send editor context", message: errorMessage(error), variant: "error" })
+            restoreEntry()
+            return false
+          }
         }
       }
       // The data layer admits optimistically: the prompt renders immediately
       // and rolls back if the server rejects it, so submission does not wait
       // on the network. On rejection the row is already rolled back; restore
       // the composer unless the user has started typing something new.
-      const entry = { ...store.prompt, mode: currentMode }
       data.session
         .prompt({
-          sessionID,
+          sessionID: target,
           text: inputText,
-          files: store.prompt.files,
-          agents: store.prompt.agents,
-          skills: store.prompt.skills?.length ? store.prompt.skills : undefined,
+          files: entry.files,
+          agents: entry.agents,
+          skills: entry.skills?.length ? entry.skills : undefined,
           delivery,
+          gate: newSession?.gate,
         })
         .catch((error) => {
+          if (newSession) return newSession.recover(error)
           toast.show({ title: "Failed to send prompt", message: errorMessage(error), variant: "error" })
-          if (disposed || input.isDestroyed || input.plainText !== "") return
-          input.setText(entry.text)
-          setStore("prompt", entry)
-          setStore("mode", entry.mode ?? "normal")
-          restoreExtmarksFromPrompt(entry)
-          input.cursorOffset = entry.text.length
+          restoreEntry()
         })
       if (pendingEditorSelection) editor.markSelectionSent()
     }
-    history.append({
-      ...store.prompt,
-      mode: currentMode,
-    })
-    input.extmarks.clear()
-    setStore("prompt", emptyPrompt())
-    setStore("extmarkToPart", new Map())
-    props.onSubmit?.()
 
     // Optimistic admission puts the message in the store synchronously, so
     // the session view renders it on arrival.
     if (!props.sessionID) {
       if (pendingEditorSelection) editor.preserveSelectionFromNewSession()
+      // Text typed while session creation was in flight lives in this (home)
+      // prompt, which unmounts on navigation and would stash it under the
+      // home key. Re-stash it under the new session so that composer restores
+      // it, and clear it here so onCleanup does not also stash it for home.
+      if (!disposed && !input.isDestroyed && store.prompt.text) {
+        // Copy before clearing: unwrap returns the live store target, and the
+        // resetComposer store write merges into that same object.
+        saveDraft(sessionID, { prompt: { ...unwrap(store.prompt) }, cursor: input.cursorOffset })
+        resetComposer()
+      }
       route.navigate({
         type: "session",
         sessionID,
       })
     }
-    input.clear()
     if (finishMoveProgress) move.finishSubmit()
     return true
   }
@@ -1509,10 +1562,7 @@ export function Prompt(props: PromptProps) {
         mode: store.mode,
       })
     }
-    input.clear()
-    input.extmarks.clear()
-    setStore("prompt", emptyPrompt())
-    setStore("extmarkToPart", new Map())
+    resetComposer()
   }
 
   const highlight = createMemo(() => {