feat(app): follow-up behavior (#17233)

packages/app/src/components/prompt-input/submit.ts

@@ -9,7 +9,7 @@ import { useLanguage } from "@/context/language"
 import { useLayout } from "@/context/layout"
 import { useLocal } from "@/context/local"
 import { usePermission } from "@/context/permission"
-import { type ImageAttachmentPart, type Prompt, usePrompt } from "@/context/prompt"
+import { type ContextItem, type ImageAttachmentPart, type Prompt, usePrompt } from "@/context/prompt"
 import { useSDK } from "@/context/sdk"
 import { useSync } from "@/context/sync"
 import { Identifier } from "@/utils/id"
@@ -25,6 +25,145 @@ type PendingPrompt = {
 
 const pending = new Map<string, PendingPrompt>()
 
+export type FollowupDraft = {
+  sessionID: string
+  sessionDirectory: string
+  prompt: Prompt
+  context: (ContextItem & { key: string })[]
+  agent: string
+  model: { providerID: string; modelID: string }
+  variant?: string
+}
+
+type FollowupSendInput = {
+  client: ReturnType<typeof useSDK>["client"]
+  globalSync: ReturnType<typeof useGlobalSync>
+  sync: ReturnType<typeof useSync>
+  draft: FollowupDraft
+  messageID?: string
+  optimisticBusy?: boolean
+  before?: () => Promise<boolean> | boolean
+}
+
+const draftText = (prompt: Prompt) => prompt.map((part) => ("content" in part ? part.content : "")).join("")
+
+const draftImages = (prompt: Prompt) => prompt.filter((part): part is ImageAttachmentPart => part.type === "image")
+
+export async function sendFollowupDraft(input: FollowupSendInput) {
+  const text = draftText(input.draft.prompt)
+  const images = draftImages(input.draft.prompt)
+  const [, setStore] = input.globalSync.child(input.draft.sessionDirectory)
+
+  const setBusy = () => {
+    if (!input.optimisticBusy) return
+    setStore("session_status", input.draft.sessionID, { type: "busy" })
+  }
+
+  const setIdle = () => {
+    if (!input.optimisticBusy) return
+    setStore("session_status", input.draft.sessionID, { type: "idle" })
+  }
+
+  const wait = async () => {
+    const ok = await input.before?.()
+    if (ok === false) return false
+    return true
+  }
+
+  const [head, ...tail] = text.split(" ")
+  const cmd = head?.startsWith("/") ? head.slice(1) : undefined
+  if (cmd && input.sync.data.command.find((item) => item.name === cmd)) {
+    setBusy()
+    try {
+      if (!(await wait())) {
+        setIdle()
+        return false
+      }
+
+      await input.client.session.command({
+        sessionID: input.draft.sessionID,
+        command: cmd,
+        arguments: tail.join(" "),
+        agent: input.draft.agent,
+        model: `${input.draft.model.providerID}/${input.draft.model.modelID}`,
+        variant: input.draft.variant,
+        parts: images.map((attachment) => ({
+          id: Identifier.ascending("part"),
+          type: "file" as const,
+          mime: attachment.mime,
+          url: attachment.dataUrl,
+          filename: attachment.filename,
+        })),
+      })
+      return true
+    } catch (err) {
+      setIdle()
+      throw err
+    }
+  }
+
+  const messageID = input.messageID ?? Identifier.ascending("message")
+  const { requestParts, optimisticParts } = buildRequestParts({
+    prompt: input.draft.prompt,
+    context: input.draft.context,
+    images,
+    text,
+    sessionID: input.draft.sessionID,
+    messageID,
+    sessionDirectory: input.draft.sessionDirectory,
+  })
+
+  const message: Message = {
+    id: messageID,
+    sessionID: input.draft.sessionID,
+    role: "user",
+    time: { created: Date.now() },
+    agent: input.draft.agent,
+    model: input.draft.model,
+    variant: input.draft.variant,
+  }
+
+  const add = () =>
+    input.sync.session.optimistic.add({
+      directory: input.draft.sessionDirectory,
+      sessionID: input.draft.sessionID,
+      message,
+      parts: optimisticParts,
+    })
+
+  const remove = () =>
+    input.sync.session.optimistic.remove({
+      directory: input.draft.sessionDirectory,
+      sessionID: input.draft.sessionID,
+      messageID,
+    })
+
+  setBusy()
+  add()
+
+  try {
+    if (!(await wait())) {
+      setIdle()
+      remove()
+      return false
+    }
+
+    await input.client.session.promptAsync({
+      sessionID: input.draft.sessionID,
+      agent: input.draft.agent,
+      model: input.draft.model,
+      messageID,
+      parts: requestParts,
+      variant: input.draft.variant,
+    })
+    return true
+  } catch (err) {
+    setIdle()
+    remove()
+    throw err
+  }
+}
+
 type PromptSubmitInput = {
   info: Accessor<{ id: string } | undefined>
   imageAttachments: Accessor<ImageAttachmentPart[]>
@@ -41,6 +180,9 @@ type PromptSubmitInput = {
   setPopover: (popover: "at" | "slash" | null) => void
   newSessionWorktree?: Accessor<string | undefined>
   onNewSessionWorktreeReset?: () => void
+  shouldQueue?: Accessor<boolean>
+  onQueue?: (draft: FollowupDraft) => void
+  onAbort?: () => void
   onSubmit?: () => void
 }
 
@@ -82,6 +224,8 @@ export function createPromptSubmit(input: PromptSubmitInput) {
     const [, setStore] = globalSync.child(sdk.directory)
     setStore("todo", sessionID, [])
 
+    input.onAbort?.()
+
     const queued = pending.get(sessionID)
     if (queued) {
       queued.abort.abort()
@@ -116,6 +260,12 @@ export function createPromptSubmit(input: PromptSubmitInput) {
     }
   }
 
+  const clearContext = () => {
+    for (const item of prompt.context.items()) {
+      prompt.context.remove(item.key)
+    }
+  }
+
   const handleSubmit = async (event: Event) => {
     event.preventDefault()
 
@@ -215,14 +365,22 @@ export function createPromptSubmit(input: PromptSubmitInput) {
       return
     }
 
-    input.onSubmit?.()
-
     const model = {
       modelID: currentModel.id,
       providerID: currentModel.provider.id,
     }
     const agent = currentAgent.name
     const variant = local.model.variant.current()
+    const context = prompt.context.items().slice()
+    const draft: FollowupDraft = {
+      sessionID: session.id,
+      sessionDirectory,
+      prompt: currentPrompt,
+      context,
+      agent,
+      model,
+      variant,
+    }
 
     const clearInput = () => {
       prompt.reset()
@@ -243,6 +401,15 @@ export function createPromptSubmit(input: PromptSubmitInput) {
       })
     }
 
+    if (!isNewSession && mode === "normal" && input.shouldQueue?.()) {
+      input.onQueue?.(draft)
+      clearContext()
+      clearInput()
+      return
+    }
+
+    input.onSubmit?.()
+
     if (mode === "shell") {
       clearInput()
       client.session
@@ -295,48 +462,19 @@ export function createPromptSubmit(input: PromptSubmitInput) {
       }
     }
 
-    const context = prompt.context.items().slice()
     const commentItems = context.filter((item) => item.type === "file" && !!item.comment?.trim())
-
     const messageID = Identifier.ascending("message")
-    const { requestParts, optimisticParts } = buildRequestParts({
-      prompt: currentPrompt,
-      context,
-      images,
-      text,
-      sessionID: session.id,
-      messageID,
-      sessionDirectory,
-    })
 
-    const optimisticMessage: Message = {
-      id: messageID,
-      sessionID: session.id,
-      role: "user",
-      time: { created: Date.now() },
-      agent,
-      model,
-      variant,
-    }
-
-    const addOptimisticMessage = () =>
-      sync.session.optimistic.add({
-        directory: sessionDirectory,
-        sessionID: session.id,
-        message: optimisticMessage,
-        parts: optimisticParts,
-      })
-
-    const removeOptimisticMessage = () =>
+    const removeOptimisticMessage = () => {
       sync.session.optimistic.remove({
         directory: sessionDirectory,
         sessionID: session.id,
         messageID,
       })
+    }
 
     removeCommentItems(commentItems)
     clearInput()
-    addOptimisticMessage()
 
     const waitForWorktree = async () => {
       const worktree = WorktreeState.get(sessionDirectory)
@@ -393,20 +531,15 @@ export function createPromptSubmit(input: PromptSubmitInput) {
       return true
     }
 
-    const send = async () => {
-      const ok = await waitForWorktree()
-      if (!ok) return
-      await client.session.promptAsync({
-        sessionID: session.id,
-        agent,
-        model,
-        messageID,
-        parts: requestParts,
-        variant,
-      })
-    }
-
-    void send().catch((err) => {
+    void sendFollowupDraft({
+      client,
+      sync,
+      globalSync,
+      draft,
+      messageID,
+      optimisticBusy: sessionDirectory === projectDirectory,
+      before: waitForWorktree,
+    }).catch((err) => {
       pending.delete(session.id)
       if (sessionDirectory === projectDirectory) {
         sync.set("session_status", session.id, { type: "idle" })