refactor(vscode): extract prompt sandbox response handling

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("<SandboxButtonBase")
@@ -108,7 +114,26 @@ describe("PromptInput sandbox toggle", () => {
     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", () => {

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<string, SandboxState>
+    requests?: Record<string, string>
+  } = {},
+) {
+  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],
+            ]
+          : [],
+      )
+    },
+  )
+})

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<PromptInputProps> = (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

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<string, SandboxState>
+  setStates: (states: Record<string, SandboxState>) => 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
+    }
+  }
+}