fix(vscode): keep queued messages visible

.changeset/show-queued-prompts.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Keep queued prompts visible while the server confirms and loads their message parts.

packages/kilo-vscode/tests/unit/prompt-send-contract.test.ts

@@ -89,6 +89,16 @@ describe("sendCommand dismisses pending tool requests", () => {
   })
 })
 
+describe("confirmed queued prompts retain optimistic parts", () => {
+  const source = readFile(SESSION_FILE)
+  const body = extractFunctionBody(source, "handleMessageCreated")
+
+  it("does not clear optimistic parts before canonical part events arrive", () => {
+    expect(body).toContain("Keep placeholder parts until their canonical part.updated events arrive")
+    expect(body).not.toContain("delete p[message.id]")
+  })
+})
+
 describe("static command completion contract", () => {
   const source = readFile(SESSION_FILE)
 

packages/kilo-vscode/tests/unit/session-parts.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, it } from "bun:test"
-import { mergeParts, sameParts } from "../../webview-ui/src/context/session-parts"
+import { mergeOptimisticPart, mergeParts, sameParts } from "../../webview-ui/src/context/session-parts"
 import type { Part } from "../../webview-ui/src/types/messages"
 
 function text(id: string, value: string, time: { start?: number; end?: number } = {}): Part {
@@ -10,6 +10,10 @@ function tool(id: string): Part {
   return { id, messageID: "m1", type: "tool", tool: "bash", state: { status: "pending", input: {} } }
 }
 
+function file(id: string): Part {
+  return { id, messageID: "m1", type: "file", mime: "text/plain", url: "data:,file" }
+}
+
 function value(parts: Part[], id: string) {
   const part = parts.find((item) => item.id === id)
   if (!part || part.type !== "text") return
@@ -100,6 +104,24 @@ describe("mergeParts", () => {
   })
 })
 
+describe("mergeOptimisticPart", () => {
+  it("replaces the optimistic user part when its canonical event arrives", () => {
+    const current = [text("client", "queued prompt")]
+    const result = mergeOptimisticPart(current, new Set(["client"]), text("server", "queued prompt"))
+
+    expect(result.parts).toEqual([text("server", "queued prompt")])
+    expect(result.replaced).toBe("client")
+  })
+
+  it("keeps unmatched optimistic parts while canonical attachments arrive", () => {
+    const current = [text("client-text", "queued prompt"), file("client-file")]
+    const result = mergeOptimisticPart(current, new Set(["client-text", "client-file"]), file("server-file"))
+
+    expect(result.parts.map((part) => part.id)).toEqual(["client-text", "server-file"])
+    expect(result.replaced).toBe("client-file")
+  })
+})
+
 describe("sameParts", () => {
   it("accepts equal hydrated and snapshot parts", () => {
     expect(sameParts([text("p1", "done", { end: 2 })], [text("p1", "done", { end: 2 })])).toBe(true)

packages/kilo-vscode/webview-ui/src/context/session-parts.ts

@@ -25,6 +25,15 @@ export function sameParts(local: Part[] = [], snapshot: Part[] = []): boolean {
   return true
 }
 
+export function mergeOptimisticPart(current: Part[], ids: ReadonlySet<string>, part: Part) {
+  const index = current.findIndex((item) => ids.has(item.id) && item.type === part.type)
+  if (index < 0) return { parts: [...current, part] }
+  const old = current[index]!
+  const next = current.slice()
+  next[index] = part
+  return { parts: next, replaced: old.id }
+}
+
 /**
  * Reconcile snapshots may be older than in-flight streaming deltas. Preserve
  * only appended streamed tail parts and open prefix extensions while still

packages/kilo-vscode/webview-ui/src/context/session.tsx

@@ -77,7 +77,7 @@ import { getAgentModel } from "./session-model-store"
 import { resolveMessagePrefs } from "./session-preferences"
 import { errorIDs, preserveSessionErrors, withoutResolvedSessionErrors } from "./session-errors"
 import { PartStash } from "./part-stash"
-import { mergeParts } from "./session-parts"
+import { mergeOptimisticPart, mergeParts } from "./session-parts"
 import { mergeMessages, sameReconcileShape } from "./session-merge"
 import { state as todoState } from "./todo-revert"
 import { sessionVariantKeys, transferVariants, variantKey } from "./session-variant-store"
@@ -457,6 +457,9 @@ export const SessionProvider: ParentComponent = (props) => {
   // Tracks optimistic messageIDs that haven't been confirmed by the server yet.
   // Prevents handleMessagesLoaded from wiping them when it replaces the array.
   const pendingOptimistic = new Map<string, Set<string>>()
+  // Keeps optimistic parts visible between message.updated and their canonical
+  // message.part.updated events.
+  const optimisticParts = new Map<string, Set<string>>()
   // Sessions can be created/imported while an older list request is still in flight.
   // Keep them until a later list payload confirms them or deletion arrives.
   const freshSessions = new Set<string>()
@@ -1478,13 +1481,14 @@ export const SessionProvider: ParentComponent = (props) => {
       for (let i = 0; i < messages.length; i++) {
         const msg = messages[i]!
         const parts = msg.parts ?? []
-        if (mode === "reconcile" && store.parts[msg.id]) {
+        if (mode === "reconcile" && store.parts[msg.id] && !optimisticParts.has(msg.id)) {
           const merged = mergeParts(store.parts[msg.id], parts, input.since ?? Number.POSITIVE_INFINITY)
           setStore("parts", msg.id, reconcile(merged, { key: "id" }))
           stash.remove(msg.id)
           continue
         }
         if (parts.length > 0) {
+          optimisticParts.delete(msg.id)
           loadedParts[msg.id] = parts
           if (i >= cutoff) {
             setStore("parts", msg.id, parts)
@@ -1537,22 +1541,12 @@ export const SessionProvider: ParentComponent = (props) => {
   function handleMessageCreated(message: Message) {
     if (message.role === "assistant") clearSessionDraftDiscarded(message.sessionID)
     // Message confirmed by server — no longer optimistic.
-    // Clear placeholder parts so they don't duplicate alongside real parts
-    // arriving via individual part.updated events (the server's message.updated
-    // SSE event does NOT include parts).
+    // Keep placeholder parts until their canonical part.updated events arrive.
+    // The message.updated SSE event does not include parts, so clearing them
+    // here makes a queued prompt render only its status during that gap.
     const pending = pendingOptimistic.get(message.sessionID)
-    const wasOptimistic = pending?.has(message.id)
     pending?.delete(message.id)
 
-    if (wasOptimistic) {
-      setStore(
-        "parts",
-        produce((p) => {
-          delete p[message.id]
-        }),
-      )
-    }
-
     const exists = (store.messages[message.sessionID] ?? []).some((msg) => msg.id === message.id)
     setStore("messages", message.sessionID, (msgs = []) => {
       if (message.sessionErrorID && msgs.some((msg) => msg.sessionErrorID === message.sessionErrorID)) return msgs
@@ -1573,6 +1567,7 @@ export const SessionProvider: ParentComponent = (props) => {
     recoverPrefs(message.sessionID, [message])
 
     if (message.parts && message.parts.length > 0) {
+      optimisticParts.delete(message.id)
       stash.remove(message.id)
       setStore("parts", message.id, message.parts)
     }
@@ -1608,6 +1603,19 @@ export const SessionProvider: ParentComponent = (props) => {
       setStore("parts", effectiveMessageID, stashed)
     }
 
+    const current = store.parts[effectiveMessageID] ?? []
+    const index = current.findIndex((item) => item.id === part.id)
+    const pending = optimisticParts.get(effectiveMessageID)
+    if (index < 0 && pending) {
+      const merged = mergeOptimisticPart(current, pending, part)
+      setStore("parts", effectiveMessageID, merged.parts)
+      if (merged.replaced) {
+        pending.delete(merged.replaced)
+        if (pending.size === 0) optimisticParts.delete(effectiveMessageID)
+      }
+      return
+    }
+
     setStore(
       "parts",
       produce((parts) => {
@@ -1689,6 +1697,7 @@ export const SessionProvider: ParentComponent = (props) => {
           delete map[sessionID]
         }),
       )
+      for (const msg of store.messages[sessionID] ?? []) optimisticParts.delete(msg.id)
       // Session is idle - any remaining pending optimistic IDs are either
       // already confirmed (messageCreated removed them) or orphaned (queued
       // callbacks were dropped on abort). Clean up the tracking set; the
@@ -1817,6 +1826,7 @@ export const SessionProvider: ParentComponent = (props) => {
     if (!message.messageID && sid) aborts.clear(sid)
     if (sid && message.messageID) {
       pendingOptimistic.get(sid)?.delete(message.messageID)
+      optimisticParts.delete(message.messageID)
       stash.remove(message.messageID)
       batch(() => {
         setStore("messages", sid, (msgs = []) => msgs.filter((m) => m.id !== message.messageID))
@@ -1978,6 +1988,7 @@ export const SessionProvider: ParentComponent = (props) => {
       // Collect message IDs so we can clean up their parts (store + stash)
       const msgs = store.messages[sessionID] ?? []
       const msgIds = msgs.map((m) => m.id)
+      for (const id of msgIds) optimisticParts.delete(id)
       for (const id of msgIds) stash.remove(id)
       clearHiddenErrors(msgIds)
 
@@ -2041,6 +2052,7 @@ export const SessionProvider: ParentComponent = (props) => {
 
   // Splices the message from the store and deletes its parts.
   function handleMessageRemoved(sessionID: string, messageID: string) {
+    optimisticParts.delete(messageID)
     setStore("messages", sessionID, (msgs = []) => msgs.filter((m) => m.id !== messageID))
     dropMessageTools(sessionID, messageID)
     clearHiddenErrors([messageID])
@@ -2220,6 +2232,7 @@ export const SessionProvider: ParentComponent = (props) => {
 
     setStore("messages", sid, (msgs = []) => [...msgs, temp])
     setStore("parts", messageID, parts)
+    if (parts.length > 0) optimisticParts.set(messageID, new Set(parts.map((part) => part.id)))
     patchPage(sid, { lastMutation: "append" })
     queueMicrotask(() => window.dispatchEvent(new CustomEvent("resumeAutoScroll")))
   }