From 618d956a9c1341bf03be5cb438a1b2038d08e665 Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Fri, 28 Aug 2026 11:53:09 +0200 Subject: [PATCH] refactor(vscode): extract prompt sandbox response handling --- .../prompt-input-connection-guard.test.ts | 43 +++- .../unit/prompt-sandbox-messages.test.ts | 241 ++++++++++++++++++ .../src/components/chat/PromptInput.tsx | 112 ++------ .../chat/prompt-sandbox-messages.ts | 98 +++++++ 4 files changed, 398 insertions(+), 96 deletions(-) create mode 100644 packages/kilo-vscode/tests/unit/prompt-sandbox-messages.test.ts create mode 100644 packages/kilo-vscode/webview-ui/src/components/chat/prompt-sandbox-messages.ts diff --git a/packages/kilo-vscode/tests/unit/prompt-input-connection-guard.test.ts b/packages/kilo-vscode/tests/unit/prompt-input-connection-guard.test.ts index fe8d3cd3122f..ec6b4fbd84f6 100644 --- a/packages/kilo-vscode/tests/unit/prompt-input-connection-guard.test.ts +++ b/packages/kilo-vscode/tests/unit/prompt-input-connection-guard.test.ts @@ -6,6 +6,10 @@ const path = join(__dirname, "..", "..", "webview-ui", "src", "components", "cha const buttonPath = join(__dirname, "..", "..", "webview-ui", "src", "components", "shared", "SandboxButton.tsx") const iconPath = join(__dirname, "..", "..", "..", "kilo-ui", "src", "components", "icon.tsx") const src = readFileSync(path, "utf8") +const responses = readFileSync( + join(__dirname, "..", "..", "webview-ui", "src", "components", "chat", "prompt-sandbox-messages.ts"), + "utf8", +) const button = readFileSync(buttonPath, "utf8") const icons = readFileSync(iconPath, "utf8") @@ -79,7 +83,9 @@ describe("PromptInput sandbox toggle", () => { expect(src).toContain('const sandboxRequest = (sessionID?: string) => sandboxRequests()[sessionID ?? ""]') expect(src).toContain("sandboxRequest(sandboxID()) !== undefined") expect(src).toContain("if (current[key] !== requestID) return current") - expect(src).toContain("clearSandboxRequest(message.sessionID, message.requestID!)") + expect(src).toContain("pending: sandboxRequest") + expect(src).toContain("clear: clearSandboxRequest") + expect(responses).toContain("input.clear(message.sessionID, message.requestID!)") expect(src).not.toContain("setSandboxTarget") }) @@ -92,13 +98,13 @@ describe("PromptInput sandbox toggle", () => { expect(src).toContain("{ action: toggleSandbox, enabled: () => sandboxVisible() && !sandboxDisabled() }") expect(src).toContain('if (!sandboxVisible()) hidden.add("sandbox")') expect(src).toContain("onToggle={toggleSandbox}") - expect(src).toContain('message.type === "sandboxStatus"') - expect(src).not.toContain("message.sessionID !== sandboxID() && !matching") - expect(src).toContain("const next = applySandboxStates(current, message)") - expect(src).toContain("if (next !== current) setSandboxes(next)") - expect(src).toContain("message.requestID === sandboxRequest(message.sessionID)") - expect(src).toContain("if (message.sessionID === sandboxID())") - expect(src).toContain("if (message.sessionID === sandboxID()) retrySandbox(message.sessionID)") + expect(responses).toContain('case "sandboxStatus":') + expect(responses).not.toContain("message.sessionID !== input.session() && !matching") + expect(responses).toContain("const next = applySandboxStates(current, message)") + expect(responses).toContain("if (next !== current) input.setStates(next)") + expect(responses).toContain("message.requestID === input.pending(message.sessionID)") + expect(responses).toContain("if (message.sessionID === input.session()) input.reset()") + expect(responses).toContain("if (message.sessionID === input.session()) input.retry(message.sessionID)") expect(src).toContain("sandboxID() ? sandbox()?.enabled : sandboxDefault()?.enabled") expect(src).toContain('type: "requestSandboxDefault", agentManagerContext: ctx()') expect(src).toContain(" { expect(button).toContain("aria-pressed={props.enabled}") expect(button).toContain('class={`prompt-status-button ${props.enabled ? "prompt-status-button--active" : ""}`}') expect(src).toContain("if (sandboxRequest(undefined)) return") - expect(src).not.toContain("if (state === current) return true") + expect(responses).not.toContain("if (state === current) return true") + }) + + it("keeps the response subscription and retry timer in the prompt", () => { + expect(src).toContain("const handleSandboxMessage = sandboxMessages({") + expect(src).toContain("connected: server.isConnected") + expect(src).toContain("session: sandboxID") + expect(src).toContain("defaults: sandboxDefault") + expect(src).toContain("setDefault: setSandboxDefault") + expect(src).toContain("states: sandboxes") + expect(src).toContain("setStates: setSandboxes") + expect(src).toContain("retry: retrySandbox") + expect(src).toContain("refresh: requestSandbox") + expect(src).toContain( + "reset: () => {\n sandboxAttempts = 0\n if (sandboxRetry) clearTimeout(sandboxRetry)\n sandboxRetry = undefined", + ) + expect(src).toContain( + "const unsubscribe = vscode.onMessage((message) => {\n if (handleSandboxMessage(message)) return", + ) + expect(src).toContain("unsubscribe()") }) it("preserves the draft when the sandbox control is disabled", () => { diff --git a/packages/kilo-vscode/tests/unit/prompt-sandbox-messages.test.ts b/packages/kilo-vscode/tests/unit/prompt-sandbox-messages.test.ts new file mode 100644 index 000000000000..4810b7af1dee --- /dev/null +++ b/packages/kilo-vscode/tests/unit/prompt-sandbox-messages.test.ts @@ -0,0 +1,241 @@ +import { describe, expect, it } from "bun:test" +import { sandboxMessages } from "../../webview-ui/src/components/chat/prompt-sandbox-messages" +import type { SandboxDefaultState, SandboxState } from "../../webview-ui/src/components/chat/prompt-input-utils" +import type { + SandboxDefaultStatusMessage, + SandboxStatusMessage, + SandboxStatusErrorMessage, +} from "../../webview-ui/src/types/messages" + +const preference: SandboxDefaultState = { + desired: true, + enabled: true, + available: true, + reason: undefined, + revision: 3, +} +const defaults: SandboxDefaultStatusMessage = { type: "sandboxDefaultStatus", ...preference } +const status: SandboxStatusMessage = { + type: "sandboxStatus", + sessionID: "ses_1", + directory: "/repo", + enabled: true, + available: true, + version: 2, + revision: 3, +} +const failure: SandboxStatusErrorMessage = { + type: "sandboxStatusError", + sessionID: "ses_1", + directory: "/repo", + message: "Sandbox unavailable", + revision: 3, +} + +function setup( + input: { + connected?: boolean + session?: string + defaults?: SandboxDefaultState + states?: Record + requests?: Record + } = {}, +) { + const state = { + ...input, + connected: input.connected ?? true, + states: input.states ?? {}, + requests: input.requests ?? {}, + } + const calls: unknown[][] = [] + const handle = sandboxMessages({ + connected: () => state.connected, + session: () => state.session, + pending: (id) => state.requests[id ?? ""], + clear: (id, request) => calls.push(["clear", id, request]), + defaults: () => state.defaults, + setDefault: (value) => { + state.defaults = value + calls.push(["default"]) + }, + states: () => state.states, + setStates: (value) => { + state.states = value + calls.push(["states"]) + }, + reset: () => calls.push(["reset"]), + retry: (id) => calls.push(["retry", id]), + refresh: () => calls.push(["refresh"]), + error: (reason) => calls.push(["error", reason]), + }) + return { state, calls, handle } +} + +describe("sandboxMessages", () => { + it("dispatches config updates and leaves unrelated messages unhandled", () => { + const fixture = setup({ connected: false }) + expect(fixture.handle({ type: "configUpdated", config: {} })).toBe(true) + expect(fixture.handle({ type: "autoApproveState", active: true })).toBe(false) + expect(fixture.calls).toEqual([["refresh"]]) + }) + + it.each([defaults, status, failure])("ignores $type responses while disconnected", (message) => { + const fixture = setup({ connected: false, requests: { "": "request", ses_1: "request" } }) + const states = fixture.state.states + expect(fixture.handle({ ...message, requestID: "request" })).toBe(true) + expect(fixture.state.states).toBe(states) + expect(fixture.state.defaults).toBeUndefined() + expect(fixture.calls).toEqual([]) + }) + + it.each([undefined, "unmatched"])("loads defaults for a draft with request ID %s", (requestID) => { + const fixture = setup({ requests: { "": "pending" } }) + expect(fixture.handle({ ...defaults, requestID })).toBe(true) + expect(fixture.state.defaults).toEqual(preference) + expect(fixture.calls).toEqual([["default"]]) + }) + + it.each([true, false])("leaves unmatched defaults unhandled in a session when connected is %s", (connected) => { + const fixture = setup({ session: "ses_1", connected, requests: { "": "pending" } }) + expect(fixture.handle({ ...defaults, requestID: "unmatched" })).toBe(false) + expect(fixture.state.defaults).toBeUndefined() + expect(fixture.calls).toEqual([]) + }) + + it("handles a matching default response after switching to a session", () => { + const fixture = setup({ requests: { "": "request" } }) + fixture.state.session = "ses_1" + expect(fixture.handle({ ...defaults, available: false, reason: "Unavailable", requestID: "request" })).toBe(true) + expect(fixture.state.defaults).toEqual({ ...preference, available: false, reason: "Unavailable" }) + expect(fixture.calls).toEqual([["clear", undefined, "request"], ["default"], ["error", "Unavailable"]]) + }) + + it("keeps newer defaults but still clears and reports a matching stale failure", () => { + const current = { ...preference, revision: 4 } + const fixture = setup({ defaults: current, requests: { "": "request" } }) + expect(fixture.handle({ ...defaults, available: false, requestID: "request" })).toBe(true) + expect(fixture.state.defaults).toBe(current) + expect(fixture.calls).toEqual([ + ["clear", undefined, "request"], + ["error", undefined], + ]) + }) + + it("accepts equal default revisions without reporting unmatched unavailability", () => { + const fixture = setup({ defaults: preference }) + expect(fixture.handle({ ...defaults, desired: false, enabled: false, available: false })).toBe(true) + expect(fixture.state.defaults).toEqual({ ...preference, desired: false, enabled: false, available: false }) + expect(fixture.calls).toEqual([["default"]]) + }) + + it.each([undefined, "unmatched"])("updates session status without clearing request ID %s", (requestID) => { + const fixture = setup({ session: "ses_1", requests: { ses_1: "pending" } }) + const message = { ...status, available: false, requestID } + expect(fixture.handle(message)).toBe(true) + expect(fixture.state.states).toEqual({ ses_1: message }) + expect(fixture.calls).toEqual([["states"], ["reset"]]) + }) + + it("clears matching toggles and reports the applied session status", () => { + const fixture = setup({ session: "ses_1", requests: { ses_1: "request" } }) + const message = { ...status, available: false, reason: "Unavailable", requestID: "request" } + expect(fixture.handle(message)).toBe(true) + expect(fixture.state.states).toEqual({ ses_1: message }) + expect(fixture.calls).toEqual([["clear", "ses_1", "request"], ["states"], ["reset"], ["error", "Unavailable"]]) + }) + + it("caches other sessions without resetting the active session retry", () => { + const fixture = setup({ session: "ses_1", states: { ses_1: status }, requests: { ses_2: "request" } }) + const message = { ...status, sessionID: "ses_2", requestID: "request" } + expect(fixture.handle(message)).toBe(true) + expect(fixture.state.states).toEqual({ ses_1: status, ses_2: message }) + expect(fixture.calls).toEqual([["clear", "ses_2", "request"], ["states"]]) + }) + + it.each([{ revision: 4 }, { version: 3 }])("keeps newer session state ordered by %j", (ordering) => { + const current = { ...status, ...ordering } + const fixture = setup({ session: "ses_1", states: { ses_1: current }, requests: { ses_1: "request" } }) + const states = fixture.state.states + expect(fixture.handle({ ...status, available: false, requestID: "request" })).toBe(true) + expect(fixture.state.states).toBe(states) + expect(fixture.calls).toEqual([["clear", "ses_1", "request"], ["reset"]]) + }) + + it("reports retained unavailable state rather than a stale status reason", () => { + const current = { ...status, available: false, reason: "Current failure", revision: 4 } + const fixture = setup({ states: { ses_1: current }, requests: { ses_1: "request" } }) + expect(fixture.handle({ ...status, requestID: "request" })).toBe(true) + expect(fixture.state.states.ses_1).toBe(current) + expect(fixture.calls).toEqual([ + ["clear", "ses_1", "request"], + ["error", "Current failure"], + ]) + }) + + it("preserves same-directory state on status errors and retries only the active session", () => { + const fixture = setup({ session: "ses_1", states: { ses_1: status } }) + expect(fixture.handle(failure)).toBe(true) + expect(fixture.state.states.ses_1).toEqual({ + sessionID: "ses_1", + directory: "/repo", + enabled: true, + available: false, + reason: failure.message, + version: 2, + revision: 3, + }) + expect(fixture.calls).toEqual([["states"], ["retry", "ses_1"]]) + fixture.calls.length = 0 + fixture.state.session = "ses_2" + expect(fixture.handle({ ...failure, revision: 4 })).toBe(true) + expect(fixture.calls).toEqual([["states"]]) + }) + + it.each([undefined, { ...status, directory: "/other" }])( + "resets unknown or different-directory error state %j", + (current) => { + const fixture = setup({ states: current ? { ses_1: current } : {} }) + expect(fixture.handle(failure)).toBe(true) + expect(fixture.state.states.ses_1).toEqual({ + sessionID: "ses_1", + directory: "/repo", + enabled: false, + available: false, + reason: failure.message, + version: 0, + revision: 3, + }) + expect(fixture.calls).toEqual([["states"]]) + }, + ) + + it.each([undefined, "unmatched", "request"])( + "ignores stale errors with request ID %s after matching cleanup", + (requestID) => { + const current = { ...status, revision: 4 } + const fixture = setup({ session: "ses_1", states: { ses_1: current }, requests: { ses_1: "request" } }) + const states = fixture.state.states + expect(fixture.handle({ ...failure, requestID })).toBe(true) + expect(fixture.state.states).toBe(states) + expect(fixture.calls).toEqual(requestID === "request" ? [["clear", "ses_1", "request"]] : []) + }, + ) + + it.each(["unmatched", "request"])( + "handles toggle errors with request ID %s without state updates or retry", + (requestID) => { + const fixture = setup({ session: "ses_1", states: { ses_1: status }, requests: { ses_1: "request" } }) + const states = fixture.state.states + expect(fixture.handle({ ...failure, requestID })).toBe(true) + expect(fixture.state.states).toBe(states) + expect(fixture.calls).toEqual( + requestID === "request" + ? [ + ["clear", "ses_1", "request"], + ["error", failure.message], + ] + : [], + ) + }, + ) +}) 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 16a7b15e4e1e..a133944d3933 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx +++ b/packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx @@ -52,12 +52,12 @@ import { insertSpacedText, isPromptBusy, isPathMention, - applySandboxStates, memoryRest, type SandboxDefaultState, type SandboxState, } from "./prompt-input-utils" -import type { ExtensionMessage, ReviewCommentEntry, SendMessageFailedMessage, TextPart } from "../../types/messages" +import { sandboxMessages } from "./prompt-sandbox-messages" +import type { ReviewCommentEntry, SendMessageFailedMessage, TextPart } from "../../types/messages" import { formatReviewCommentsMarkdown } from "../../utils/review-comment-markdown" import { createdDraftKey, @@ -632,91 +632,29 @@ export const PromptInput: Component = (props) => { imageDrafts.set(target, images) } - const handleSandboxMessage = (message: ExtensionMessage) => { - if (message.type === "sandboxDefaultStatus") { - const matching = message.requestID !== undefined && message.requestID === sandboxRequest(undefined) - if (sandboxID() && !matching) return false - if (!server.isConnected()) return true - if (matching) clearSandboxRequest(undefined, message.requestID!) - const current = sandboxDefault() - if (!current || current.revision <= message.revision) { - setSandboxDefault({ - desired: message.desired, - enabled: message.enabled, - available: message.available, - reason: message.reason, - revision: message.revision, - }) - } - if (matching && !message.available) { - showToast({ - variant: "error", - title: language.t("common.requestFailed"), - description: message.reason, - }) - } - return true - } - - if (message.type === "sandboxStatus") { - const matching = message.requestID !== undefined && message.requestID === sandboxRequest(message.sessionID) - if (!server.isConnected()) return true - const current = sandboxes() - if (matching) clearSandboxRequest(message.sessionID, message.requestID!) - const next = applySandboxStates(current, message) - if (next !== current) setSandboxes(next) - const state = next[message.sessionID] - if (message.sessionID === sandboxID()) { - sandboxAttempts = 0 - if (sandboxRetry) clearTimeout(sandboxRetry) - sandboxRetry = undefined - } - if (matching && !state.available) { - showToast({ - variant: "error", - title: language.t("common.requestFailed"), - description: state.reason, - }) - } - return true - } - - if (message.type === "sandboxStatusError") { - const matching = message.requestID !== undefined && message.requestID === sandboxRequest(message.sessionID) - if (!server.isConnected()) return true - const current = sandboxes() - const state = current[message.sessionID] - if (matching) clearSandboxRequest(message.sessionID, message.requestID!) - if ((state?.revision ?? -1) > message.revision) return true - if (!message.requestID) { - const same = state?.directory === message.directory - setSandboxes( - applySandboxStates(current, { - sessionID: message.sessionID, - directory: message.directory, - enabled: same ? state.enabled : false, - available: false, - reason: message.message, - version: same ? state.version : 0, - revision: message.revision, - }), - ) - if (message.sessionID === sandboxID()) retrySandbox(message.sessionID) - } - if (matching) { - showToast({ - variant: "error", - title: language.t("common.requestFailed"), - description: message.message, - }) - } - return true - } - - if (message.type !== "configUpdated") return false - requestSandbox() - return true - } + const handleSandboxMessage = sandboxMessages({ + connected: server.isConnected, + session: sandboxID, + pending: sandboxRequest, + clear: clearSandboxRequest, + defaults: sandboxDefault, + setDefault: setSandboxDefault, + states: sandboxes, + setStates: setSandboxes, + reset: () => { + sandboxAttempts = 0 + if (sandboxRetry) clearTimeout(sandboxRetry) + sandboxRetry = undefined + }, + retry: retrySandbox, + refresh: requestSandbox, + error: (reason) => + showToast({ + variant: "error", + title: language.t("common.requestFailed"), + description: reason, + }), + }) const unsubscribe = vscode.onMessage((message) => { if (handleSandboxMessage(message)) return diff --git a/packages/kilo-vscode/webview-ui/src/components/chat/prompt-sandbox-messages.ts b/packages/kilo-vscode/webview-ui/src/components/chat/prompt-sandbox-messages.ts new file mode 100644 index 000000000000..3a4241b8865c --- /dev/null +++ b/packages/kilo-vscode/webview-ui/src/components/chat/prompt-sandbox-messages.ts @@ -0,0 +1,98 @@ +import type { + ExtensionMessage, + SandboxDefaultStatusMessage, + SandboxStatusMessage, + SandboxStatusErrorMessage, +} from "../../types/messages" +import { applySandboxStates, type SandboxDefaultState, type SandboxState } from "./prompt-input-utils" + +type Input = { + connected: () => boolean + session: () => string | null | undefined + pending: (sessionID?: string) => string | undefined + clear: (sessionID: string | undefined, requestID: string) => void + defaults: () => SandboxDefaultState | undefined + setDefault: (state: SandboxDefaultState) => void + states: () => Record + setStates: (states: Record) => void + reset: () => void + retry: (sessionID: string) => void + refresh: () => void + error: (reason: string | undefined) => void +} + +function defaults(message: SandboxDefaultStatusMessage, input: Input) { + const matching = message.requestID !== undefined && message.requestID === input.pending(undefined) + if (input.session() && !matching) return false + if (!input.connected()) return true + if (matching) input.clear(undefined, message.requestID!) + const current = input.defaults() + if (!current || current.revision <= message.revision) { + input.setDefault({ + desired: message.desired, + enabled: message.enabled, + available: message.available, + reason: message.reason, + revision: message.revision, + }) + } + if (matching && !message.available) input.error(message.reason) + return true +} + +function status(message: SandboxStatusMessage, input: Input) { + const matching = message.requestID !== undefined && message.requestID === input.pending(message.sessionID) + if (!input.connected()) return true + const current = input.states() + if (matching) input.clear(message.sessionID, message.requestID!) + const next = applySandboxStates(current, message) + if (next !== current) input.setStates(next) + const state = next[message.sessionID] + if (message.sessionID === input.session()) input.reset() + if (matching && !state.available) input.error(state.reason) + return true +} + +function failure(message: SandboxStatusErrorMessage, input: Input) { + const matching = message.requestID !== undefined && message.requestID === input.pending(message.sessionID) + if (!input.connected()) return true + const current = input.states() + const state = current[message.sessionID] + if (matching) input.clear(message.sessionID, message.requestID!) + if ((state?.revision ?? -1) > message.revision) return true + if (!message.requestID) { + const same = state?.directory === message.directory + input.setStates( + applySandboxStates(current, { + sessionID: message.sessionID, + directory: message.directory, + enabled: same ? state.enabled : false, + available: false, + reason: message.message, + version: same ? state.version : 0, + revision: message.revision, + }), + ) + if (message.sessionID === input.session()) input.retry(message.sessionID) + } + if (matching) input.error(message.message) + return true +} + +export function sandboxMessages(input: Input) { + return (message: ExtensionMessage) => { + switch (message.type) { + case "sandboxDefaultStatus": + return defaults(message, input) + case "sandboxStatus": + return status(message, input) + case "sandboxStatusError": + return failure(message, input) + case "configUpdated": + input.refresh() + return true + default: + return false + } + } +}