diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index 7b286eb0125..5f5d6f3aeaf 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -336,6 +336,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper .parse(message.files) await this.handleSendMessage( message.text, + typeof message.messageID === "string" ? message.messageID : undefined, message.sessionID, message.providerID, message.modelID, @@ -362,8 +363,6 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper break case "clearSession": this.currentSession = null - this.trackedSessionIds.clear() - this.syncedChildSessions.clear() break case "loadMessages": // Don't await: allow parallel loads so rapid session switching @@ -546,6 +545,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper void this.handleImportAndSend( message.cloudSessionId, message.text, + typeof message.messageID === "string" ? message.messageID : undefined, message.providerID, message.modelID, message.agent, @@ -1272,6 +1272,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper private async handleImportAndSend( cloudSessionId: string, text: string, + messageID?: string, providerID?: string, modelID?: string, agent?: string, @@ -1340,10 +1341,15 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper try { const editorContext = await this.gatherEditorContext() - await this.client.session.prompt( + if (messageID) { + this.connectionService.recordMessageSessionId(messageID, session.id) + } + + await this.client.session.promptAsync( { sessionID: session.id, directory: workspaceDir, + messageID, parts, model: providerID && modelID ? { providerID, modelID } : undefined, agent, @@ -1355,9 +1361,12 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper } catch (err) { console.error("[Kilo New] Failed to send message after cloud import:", err) this.postMessage({ - type: "error", - message: err instanceof Error ? err.message : "Failed to send message after import", + type: "sendMessageFailed", + error: err instanceof Error ? err.message : "Failed to send message after import", + text, sessionID: session.id, + messageID, + files, }) } } @@ -1425,9 +1434,12 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper /** * Handle sending a message from the webview. + * Uses promptAsync (fire-and-forget) — the server queues concurrent prompts + * internally, so the client doesn't need to block on busy sessions. */ private async handleSendMessage( text: string, + messageID?: string, sessionID?: string, providerID?: string, modelID?: string, @@ -1437,12 +1449,17 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper ): Promise { if (!this.client) { this.postMessage({ - type: "error", - message: "Not connected to CLI backend", + type: "sendMessageFailed", + error: "Not connected to CLI backend", + text, + sessionID, + messageID, + files, }) return } + let targetSessionID = sessionID try { const workspaceDir = this.getWorkspaceDirectory(sessionID || this.currentSession?.id) @@ -1461,10 +1478,11 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper }) } - const targetSessionID = sessionID || this.currentSession?.id + targetSessionID = sessionID || this.currentSession?.id if (!targetSessionID) { throw new Error("No session available") } + this.trackedSessionIds.add(targetSessionID) // Build parts array with file context and user text const parts: Array = [] @@ -1480,10 +1498,15 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper const editorContext = await this.gatherEditorContext() - await this.client.session.prompt( + if (messageID) { + this.connectionService.recordMessageSessionId(messageID, targetSessionID) + } + + await this.client.session.promptAsync( { sessionID: targetSessionID, directory: workspaceDir, + messageID, parts, model: providerID && modelID ? { providerID, modelID } : undefined, agent, @@ -1495,8 +1518,12 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper } catch (error) { console.error("[Kilo New] KiloProvider: Failed to send message:", error) this.postMessage({ - type: "error", - message: getErrorMessage(error) || "Failed to send message", + type: "sendMessageFailed", + error: getErrorMessage(error) || "Failed to send message", + text, + sessionID: targetSessionID, + messageID, + files, }) } } diff --git a/packages/kilo-vscode/webview-ui/src/components/chat/MessageList.tsx b/packages/kilo-vscode/webview-ui/src/components/chat/MessageList.tsx index b79ceab597f..df3d4c3d79d 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/MessageList.tsx +++ b/packages/kilo-vscode/webview-ui/src/components/chat/MessageList.tsx @@ -21,6 +21,7 @@ import { CloudImportDialog } from "./CloudImportDialog" import { FeedbackDialog } from "./FeedbackDialog" import { VscodeSessionTurn } from "./VscodeSessionTurn" import { WorkingIndicator } from "../shared/WorkingIndicator" +import { activeUserMessageID as getActiveUserMessageID } from "../../context/session-queue" const KiloLogo = (): JSX.Element => { const iconsBaseUri = (window as { ICONS_BASE_URI?: string }).ICONS_BASE_URI || "" @@ -72,7 +73,13 @@ export const MessageList: Component = (props) => { .slice(0, 3), ) - const lastUserMessageID = createMemo(() => userMessages().at(-1)?.id) + const activeUserID = createMemo(() => getActiveUserMessageID(session.messages(), session.statusInfo())) + + const activeUserIndex = createMemo(() => { + const active = activeUserID() + if (!active) return -1 + return userMessages().findIndex((msg) => msg.id === active) + }) return (
@@ -130,13 +137,21 @@ export const MessageList: Component = (props) => { - {(msg) => ( - - )} + {(msg, index) => { + const queued = createMemo(() => { + const active = activeUserIndex() + if (active === -1) return false + return index() > active + }) + + return ( + + ) + }} diff --git a/packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx b/packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx index d0afd90669a..0a8332605d5 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx +++ b/packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx @@ -171,10 +171,9 @@ export const PromptInput: Component = () => { const isBusy = () => session.status() === "busy" const isDisabled = () => !server.isConnected() - const canSend = () => - (text().trim().length > 0 || imageAttach.images().length > 0 || reviewComments().length > 0) && - !isBusy() && - !isDisabled() + const hasInput = () => text().trim().length > 0 || imageAttach.images().length > 0 || reviewComments().length > 0 + const canSend = () => hasInput() && !isDisabled() + const showStop = () => isBusy() && !hasInput() const placeholder = () => { switch (server.connectionState()) { case "connecting": @@ -228,11 +227,38 @@ export const PromptInput: Component = () => { } if (message.type === "triggerTask") { - if (isBusy() || isDisabled()) return + if (isDisabled()) return const sel = session.selected() session.sendMessage(message.text, sel?.providerID, sel?.modelID) } + if (message.type === "sendMessageFailed") { + const failed = message as import("../../types/messages").SendMessageFailedMessage + // Only restore draft if the failure is for the current session and the + // input is empty (user hasn't started typing something new). + const target = failed.sessionID ?? "__new__" + if (target === sessionKey() && !text().trim() && imageAttach.images().length === 0) { + if (failed.text) { + setText(failed.text) + setGhostText("") + if (textareaRef) { + textareaRef.value = failed.text + adjustHeight() + textareaRef.focus() + } + } + const images = (failed.files ?? []) + .filter((f) => f.mime.startsWith("image/") && f.url.startsWith("data:")) + .map((f) => ({ + id: crypto.randomUUID(), + filename: f.filename ?? "image", + mime: f.mime, + dataUrl: f.url, + })) + if (images.length > 0) imageAttach.replace(images) + } + } + if (message.type === "action" && message.action === "focusInput") { textareaRef?.focus() } @@ -424,10 +450,10 @@ export const PromptInput: Component = () => { const pending = reviewComments() const review = pending.length > 0 ? formatReviewCommentsMarkdown(pending) : "" const message = draft && review ? `${review}\n\n${draft}` : draft || review - if ((!message && imgs.length === 0) || isBusy() || isDisabled()) return + if ((!message && imgs.length === 0) || isDisabled()) return const mentionFiles = mention.parseFileAttachments(draft) - const imgFiles = imgs.map((img) => ({ mime: img.mime, url: img.dataUrl })) + const imgFiles = imgs.map((img) => ({ mime: img.mime, url: img.dataUrl, filename: img.filename })) const allFiles = [...mentionFiles, ...imgFiles] const sel = session.selected() @@ -611,7 +637,7 @@ export const PromptInput: Component = () => {
diff --git a/packages/kilo-vscode/webview-ui/src/context/session-queue.ts b/packages/kilo-vscode/webview-ui/src/context/session-queue.ts new file mode 100644 index 00000000000..ac64139585b --- /dev/null +++ b/packages/kilo-vscode/webview-ui/src/context/session-queue.ts @@ -0,0 +1,25 @@ +import type { Message, SessionStatusInfo } from "../types/messages" + +// Find the user message whose turn the server is actively processing. +// Any user message after this one is "queued" (waiting for its turn). +export function activeUserMessageID(messages: Message[], status: SessionStatusInfo) { + // Walk backward to find a non-completed assistant — its parent is the active turn + for (let i = messages.length - 1; i >= 0; i -= 1) { + const msg = messages[i] + if (msg.role !== "assistant") continue + if (typeof msg.time?.completed === "number") continue + if (!msg.parentID) break + const parent = messages.find((item) => item.id === msg.parentID) + if (parent?.role === "user") return parent.id + break + } + + // No pending assistant found — if busy, the last user message is the active turn + if (status.type === "idle") return undefined + + for (let i = messages.length - 1; i >= 0; i -= 1) { + if (messages[i].role === "user") return messages[i].id + } + + return undefined +} diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index 4708af9a4e3..24219c433a8 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/session.tsx @@ -38,9 +38,11 @@ import type { AgentInfo, ExtensionMessage, FileAttachment, + SendMessageFailedMessage, } from "../types/messages" import { removeSessionPermissions, upsertPermission } from "./permission-queue" import { computeStatus, calcTotalCost, calcContextUsage } from "./session-utils" +import { Identifier } from "../utils/id" // Store structure for messages and parts interface SessionStore { @@ -208,6 +210,10 @@ export const SessionProvider: ParentComponent = (props) => { // Cloud session preview state const [cloudPreviewId, setCloudPreviewId] = createSignal(null) + // 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>() + // Store for sessions, messages, parts, todos, modelSelections, agentSelections const [store, setStore] = createStore({ sessions: {}, @@ -450,6 +456,10 @@ export const SessionProvider: ParentComponent = (props) => { if (!message.sessionID || message.sessionID === currentSessionID()) setLoading(false) break + case "sendMessageFailed": + handleSendMessageFailed(message as unknown as SendMessageFailedMessage) + break + case "cloudSessionDataLoaded": handleCloudSessionDataLoaded(message.cloudSessionId, message.title, message.messages) break @@ -503,7 +513,19 @@ export const SessionProvider: ParentComponent = (props) => { function handleMessagesLoaded(sessionID: string, messages: Message[]) { batch(() => { if (sessionID === currentSessionID()) setLoading(false) - setStore("messages", sessionID, messages) + + // Preserve optimistic messages that haven't been confirmed yet. + // The server may not have created the message record by the time + // this session's messages are loaded (e.g. on session switch). + const pending = pendingOptimistic.get(sessionID) + if (pending && pending.size > 0) { + const loadedIds = new Set(messages.map((m) => m.id)) + const current = store.messages[sessionID] ?? [] + const orphans = current.filter((m) => pending.has(m.id) && !loadedIds.has(m.id)) + setStore("messages", sessionID, [...messages, ...orphans]) + } else { + setStore("messages", sessionID, messages) + } // Also extract parts from messages for (const msg of messages) { @@ -515,31 +537,33 @@ export const SessionProvider: ParentComponent = (props) => { } function handleMessageCreated(message: Message) { + // 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). + 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] + }), + ) + } + setStore("messages", message.sessionID, (msgs = []) => { - // Check if message already exists (update case) - const existingIndex = msgs.findIndex((m) => m.id === message.id) - if (existingIndex >= 0) { + // Check if message already exists (optimistic or update case). + // Since we now use the same messageID for optimistic and server messages, + // this naturally handles the optimistic→real transition. + const idx = msgs.findIndex((m) => m.id === message.id) + if (idx >= 0) { const updated = [...msgs] - updated[existingIndex] = { ...msgs[existingIndex], ...message } + updated[idx] = { ...msgs[idx], ...message } return updated } - // Replace optimistic user message if one exists - if (message.role === "user") { - const optimisticIdx = msgs.findIndex((m) => m.id.startsWith("optimistic-") && m.role === "user") - if (optimisticIdx >= 0) { - const updated = [...msgs] - // Clean up optimistic parts - const old = msgs[optimisticIdx] - setStore( - "parts", - produce((parts) => { - delete parts[old.id] - }), - ) - updated[optimisticIdx] = message - return updated - } - } return [...msgs, message] }) @@ -615,6 +639,11 @@ export const SessionProvider: ParentComponent = (props) => { delete map[sessionID] }), ) + // 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 + // messages themselves will be reconciled on the next messagesLoaded. + pendingOptimistic.delete(sessionID) } } @@ -669,6 +698,32 @@ export const SessionProvider: ParentComponent = (props) => { setQuestionErrors((prev) => new Set(prev).add(requestID)) } + /** + * Handle a failed send: remove the optimistic message from the store + * and show a toast. The PromptInput restores the draft text separately + * by listening for the same sendMessageFailed event. + */ + function handleSendMessageFailed(message: SendMessageFailedMessage) { + if (message.sessionID && message.messageID) { + pendingOptimistic.get(message.sessionID)?.delete(message.messageID) + batch(() => { + setStore("messages", message.sessionID!, (msgs = []) => msgs.filter((m) => m.id !== message.messageID)) + setStore( + "parts", + produce((parts) => { + delete parts[message.messageID!] + }), + ) + }) + } + + showToast({ + variant: "error", + title: language.t("prompt.toast.promptSendFailed.title") ?? "Failed to send message", + description: message.error, + }) + } + /** * BFS walk over message parts to discover all session IDs in a session's * family tree (self + subagents + sub-subagents). Reads directly from the @@ -736,6 +791,7 @@ export const SessionProvider: ParentComponent = (props) => { } function handleSessionDeleted(sessionID: string) { + pendingOptimistic.delete(sessionID) batch(() => { // Collect message IDs so we can clean up their parts const msgs = store.messages[sessionID] ?? [] @@ -907,6 +963,8 @@ export const SessionProvider: ParentComponent = (props) => { return } + const messageID = Identifier.ascending("message") + const preview = cloudPreviewId() if (preview) { const agent = selectedAgentName() !== defaultAgent() ? selectedAgentName() : undefined @@ -914,6 +972,7 @@ export const SessionProvider: ParentComponent = (props) => { type: "importAndSend", cloudSessionId: preview, text, + messageID, providerID, modelID, agent, @@ -925,17 +984,36 @@ export const SessionProvider: ParentComponent = (props) => { const sid = currentSessionID() if (sid) { - const tempId = `optimistic-${crypto.randomUUID()}` const now = Date.now() const temp: Message = { - id: tempId, + id: messageID, sessionID: sid, role: "user", createdAt: new Date(now).toISOString(), time: { created: now }, } + // Track this optimistic message so handleMessagesLoaded preserves it + const pending = pendingOptimistic.get(sid) ?? new Set() + pending.add(messageID) + pendingOptimistic.set(sid, pending) + + const parts: Part[] = [] + if (text) { + parts.push({ type: "text" as const, id: Identifier.ascending("part"), messageID, text }) + } + for (const file of files ?? []) { + parts.push({ + type: "file" as const, + id: Identifier.ascending("part"), + messageID, + mime: file.mime, + url: file.url, + filename: file.filename, + }) + } + setStore("messages", sid, (msgs = []) => [...msgs, temp]) - setStore("parts", tempId, [{ type: "text" as const, id: `${tempId}-text`, text }]) + setStore("parts", messageID, parts) // The optimistic message is now in the DOM but the session status is // still "idle" (the CLI backend hasn't started yet), so the auto-scroll // ResizeObserver won't scroll on its own. Force scroll to bottom so the @@ -948,6 +1026,7 @@ export const SessionProvider: ParentComponent = (props) => { vscode.postMessage({ type: "sendMessage", text, + messageID, sessionID: sid, providerID, modelID, @@ -1056,10 +1135,6 @@ export const SessionProvider: ParentComponent = (props) => { setCurrentSessionID(undefined) setCloudPreviewId(null) setLoading(false) - setPermissions([]) - setRespondingPermissions(new Set()) - setQuestions([]) - setQuestionErrors(new Set()) setPendingAgentSelection(defaultAgent()) vscode.postMessage({ type: "clearSession" }) } diff --git a/packages/kilo-vscode/webview-ui/src/hooks/useImageAttachments.ts b/packages/kilo-vscode/webview-ui/src/hooks/useImageAttachments.ts index c1c3615a3ac..c72d6f123fd 100644 --- a/packages/kilo-vscode/webview-ui/src/hooks/useImageAttachments.ts +++ b/packages/kilo-vscode/webview-ui/src/hooks/useImageAttachments.ts @@ -33,6 +33,8 @@ export function useImageAttachments() { const clear = () => setImages([]) + const replace = (next: ImageAttachment[]) => setImages(next) + const handlePaste = (event: ClipboardEvent) => { const items = Array.from(event.clipboardData?.items ?? []) const imageItems = items.filter((item) => item.kind === "file" && ACCEPTED_IMAGE_TYPES.includes(item.type)) @@ -71,6 +73,7 @@ export function useImageAttachments() { add, remove, clear, + replace, handlePaste, handleDragOver, handleDragLeave, diff --git a/packages/kilo-vscode/webview-ui/src/types/messages.ts b/packages/kilo-vscode/webview-ui/src/types/messages.ts index 263ea97a84d..8b4543fc5d9 100644 --- a/packages/kilo-vscode/webview-ui/src/types/messages.ts +++ b/packages/kilo-vscode/webview-ui/src/types/messages.ts @@ -34,6 +34,13 @@ export interface TextPart extends BasePart { text: string } +export interface FilePart extends BasePart { + type: "file" + mime: string + url: string + filename?: string +} + export interface ToolPart extends BasePart { type: "tool" tool: string @@ -62,7 +69,7 @@ export interface StepFinishPart extends BasePart { } } -export type Part = TextPart | ToolPart | ReasoningPart | StepStartPart | StepFinishPart +export type Part = TextPart | FilePart | ToolPart | ReasoningPart | StepStartPart | StepFinishPart // Part delta for streaming updates export interface PartDelta { @@ -374,6 +381,15 @@ export interface ErrorMessage { sessionID?: string } +export interface SendMessageFailedMessage { + type: "sendMessageFailed" + error: string + text: string + sessionID?: string + messageID?: string + files?: FileAttachment[] +} + export interface PartUpdatedMessage { type: "partUpdated" sessionID?: string @@ -1018,6 +1034,7 @@ export type ExtensionMessage = | ReadyMessage | ConnectionStateMessage | ErrorMessage + | SendMessageFailedMessage | PartUpdatedMessage | SessionStatusMessage | PermissionRequestMessage @@ -1100,11 +1117,13 @@ export type ExtensionMessage = export interface FileAttachment { mime: string url: string + filename?: string } export interface SendMessageRequest { type: "sendMessage" text: string + messageID?: string sessionID?: string providerID?: string modelID?: string @@ -1164,6 +1183,7 @@ export interface ImportAndSendMessage { type: "importAndSend" cloudSessionId: string text: string + messageID?: string providerID?: string modelID?: string agent?: string diff --git a/packages/kilo-vscode/webview-ui/src/utils/id.ts b/packages/kilo-vscode/webview-ui/src/utils/id.ts new file mode 100644 index 00000000000..c31726bb7a2 --- /dev/null +++ b/packages/kilo-vscode/webview-ui/src/utils/id.ts @@ -0,0 +1,60 @@ +type Prefix = "message" | "part" + +const prefixes: Record = { + message: "msg", + part: "prt", +} + +const RANDOM_LENGTH = 14 +let lastTimestamp = 0 +let counter = 0 + +// Browser-side sortable IDs for optimistic queued messages/parts. +// They need to sort the same way as backend-generated IDs so queued turns +// stay in order before SSE catches up. +export const Identifier = { + ascending(prefix: Prefix) { + const currentTimestamp = Date.now() + + if (currentTimestamp !== lastTimestamp) { + lastTimestamp = currentTimestamp + counter = 0 + } + + counter += 1 + + // Server uses 6 bytes big-endian (48 bits). Truncate to 12 hex chars + // to match — without this, the full BigInt produces 14 chars which + // breaks the server's lexicographic exit condition (lastUser < lastAssistant). + const sortable = (BigInt(currentTimestamp) * BigInt(0x1000) + BigInt(counter)) + .toString(16) + .slice(-12) + .padStart(12, "0") + return `${prefixes[prefix]}_${sortable}${randomBase62(RANDOM_LENGTH)}` + }, +} + +function randomBase62(length: number) { + const chars = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" + const bytes = getRandomBytes(length) + let result = "" + for (let i = 0; i < length; i += 1) { + result += chars[bytes[i] % 62] + } + return result +} + +function getRandomBytes(length: number) { + const bytes = new Uint8Array(length) + + if (typeof globalThis.crypto?.getRandomValues === "function") { + globalThis.crypto.getRandomValues(bytes) + return bytes + } + + for (let i = 0; i < length; i += 1) { + bytes[i] = Math.floor(Math.random() * 256) + } + + return bytes +}