diff --git a/.changeset/show-queued-prompts.md b/.changeset/show-queued-prompts.md new file mode 100644 index 0000000000..7431d4ee06 --- /dev/null +++ b/.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. diff --git a/packages/kilo-vscode/tests/unit/prompt-send-contract.test.ts b/packages/kilo-vscode/tests/unit/prompt-send-contract.test.ts index cc44d15c8d..67ddbe6a77 100644 --- a/packages/kilo-vscode/tests/unit/prompt-send-contract.test.ts +++ b/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) diff --git a/packages/kilo-vscode/tests/unit/session-parts.test.ts b/packages/kilo-vscode/tests/unit/session-parts.test.ts index 9acaa16390..16cebaf7ac 100644 --- a/packages/kilo-vscode/tests/unit/session-parts.test.ts +++ b/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) diff --git a/packages/kilo-vscode/webview-ui/src/context/session-parts.ts b/packages/kilo-vscode/webview-ui/src/context/session-parts.ts index 2733d73775..883a2b7148 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session-parts.ts +++ b/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, 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 diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index ed5b5d07ff..d6f709fc16 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/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>() + // Keeps optimistic parts visible between message.updated and their canonical + // message.part.updated events. + const optimisticParts = new Map>() // 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() @@ -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"))) }