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

1 files changed, 69 insertions, 2 deletions

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" },
+    ])
+  })
 })