fix(app): better optimistic prompt submit (#17337)

4 files changed, 254 insertions, 9 deletions

diff --git a/packages/app/src/components/prompt-input/submit.test.ts b/packages/app/src/components/prompt-input/submit.test.ts
index 4109417d2..9f7fac69d 100644
--- a/packages/app/src/components/prompt-input/submit.test.ts
+++ b/packages/app/src/components/prompt-input/submit.test.ts
@@ -7,12 +7,16 @@ const createdClients: string[] = []
 const createdSessions: string[] = []
 const enabledAutoAccept: Array<{ sessionID: string; directory: string }> = []
 const optimistic: Array<{
+  directory?: string
+  sessionID?: string
   message: {
     agent: string
     model: { providerID: string; modelID: string }
     variant?: string
   }
 }> = []
+const optimisticSeeded: boolean[] = []
+const storedSessions: Record<string, Array<{ id: string; title?: string }>> = {}
 const sentShell: string[] = []
 const syncedDirectories: string[] = []
 
@@ -28,7 +32,12 @@ const clientFor = (directory: string) => {
     session: {
       create: async () => {
         createdSessions.push(directory)
-        return { data: { id: `session-${createdSessions.length}` } }
+        return {
+          data: {
+            id: `session-${createdSessions.length}`,
+            title: `New session ${createdSessions.length}`,
+          },
+        }
       },
       shell: async () => {
         sentShell.push(directory)
@@ -129,9 +138,16 @@ beforeAll(async () => {
       session: {
         optimistic: {
           add: (value: {
+            directory?: string
+            sessionID?: string
             message: { agent: string; model: { providerID: string; modelID: string }; variant?: string }
           }) => {
             optimistic.push(value)
+            optimisticSeeded.push(
+              !!value.directory &&
+                !!value.sessionID &&
+                !!storedSessions[value.directory]?.find((item) => item.id === value.sessionID)?.title,
+            )
           },
           remove: () => undefined,
         },
@@ -144,7 +160,21 @@ beforeAll(async () => {
     useGlobalSync: () => ({
       child: (directory: string) => {
         syncedDirectories.push(directory)
-        return [{}, () => undefined]
+        storedSessions[directory] ??= []
+        return [
+          { session: storedSessions[directory] },
+          (...args: unknown[]) => {
+            if (args[0] !== "session") return
+            const next = args[1]
+            if (typeof next === "function") {
+              storedSessions[directory] = next(storedSessions[directory]) as Array<{ id: string; title?: string }>
+              return
+            }
+            if (Array.isArray(next)) {
+              storedSessions[directory] = next as Array<{ id: string; title?: string }>
+            }
+          },
+        ]
       },
     }),
   }))
@@ -170,11 +200,13 @@ beforeEach(() => {
   createdSessions.length = 0
   enabledAutoAccept.length = 0
   optimistic.length = 0
+  optimisticSeeded.length = 0
   params = {}
   sentShell.length = 0
   syncedDirectories.length = 0
   selected = "/repo/worktree-a"
   variant = undefined
+  for (const key of Object.keys(storedSessions)) delete storedSessions[key]
 })
 
 describe("prompt submit worktree selection", () => {
@@ -207,7 +239,7 @@ describe("prompt submit worktree selection", () => {
     expect(createdClients).toEqual(["/repo/worktree-a", "/repo/worktree-b"])
     expect(createdSessions).toEqual(["/repo/worktree-a", "/repo/worktree-b"])
     expect(sentShell).toEqual(["/repo/worktree-a", "/repo/worktree-b"])
-    expect(syncedDirectories).toEqual(["/repo/worktree-a", "/repo/worktree-b"])
+    expect(syncedDirectories).toEqual(["/repo/worktree-a", "/repo/worktree-a", "/repo/worktree-b", "/repo/worktree-b"])
   })
 
   test("applies auto-accept to newly created sessions", async () => {
@@ -271,4 +303,32 @@ describe("prompt submit worktree selection", () => {
       },
     })
   })
+
+  test("seeds new sessions before optimistic prompts are added", async () => {
+    const submit = createPromptSubmit({
+      info: () => undefined,
+      imageAttachments: () => [],
+      commentCount: () => 0,
+      autoAccept: () => false,
+      mode: () => "normal",
+      working: () => false,
+      editor: () => undefined,
+      queueScroll: () => undefined,
+      promptLength: (value) => value.reduce((sum, part) => sum + ("content" in part ? part.content.length : 0), 0),
+      addToHistory: () => undefined,
+      resetHistoryNavigation: () => undefined,
+      setMode: () => undefined,
+      setPopover: () => undefined,
+      newSessionWorktree: () => selected,
+      onNewSessionWorktreeReset: () => undefined,
+      onSubmit: () => undefined,
+    })
+
+    const event = { preventDefault: () => undefined } as unknown as Event
+
+    await submit.handleSubmit(event)
+
+    expect(storedSessions["/repo/worktree-a"]).toEqual([{ id: "session-1", title: "New session 1" }])
+    expect(optimisticSeeded).toEqual([true])
+  })
 })
diff --git a/packages/app/src/components/prompt-input/submit.ts b/packages/app/src/components/prompt-input/submit.ts
index eb3e0c82f..e8d765cd9 100644
--- a/packages/app/src/components/prompt-input/submit.ts
+++ b/packages/app/src/components/prompt-input/submit.ts
@@ -1,6 +1,7 @@
-import type { Message } from "@opencode-ai/sdk/v2/client"
+import type { Message, Session } from "@opencode-ai/sdk/v2/client"
 import { showToast } from "@opencode-ai/ui/toast"
 import { base64Encode } from "@opencode-ai/util/encode"
+import { Binary } from "@opencode-ai/util/binary"
 import { useNavigate, useParams } from "@solidjs/router"
 import type { Accessor } from "solid-js"
 import type { FileSelection } from "@/context/file"
@@ -266,6 +267,20 @@ export function createPromptSubmit(input: PromptSubmitInput) {
     }
   }
 
+  const seed = (dir: string, info: Session) => {
+    const [, setStore] = globalSync.child(dir)
+    setStore("session", (list: Session[]) => {
+      const result = Binary.search(list, info.id, (item) => item.id)
+      const next = [...list]
+      if (result.found) {
+        next[result.index] = info
+        return next
+      }
+      next.splice(result.index, 0, info)
+      return next
+    })
+  }
+
   const handleSubmit = async (event: Event) => {
     event.preventDefault()
 
@@ -341,7 +356,7 @@ export function createPromptSubmit(input: PromptSubmitInput) {
 
     let session = input.info()
     if (!session && isNewSession) {
-      session = await client.session
+      const created = await client.session
         .create()
         .then((x) => x.data ?? undefined)
         .catch((err) => {
@@ -351,7 +366,9 @@ export function createPromptSubmit(input: PromptSubmitInput) {
           })
           return undefined
         })
-      if (session) {
+      if (created) {
+        seed(sessionDirectory, created)
+        session = created
         if (shouldAutoAccept) permission.enableAutoAccept(session.id, sessionDirectory)
         layout.handoff.setTabs(base64Encode(sessionDirectory), session.id)
         navigate(`/${base64Encode(sessionDirectory)}/session/${session.id}`)
diff --git a/packages/app/src/context/sync-optimistic.test.ts b/packages/app/src/context/sync-optimistic.test.ts
index 7deeddd6e..94324f8a0 100644
--- a/packages/app/src/context/sync-optimistic.test.ts
+++ b/packages/app/src/context/sync-optimistic.test.ts
@@ -1,6 +1,8 @@
 import { describe, expect, test } from "bun:test"
 import type { Message, Part } from "@opencode-ai/sdk/v2/client"
-import { applyOptimisticAdd, applyOptimisticRemove } from "./sync"
+import { applyOptimisticAdd, applyOptimisticRemove, mergeOptimisticPage } from "./sync"
+
+type Text = Extract<Part, { type: "text" }>
 
 const userMessage = (id: string, sessionID: string): Message => ({
   id,
@@ -11,7 +13,7 @@ const userMessage = (id: string, sessionID: string): Message => ({
   model: { providerID: "openai", modelID: "gpt" },
 })
 
-const textPart = (id: string, sessionID: string, messageID: string): Part => ({
+const textPart = (id: string, sessionID: string, messageID: string): Text => ({
   id,
   sessionID,
   messageID,
@@ -53,4 +55,69 @@ describe("sync optimistic reducers", () => {
     expect(draft.part.msg_1).toBeUndefined()
     expect(draft.part.msg_2).toHaveLength(1)
   })
+
+  test("mergeOptimisticPage keeps pending messages in fetched timelines", () => {
+    const sessionID = "ses_1"
+    const page = mergeOptimisticPage(
+      {
+        session: [userMessage("msg_1", sessionID)],
+        part: [{ id: "msg_1", part: [textPart("prt_1", sessionID, "msg_1")] }],
+        complete: true,
+      },
+      [{ message: userMessage("msg_2", sessionID), parts: [textPart("prt_2", sessionID, "msg_2")] }],
+    )
+
+    expect(page.session.map((x) => x.id)).toEqual(["msg_1", "msg_2"])
+    expect(page.part.find((x) => x.id === "msg_2")?.part.map((x) => x.id)).toEqual(["prt_2"])
+    expect(page.confirmed).toEqual([])
+    expect(page.complete).toBe(true)
+  })
+
+  test("mergeOptimisticPage keeps missing optimistic parts until the server has them", () => {
+    const sessionID = "ses_1"
+    const page = mergeOptimisticPage(
+      {
+        session: [userMessage("msg_2", sessionID)],
+        part: [{ id: "msg_2", part: [textPart("prt_2", sessionID, "msg_2")] }],
+        complete: true,
+      },
+      [
+        {
+          message: userMessage("msg_2", sessionID),
+          parts: [textPart("prt_1", sessionID, "msg_2"), textPart("prt_2", sessionID, "msg_2")],
+        },
+      ],
+    )
+
+    expect(page.part.find((x) => x.id === "msg_2")?.part.map((x) => x.id)).toEqual(["prt_1", "prt_2"])
+    expect(page.confirmed).toEqual([])
+  })
+
+  test("mergeOptimisticPage confirms echoed messages once all parts arrive", () => {
+    const sessionID = "ses_1"
+    const page = mergeOptimisticPage(
+      {
+        session: [userMessage("msg_2", sessionID)],
+        part: [
+          {
+            id: "msg_2",
+            part: [{ ...textPart("prt_1", sessionID, "msg_2"), text: "server" }, textPart("prt_2", sessionID, "msg_2")],
+          },
+        ],
+        complete: true,
+      },
+      [
+        {
+          message: userMessage("msg_2", sessionID),
+          parts: [textPart("prt_1", sessionID, "msg_2"), textPart("prt_2", sessionID, "msg_2")],
+        },
+      ],
+    )
+
+    expect(page.confirmed).toEqual(["msg_2"])
+    expect(page.part.find((x) => x.id === "msg_2")?.part).toMatchObject([
+      { id: "prt_1", type: "text", text: "server" },
+      { id: "prt_2", type: "text", text: "prt_2" },
+    ])
+  })
 })
diff --git a/packages/app/src/context/sync.tsx b/packages/app/src/context/sync.tsx
index 9dc6623a7..0f2008723 100644
--- a/packages/app/src/context/sync.tsx
+++ b/packages/app/src/context/sync.tsx
@@ -54,6 +54,67 @@ type OptimisticRemoveInput = {
   messageID: string
 }
 
+type OptimisticItem = {
+  message: Message
+  parts: Part[]
+}
+
+type MessagePage = {
+  session: Message[]
+  part: { id: string; part: Part[] }[]
+  cursor?: string
+  complete: boolean
+}
+
+const hasParts = (parts: Part[] | undefined, want: Part[]) => {
+  if (!parts) return want.length === 0
+  return want.every((part) => Binary.search(parts, part.id, (item) => item.id).found)
+}
+
+const mergeParts = (parts: Part[] | undefined, want: Part[]) => {
+  if (!parts) return sortParts(want)
+  const next = [...parts]
+  let changed = false
+  for (const part of want) {
+    const result = Binary.search(next, part.id, (item) => item.id)
+    if (result.found) continue
+    next.splice(result.index, 0, part)
+    changed = true
+  }
+  if (!changed) return parts
+  return next
+}
+
+export function mergeOptimisticPage(page: MessagePage, items: OptimisticItem[]) {
+  if (items.length === 0) return { ...page, confirmed: [] as string[] }
+
+  const session = [...page.session]
+  const part = new Map(page.part.map((item) => [item.id, sortParts(item.part)]))
+  const confirmed: string[] = []
+
+  for (const item of items) {
+    const result = Binary.search(session, item.message.id, (message) => message.id)
+    const found = result.found
+    if (!found) session.splice(result.index, 0, item.message)
+
+    const current = part.get(item.message.id)
+    if (found && hasParts(current, item.parts)) {
+      confirmed.push(item.message.id)
+      continue
+    }
+
+    part.set(item.message.id, mergeParts(current, item.parts))
+  }
+
+  return {
+    cursor: page.cursor,
+    complete: page.complete,
+    session,
+    part: [...part.entries()].sort((a, b) => cmp(a[0], b[0])).map(([id, part]) => ({ id, part })),
+    confirmed,
+  }
+}
+
 export function applyOptimisticAdd(draft: OptimisticStore, input: OptimisticAddInput) {
   const messages = draft.message[input.sessionID]
   if (messages) {
@@ -121,6 +182,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
     const inflight = new Map<string, Promise<void>>()
     const inflightDiff = new Map<string, Promise<void>>()
     const inflightTodo = new Map<string, Promise<void>>()
+    const optimistic = new Map<string, Map<string, OptimisticItem>>()
     const maxDirs = 30
     const seen = new Map<string, Set<string>>()
     const [meta, setMeta] = createStore({
@@ -137,6 +199,33 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
       return undefined
     }
 
+    const setOptimistic = (directory: string, sessionID: string, item: OptimisticItem) => {
+      const key = keyFor(directory, sessionID)
+      const list = optimistic.get(key)
+      if (list) {
+        list.set(item.message.id, { message: item.message, parts: sortParts(item.parts) })
+        return
+      }
+      optimistic.set(key, new Map([[item.message.id, { message: item.message, parts: sortParts(item.parts) }]]))
+    }
+
+    const clearOptimistic = (directory: string, sessionID: string, messageID?: string) => {
+      const key = keyFor(directory, sessionID)
+      if (!messageID) {
+        optimistic.delete(key)
+        return
+      }
+
+      const list = optimistic.get(key)
+      if (!list) return
+      list.delete(messageID)
+      if (list.size === 0) optimistic.delete(key)
+    }
+
+    const getOptimistic = (directory: string, sessionID: string) => [
+      ...(optimistic.get(keyFor(directory, sessionID))?.values() ?? []),
+    ]
+
     const seenFor = (directory: string) => {
       const existing = seen.get(directory)
       if (existing) {
@@ -159,6 +248,9 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
 
     const clearMeta = (directory: string, sessionIDs: string[]) => {
       if (sessionIDs.length === 0) return
+      for (const sessionID of sessionIDs) {
+        clearOptimistic(directory, sessionID)
+      }
       setMeta(
         produce((draft) => {
           for (const sessionID of sessionIDs) {
@@ -232,8 +324,12 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
 
       setMeta("loading", key, true)
       await fetchMessages(input)
-        .then((next) => {
+        .then((page) => {
           if (!tracked(input.directory, input.sessionID)) return
+          const next = mergeOptimisticPage(page, getOptimistic(input.directory, input.sessionID))
+          for (const messageID of next.confirmed) {
+            clearOptimistic(input.directory, input.sessionID, messageID)
+          }
           const [store] = globalSync.child(input.directory, { bootstrap: false })
           const cached = input.mode === "prepend" ? (store.message[input.sessionID] ?? []) : []
           const message = input.mode === "prepend" ? merge(cached, next.session) : next.session
@@ -290,11 +386,15 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
         get: getSession,
         optimistic: {
           add(input: { directory?: string; sessionID: string; message: Message; parts: Part[] }) {
+            const directory = input.directory ?? sdk.directory
             const [, setStore] = target(input.directory)
+            setOptimistic(directory, input.sessionID, { message: input.message, parts: input.parts })
             setOptimisticAdd(setStore as (...args: unknown[]) => void, input)
           },
           remove(input: { directory?: string; sessionID: string; messageID: string }) {
+            const directory = input.directory ?? sdk.directory
             const [, setStore] = target(input.directory)
+            clearOptimistic(directory, input.sessionID, input.messageID)
             setOptimisticRemove(setStore as (...args: unknown[]) => void, input)
           },
         },
@@ -316,6 +416,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
             variant: input.variant,
           }
           const [, setStore] = target()
+          setOptimistic(sdk.directory, input.sessionID, { message, parts: input.parts })
           setOptimisticAdd(setStore as (...args: unknown[]) => void, {
             sessionID: input.sessionID,
             message,