diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index b08904c432..52cbca8224 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -25,6 +25,8 @@ import { mapSSEEventToWebviewMessage, getErrorMessage, isEventFromForeignProject, + MessageConfirmation, + runWithMessageConfirmation, loadSessions as loadSessionsUtil, flushPendingSessionRefresh as flushPendingSessionRefreshUtil, resolveContextDirectory, @@ -164,6 +166,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper /** Set when refreshSessions() is called before the client is ready. * Cleared and retried once the connection transitions to "connected". */ private pendingSessionRefresh = false + private readonly confirmations = new MessageConfirmation() private unsubscribeEvent: (() => void) | null = null private unsubscribeState: (() => void) | null = null /** Cached legacy migration data so migrate() doesn't re-read from disk/SecretStorage. */ // legacy-migration @@ -2297,21 +2300,12 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper /** Abort controllers for active retry loops, keyed by session ID */ private retryAbortControllers = new Map() - /** - * Execute an SDK call with exponential backoff on HTTP errors. - * Retries on 429, 5xx, and other retryable status codes. - * When the response includes `Retry-After` / `Retry-After-MS` headers, - * the delay honours that value (capped at 5 min). Otherwise uses the - * predefined backoff schedule: 5s -> 10s -> 30s -> 60s -> 300s. - * - * After MAX_RETRIES (5) attempts, automatically throws the error. - * Users can cancel via the cancel button in the UI which sends an abort - * message — this interrupts the backoff delay and stops the retry loop. - * - * The webview receives `sessionStatus` messages with a countdown so the - * user can see that a retry is in progress. - */ - private async withRetry(fn: () => Promise<{ error?: unknown; response: Response }>, sid: string): Promise { + /** Execute an SDK call with visible exponential backoff for retryable HTTP errors. */ + private async withRetry( + fn: () => Promise<{ error?: unknown; response?: Response }>, + sid: string, + messageID?: string, + ): Promise { const abortController = new AbortController() this.retryAbortControllers.set(sid, abortController) @@ -2324,6 +2318,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper const result = await fn() if (!result.error) return + if (this.confirmations.has(messageID)) return const status = result.response?.status ?? 0 @@ -2352,12 +2347,16 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper }) // Wait for delay or until aborted - await new Promise((resolve) => { - const timer = setTimeout(resolve, delay) - abortController.signal.addEventListener("abort", () => { + await new Promise((resolve) => { + const done = () => { clearTimeout(timer) - }) + abortController.signal.removeEventListener("abort", done) + resolve() + } + const timer = setTimeout(done, delay) + abortController.signal.addEventListener("abort", done, { once: true }) }) + if (this.confirmations.has(messageID)) return } } finally { this.retryAbortControllers.delete(sid) @@ -2417,19 +2416,22 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper const sid = resolved!.sid const dir = resolved!.dir - await this.withRetry( - () => - this.client!.session.promptAsync({ - sessionID: sid, - directory: dir, - messageID, - parts, - model: providerID && modelID ? { providerID, modelID } : undefined, - agent, - variant, - editorContext, - }), - sid, + await runWithMessageConfirmation(this.confirmations, messageID, "KiloProvider: Message request", () => + this.withRetry( + () => + this.client!.session.promptAsync({ + sessionID: sid, + directory: dir, + messageID, + parts, + model: providerID && modelID ? { providerID, modelID } : undefined, + agent, + variant, + editorContext, + }), + sid, + messageID, + ), ) } catch (error) { console.error("[Kilo New] KiloProvider: Failed to send message:", error) @@ -2482,20 +2484,23 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper const sid = resolved!.sid const dir = resolved!.dir - await this.withRetry( - () => - this.client!.session.command({ - sessionID: sid, - directory: dir, - command, - arguments: args, - messageID, - model: providerID && modelID ? `${providerID}/${modelID}` : undefined, - agent, - variant, - parts, - }), - sid, + await runWithMessageConfirmation(this.confirmations, messageID, "KiloProvider: Command request", () => + this.withRetry( + () => + this.client!.session.command({ + sessionID: sid, + directory: dir, + command, + arguments: args, + messageID, + model: providerID && modelID ? `${providerID}/${modelID}` : undefined, + agent, + variant, + parts, + }), + sid, + messageID, + ), ) } catch (error) { console.error("[Kilo New] KiloProvider: Failed to send command:", error) @@ -2640,6 +2645,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper postMessage: (msg) => this.postMessage(msg), getWorkspaceDirectory: (sid) => this.getWorkspaceDirectory(sid), gatherEditorContext: () => this.gatherEditorContext(), + runWithMessageConfirmation: (id, label, run) => runWithMessageConfirmation(this.confirmations, id, label, run), } } @@ -2853,6 +2859,10 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper // let a foreign session through if it was accidentally tracked. if (isEventFromForeignProject(event, this.projectID)) return + if (event.type === "message.updated") { + this.confirmations.confirm(event.properties.info.id) + } + // session.status events pass the onEventFiltered pre-filter for all providers (see line 842), // so this runs on every KiloProvider instance — including the Settings panel which has no // tracked sessions. Update sessionStatusMap and forward to webview before the diff --git a/packages/kilo-vscode/src/kilo-provider-utils.ts b/packages/kilo-vscode/src/kilo-provider-utils.ts index a7dbc21b80..8470c74918 100644 --- a/packages/kilo-vscode/src/kilo-provider-utils.ts +++ b/packages/kilo-vscode/src/kilo-provider-utils.ts @@ -58,6 +58,81 @@ export function getErrorMessage(error: unknown): string { return String(error) } +export class MessageConfirmation { + private readonly ids = new Map void> }>() + + track(id?: string): () => void { + if (!id) return () => {} + const entry = this.ids.get(id) ?? { confirmed: false, waits: new Set<() => void>() } + this.ids.set(id, entry) + return () => { + this.ids.delete(id) + } + } + + confirm(id: string): void { + const entry = this.ids.get(id) + if (!entry) return + entry.confirmed = true + for (const done of [...entry.waits]) { + done() + } + } + + has(id?: string): boolean { + if (!id) return false + return this.ids.get(id)?.confirmed ?? false + } + + wait(id?: string, timeout = 1_500): Promise { + if (!id) return Promise.resolve(false) + const entry = this.ids.get(id) + if (!entry) return Promise.resolve(false) + if (entry.confirmed) return Promise.resolve(true) + + return new Promise((resolve) => { + const timer = setTimeout(() => { + cleanup() + resolve(entry.confirmed) + }, timeout) + + const cleanup = () => { + clearTimeout(timer) + entry.waits.delete(done) + } + + const done = () => { + cleanup() + resolve(true) + } + + entry.waits.add(done) + }) + } +} + +export async function runWithMessageConfirmation( + state: MessageConfirmation, + id: string | undefined, + label: string, + run: () => Promise, +): Promise { + const release = state.track(id) + try { + return await run() + } catch (error) { + if (await state.wait(id)) { + console.warn(`[Kilo New] ${label} ended after server accepted it; ignoring transport error`, { + error: getErrorMessage(error), + }) + return undefined + } + throw error + } finally { + release() + } +} + export function sessionToWebview(session: Session) { return { id: session.id, diff --git a/packages/kilo-vscode/src/kilo-provider/handlers/cloud-session.ts b/packages/kilo-vscode/src/kilo-provider/handlers/cloud-session.ts index 20f0982818..acb0365991 100644 --- a/packages/kilo-vscode/src/kilo-provider/handlers/cloud-session.ts +++ b/packages/kilo-vscode/src/kilo-provider/handlers/cloud-session.ts @@ -19,6 +19,11 @@ export interface CloudSessionContext { postMessage(msg: unknown): void getWorkspaceDirectory(sessionId?: string): string gatherEditorContext(): Promise + runWithMessageConfirmation?( + messageID: string | undefined, + label: string, + run: () => Promise, + ): Promise } /** Fetch cloud sessions list and send to webview. */ @@ -123,6 +128,7 @@ export async function handleImportAndSend( return } + const client = ctx.client const dir = ctx.getWorkspaceDirectory() // Step 1: Import the cloud session with fresh IDs @@ -163,28 +169,32 @@ export async function handleImportAndSend( }) // Step 2: Send the user's message/command on the new local session + const run = ctx.runWithMessageConfirmation ?? ((_id, _label, fn) => fn()) try { - if (messageID) { - ctx.connectionService.recordMessageSessionId(messageID, session.id) - } + await run(messageID, "Cloud import send", async () => { + if (messageID) { + ctx.connectionService.recordMessageSessionId(messageID, session.id) + } + + if (command) { + const parts = files?.map((f) => ({ type: "file" as const, mime: f.mime, url: f.url })) + await client.session.command( + { + sessionID: session.id, + directory: dir, + command, + arguments: commandArgs ?? "", + messageID, + model: providerID && modelID ? `${providerID}/${modelID}` : undefined, + agent, + variant, + parts, + }, + { throwOnError: true }, + ) + return + } - if (command) { - const parts = files?.map((f) => ({ type: "file" as const, mime: f.mime, url: f.url })) - await ctx.client.session.command( - { - sessionID: session.id, - directory: dir, - command, - arguments: commandArgs ?? "", - messageID, - model: providerID && modelID ? `${providerID}/${modelID}` : undefined, - agent, - variant, - parts, - }, - { throwOnError: true }, - ) - } else { const parts: Array = [] if (files) { for (const f of files) { @@ -194,7 +204,7 @@ export async function handleImportAndSend( parts.push({ type: "text", text }) const editorContext = await ctx.gatherEditorContext() - await ctx.client.session.promptAsync( + await client.session.promptAsync( { sessionID: session.id, directory: dir, @@ -207,7 +217,7 @@ export async function handleImportAndSend( }, { throwOnError: true }, ) - } + }) } catch (err) { console.error("[Kilo New] Failed to send message after cloud import:", err) ctx.postMessage({ diff --git a/packages/kilo-vscode/src/services/cli-backend/sdk-sse-adapter.ts b/packages/kilo-vscode/src/services/cli-backend/sdk-sse-adapter.ts index 0021e4e7a9..2b54ac11bb 100644 --- a/packages/kilo-vscode/src/services/cli-backend/sdk-sse-adapter.ts +++ b/packages/kilo-vscode/src/services/cli-backend/sdk-sse-adapter.ts @@ -175,7 +175,10 @@ export class SdkSSEAdapter { // The SDK yields GlobalEvent = { directory, payload: Event }. const globalEvent = event as GlobalEvent - console.log("[Kilo New] SSE: 📨 Event:", globalEvent.payload.type) + const type = (globalEvent.payload as { type: string }).type + if (type !== "server.heartbeat") { + console.log("[Kilo New] SSE: 📨 Event:", type) + } this.notifyEvent(globalEvent.payload) } diff --git a/packages/kilo-vscode/tests/unit/kilo-provider-utils.test.ts b/packages/kilo-vscode/tests/unit/kilo-provider-utils.test.ts index 454997536b..f53b5a1d36 100644 --- a/packages/kilo-vscode/tests/unit/kilo-provider-utils.test.ts +++ b/packages/kilo-vscode/tests/unit/kilo-provider-utils.test.ts @@ -7,6 +7,7 @@ import { mapSSEEventToWebviewMessage, isEventFromForeignProject, mapCloudSessionMessageToWebviewMessage, + MessageConfirmation, type ProviderInfo, } from "../../src/kilo-provider-utils" import type { CloudSessionMessage } from "../../src/services/cli-backend/types" @@ -93,6 +94,44 @@ function makeAssistantMessage(overrides: Partial = {}): Assist } } +describe("MessageConfirmation", () => { + it("reports tracked confirmed messages", async () => { + const state = new MessageConfirmation() + state.track("msg-1") + state.confirm("msg-1") + + expect(state.has("msg-1")).toBe(true) + expect(await state.wait("msg-1", 1)).toBe(true) + }) + + it("resolves waiters when a message is confirmed", async () => { + const state = new MessageConfirmation() + state.track("msg-1") + const wait = state.wait("msg-1", 50) + + state.confirm("msg-1") + + expect(await wait).toBe(true) + }) + + it("returns false when confirmation does not arrive", async () => { + const state = new MessageConfirmation() + state.track("msg-1") + + expect(await state.wait("msg-1", 1)).toBe(false) + }) + + it("forgets confirmations after release", () => { + const state = new MessageConfirmation() + const release = state.track("msg-1") + state.confirm("msg-1") + + release() + + expect(state.has("msg-1")).toBe(false) + }) +}) + describe("sessionToWebview", () => { it("converts epoch timestamps to ISO strings", () => { const result = sessionToWebview(makeSession())