feat(agent-manager): add sandbox toggle to new worktree modal

.changeset/agent-manager-worktree-sandbox.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": minor
+---
+
+Add a sandbox toggle to the Agent Manager New Worktree modal so each worktree session can start sandboxed

packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts

@@ -39,6 +39,7 @@ import { startSession } from "./mcp-warmup"
 import { readTerminalFont, watchTerminalFont } from "./terminal-font"
 import { buildKeybindingMap } from "./format-keybinding"
 import { resolveVersionModels, buildInitialMessages, type CreatedVersion } from "./multi-version"
+import { ensureSandbox } from "./sandbox-bootstrap"
 import { Semaphore } from "./semaphore"
 import { PLATFORM } from "./constants"
 import type { AgentManagerOutMessage, AgentManagerInMessage } from "./types"
@@ -850,6 +851,28 @@ export class AgentManagerProvider implements Disposable {
     }
   }
 
+  /** Remove a worktree whose session could not be safely initialized. */
+  private async discardWorktree(id: string, dir: string, branch: string, sessionId?: string): Promise<void> {
+    this.getStateManager()?.removeWorktree(id)
+    this.pushState()
+
+    if (sessionId) {
+      try {
+        await this.connectionService
+          .getClient()
+          .session.delete({ sessionID: sessionId, directory: dir }, { throwOnError: true })
+      } catch (err) {
+        this.log(`Failed to delete session ${sessionId} after worktree setup failed:`, err)
+      }
+    }
+
+    try {
+      await this.getWorktreeManager()?.removeWorktree(dir, branch)
+    } catch (err) {
+      this.log(`Failed to remove worktree ${id} after setup failed:`, err)
+    }
+  }
+
   /** Send worktreeSetup.ready + sessionMeta + pushState after worktree creation. */
   private notifyWorktreeReady(sessionId: string, result: CreateWorktreeResult, worktreeId?: string): void {
     this.pushState()
@@ -1287,6 +1310,31 @@ export class AgentManagerProvider implements Disposable {
 
       const state = this.getStateManager()!
       state.addSession(session.id, wt.worktree.id)
+
+      // Sandbox must match the user's choice before this session is exposed or
+      // receives its initial prompt. A failed reconciliation aborts this version.
+      if (msg.sandbox !== undefined) {
+        try {
+          await ensureSandbox(this.connectionService.getClient(), session.id, wt.result.path, msg.sandbox)
+        } catch (error) {
+          const err = getErrorMessage(error)
+          this.log(`Failed to configure sandbox for ${session.id}: ${err}`)
+          this.postToWebview({
+            type: "agentManager.worktreeSetup",
+            status: "error",
+            message: `Failed to configure sandbox: ${err}`,
+            worktreeId: wt.worktree.id,
+          })
+          this.host.capture("Agent Manager Session Error", {
+            source: PLATFORM,
+            error: err,
+            context: "configureSandbox",
+          })
+          await this.discardWorktree(wt.worktree.id, wt.result.path, wt.result.branch, session.id)
+          continue
+        }
+      }
+
       this.registerWorktreeSession(session.id, wt.result.path)
       this.notifyWorktreeReady(session.id, wt.result, wt.worktree.id)
 

packages/kilo-vscode/src/agent-manager/sandbox-bootstrap.ts

@@ -0,0 +1,39 @@
+import type { KiloClient } from "@kilocode/sdk/v2/client"
+import { sameDirectory } from "../kilo-provider-utils"
+
+type State = {
+  directory: string
+  enabled: boolean
+  available: boolean
+  reason?: string
+  version: number
+}
+
+function unavailable(state: State) {
+  return new Error(state.reason ?? "Sandbox backend is unavailable")
+}
+
+function routed(state: State, dir: string) {
+  if (!sameDirectory(state.directory, dir)) throw new Error("Sandbox status resolved a different directory")
+}
+
+function confirm(state: State, dir: string, desired: boolean) {
+  routed(state, dir)
+  if (desired && !state.available) throw unavailable(state)
+  if (state.enabled !== desired) {
+    throw new Error(`Sandbox remained ${state.enabled ? "enabled" : "disabled"} after reconciliation`)
+  }
+  return state
+}
+
+/** Ensure a new session uses the selected sandbox state before its first prompt. */
+export async function ensureSandbox(client: KiloClient, sid: string, dir: string, desired: boolean): Promise<State> {
+  const sandbox = client.sandbox
+  const { data: current } = await sandbox.status({ sessionID: sid, directory: dir }, { throwOnError: true })
+  routed(current, dir)
+  if (current.enabled === desired) return confirm(current, dir, desired)
+  if (!current.available) throw unavailable(current)
+
+  const { data: next } = await sandbox.toggle({ sessionID: sid, directory: dir }, { throwOnError: true })
+  return confirm(next, dir, desired)
+}

packages/kilo-vscode/src/agent-manager/types.ts

@@ -439,6 +439,8 @@ interface CreateMultiVersionIn {
   baseBranch?: string
   branchName?: string
   modelAllocations?: Array<{ providerID: string; modelID: string; count: number }>
+  /** When set, reconcile each created session's sandbox override to this state. */
+  sandbox?: boolean
 }
 
 interface RenameWorktreeIn {

packages/kilo-vscode/tests/unit/prompt-input-connection-guard.test.ts

@@ -3,8 +3,10 @@ import { readFileSync } from "node:fs"
 import { join } from "node:path"
 
 const path = join(__dirname, "..", "..", "webview-ui", "src", "components", "chat", "PromptInput.tsx")
-const src = readFileSync(path, "utf8")
+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 button = readFileSync(buttonPath, "utf8")
 const icons = readFileSync(iconPath, "utf8")
 
 describe("PromptInput connection guard", () => {
@@ -62,16 +64,19 @@ describe("PromptInput sandbox toggle", () => {
     expect(src).toContain("<Show when={sandboxVisible()}>")
     expect(src).toContain("{ action: toggleSandbox, enabled: () => sandboxVisible() && !sandboxDisabled() }")
     expect(src).toContain('if (!sandboxVisible()) hidden.add("sandbox")')
-    expect(src).toContain("onClick={toggleSandbox}")
+    expect(src).toContain("onToggle={toggleSandbox}")
     expect(src).toContain('message.type === "sandboxStatus"')
     expect(src).toContain("message.sessionID !== sandboxID() && !matching")
     expect(src).toContain("setSandboxState(state)")
     expect(src).toContain("message.requestID === sandboxRequest()")
     expect(src).toContain("const target = untrack(sandboxTarget)")
     expect(src).toContain("if (target && target !== sessionID) clearSandboxRequest()")
     expect(src).toContain("sandbox()?.enabled ?? (!sandboxID() && config().experimental?.sandbox === true)")
-    expect(src).toContain("aria-pressed={sandboxEnabled()}")
+    expect(src).toContain("<SandboxButtonBase")
+    expect(src).toContain("enabled={sandboxEnabled()}")
     expect(src).toContain("!sandboxReady()")
+    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() && target === null) return")
     expect(src).not.toContain("if (state === current) return true")
   })
@@ -91,13 +96,14 @@ describe("PromptInput sandbox toggle", () => {
     expect(src).toContain(
       "const sandboxNetworkEnabled = () => config().experimental?.sandbox_restrict_network !== false",
     )
-    expect(src).toContain('<Icon name="lock" size="small" />')
     expect(src).toContain("<SandboxTooltipContent enabled={sandboxEnabled()} network={sandboxNetworkEnabled()} />")
-    expect(src).toContain('<Icon name="folder" size="small" />')
-    expect(src).toContain('<Icon name="globe" size="small" />')
-    expect(src).toContain("props.enabled && props.network")
-    expect(src).not.toContain('class="prompt-sandbox-network"')
-    expect(src).not.toContain('class="prompt-sandbox-icon"')
+    expect(src).toContain('tooltipClass="prompt-sandbox-tooltip-content"')
+    expect(button).toContain('<Icon name="lock" size="small" />')
+    expect(button).toContain('<Icon name="folder" size="small" />')
+    expect(button).toContain('<Icon name="globe" size="small" />')
+    expect(button).toContain("props.enabled && props.network")
+    expect(button).not.toContain('class="prompt-sandbox-network"')
+    expect(button).not.toContain('class="prompt-sandbox-icon"')
     expect(icons).toContain("globe: {")
   })
 })

packages/kilo-vscode/tests/unit/sandbox-bootstrap.test.ts

@@ -0,0 +1,144 @@
+import { describe, expect, test } from "bun:test"
+import { readFileSync } from "node:fs"
+import { join } from "node:path"
+import { createKiloClient } from "@kilocode/sdk/v2/client"
+import { ensureSandbox } from "../../src/agent-manager/sandbox-bootstrap"
+
+type State = {
+  directory: string
+  enabled: boolean
+  available: boolean
+  reason?: string
+  version: number
+}
+
+function setup(states: State[]) {
+  const calls: string[] = []
+  const fetch = Object.assign(
+    async (input: RequestInfo | URL, init?: RequestInit) => {
+      const request = input instanceof Request ? input : new Request(input, init)
+      calls.push(`${request.method} ${new URL(request.url).pathname}`)
+      const state = states.shift()
+      if (!state) return Response.json({ message: "Unexpected request" }, { status: 500 })
+      return Response.json(state)
+    },
+    { preconnect: globalThis.fetch.preconnect },
+  ) satisfies typeof globalThis.fetch
+
+  return {
+    calls,
+    client: createKiloClient({ baseUrl: "http://localhost", fetch }),
+  }
+}
+
+function state(enabled: boolean, available = true, directory = "/repo"): State {
+  return { directory, enabled, available, version: 1 }
+}
+
+describe("ensureSandbox", () => {
+  test("does not toggle when the effective state already matches", async () => {
+    const ctx = setup([state(true)])
+
+    const result = await ensureSandbox(ctx.client, "session-1", "/repo", true)
+
+    expect(result.enabled).toBe(true)
+    expect(ctx.calls).toEqual(["GET /session/session-1/sandbox"])
+  })
+
+  test("toggles and verifies the selected state", async () => {
+    const ctx = setup([state(false), state(true)])
+
+    const result = await ensureSandbox(ctx.client, "session-1", "/repo", true)
+
+    expect(result.enabled).toBe(true)
+    expect(ctx.calls).toEqual(["GET /session/session-1/sandbox", "POST /session/session-1/sandbox/toggle"])
+  })
+
+  test("rejects unavailable sandboxing when sandbox was requested", async () => {
+    const unavailable = { ...state(false, false), reason: "Sandbox backend unavailable" }
+    const ctx = setup([unavailable])
+
+    expect(ensureSandbox(ctx.client, "session-1", "/repo", true)).rejects.toThrow("Sandbox backend unavailable")
+    expect(ctx.calls).toEqual(["GET /session/session-1/sandbox"])
+  })
+
+  test("allows an effectively disabled sandbox when the backend is unavailable", async () => {
+    const ctx = setup([state(false, false)])
+
+    const result = await ensureSandbox(ctx.client, "session-1", "/repo", false)
+
+    expect(result.enabled).toBe(false)
+    expect(ctx.calls).toEqual(["GET /session/session-1/sandbox"])
+  })
+
+  test("rejects a toggle that does not reach the selected state", async () => {
+    const ctx = setup([state(false), state(false)])
+
+    expect(ensureSandbox(ctx.client, "session-1", "/repo", true)).rejects.toThrow(
+      "Sandbox remained disabled after reconciliation",
+    )
+  })
+
+  test("rejects status returned for a different directory without toggling", async () => {
+    const ctx = setup([state(false, true, "/other")])
+
+    expect(ensureSandbox(ctx.client, "session-1", "/repo", true)).rejects.toThrow(
+      "Sandbox status resolved a different directory",
+    )
+    expect(ctx.calls).toEqual(["GET /session/session-1/sandbox"])
+  })
+})
+
+describe("Agent Manager sandbox startup", () => {
+  const provider = readFileSync(join(__dirname, "..", "..", "src", "agent-manager", "AgentManagerProvider.ts"), "utf8")
+  const dialog = readFileSync(
+    join(__dirname, "..", "..", "webview-ui", "agent-manager", "NewWorktreeDialog.tsx"),
+    "utf8",
+  )
+
+  test("reconciles before exposing or prompting the session", () => {
+    const start = provider.indexOf("private async onCreateMultiVersion")
+    const end = provider.indexOf("\n  private ", start + 1)
+    const body = provider.slice(start, end)
+    const ensure = body.indexOf("await ensureSandbox")
+    const discard = body.indexOf("await this.discardWorktree", ensure)
+    const skip = body.indexOf("continue", discard)
+    const register = body.indexOf("this.registerWorktreeSession", ensure)
+    const ready = body.indexOf("this.notifyWorktreeReady", register)
+    const created = body.indexOf("created.push", ready)
+    const prompt = body.indexOf('type: "agentManager.sendInitialMessage"', created)
+
+    expect(ensure).toBeGreaterThan(-1)
+    expect(discard).toBeGreaterThan(ensure)
+    expect(skip).toBeGreaterThan(discard)
+    expect(register).toBeGreaterThan(skip)
+    expect(ready).toBeGreaterThan(register)
+    expect(created).toBeGreaterThan(ready)
+    expect(prompt).toBeGreaterThan(created)
+  })
+
+  test("deletes the fresh branch when sandbox setup rolls back", () => {
+    expect(provider).toContain("private async discardWorktree(id: string, dir: string, branch: string")
+    expect(provider).toContain("removeWorktree(dir, branch)")
+    expect(provider).toContain("wt.result.path, wt.result.branch, session.id")
+  })
+
+  test("uses the experiment-aware visibility condition for UI and payload", () => {
+    expect(dialog).toContain("const sandboxVisible = () => isSandboxVisible(features(), config())")
+    expect(dialog).toContain("sandbox: sandboxVisible() ? sandbox() : undefined")
+    expect(dialog).toContain("<Show when={sandboxVisible()}>")
+  })
+
+  test("places the sandbox toggle with prompt actions instead of model selectors", () => {
+    const selectors = dialog.indexOf('<div class="prompt-input-hint-selectors">')
+    const actions = dialog.indexOf('<div class="prompt-input-hint-actions">', selectors)
+    const sandbox = dialog.indexOf("<SandboxButtonBase", actions)
+    const speech = dialog.indexOf("<SpeechToTextButton", actions)
+
+    expect(selectors).toBeGreaterThan(-1)
+    expect(actions).toBeGreaterThan(selectors)
+    expect(dialog.slice(selectors, actions)).not.toContain("<SandboxButtonBase")
+    expect(sandbox).toBeGreaterThan(actions)
+    expect(speech).toBeGreaterThan(sandbox)
+  })
+})

packages/kilo-vscode/webview-ui/agent-manager/NewWorktreeDialog.tsx

@@ -20,7 +20,9 @@ import { ModelSelectorBase } from "../src/components/shared/ModelSelector"
 import { ModeSwitcherBase } from "../src/components/shared/ModeSwitcher"
 import { SpeechToTextButton } from "../src/components/speech-to-text/SpeechToTextButton"
 import { canUseSpeechToText, selectedSpeechToTextModel } from "../src/components/speech-to-text/availability"
+import { visible as isSandboxVisible } from "../src/components/settings/sandboxing"
 import { ThinkingSelectorBase } from "../src/components/shared/ThinkingSelector"
+import { SandboxButtonBase, SandboxTooltipContent } from "../src/components/shared/SandboxButton"
 import {
   MultiModelSelector,
   type ModelAllocations,
@@ -71,7 +73,7 @@ export const NewWorktreeDialog: Component<{ onClose: () => void; defaultBaseBran
   const server = useServer()
   const session = useSession()
   const provider = useProvider()
-  const { config } = useConfig()
+  const { config, features } = useConfig()
   const metrics = tracker(vscode)
   const track = (button: string, properties?: Record<string, string | number | boolean | undefined>) =>
     metrics.track(button, "configure_worktree_dialog", properties)
@@ -102,6 +104,8 @@ export const NewWorktreeDialog: Component<{ onClose: () => void; defaultBaseBran
   const [compareOpen, setCompareOpen] = createSignal(false)
   const [highlightedIndex, setHighlightedIndex] = createSignal(0)
   const [variant, setVariant] = createSignal<string | undefined>(session.currentVariant())
+  const [sandbox, setSandbox] = createSignal(config().experimental?.sandbox === true)
+  const sandboxVisible = () => isSandboxVisible(features(), config())
   const speech = useSpeechToText(vscode, server, { t })
   const canUseSpeech = () => canUseSpeechToText(config(), provider.authStates())
   const speechModel = () => selectedSpeechToTextModel(config())
@@ -246,6 +250,7 @@ export const NewWorktreeDialog: Component<{ onClose: () => void; defaultBaseBran
       baseBranch: advanced ? (baseBranch() ?? undefined) : undefined,
       branchName: customBranch,
       modelAllocations: allocations,
+      sandbox: sandboxVisible() ? sandbox() : undefined,
       files: imgFiles,
     })
 
@@ -460,6 +465,24 @@ export const NewWorktreeDialog: Component<{ onClose: () => void; defaultBaseBran
                   </Show>
                 </div>
                 <div class="prompt-input-hint-actions">
+                  <Show when={sandboxVisible()}>
+                    <SandboxButtonBase
+                      enabled={sandbox()}
+                      tooltip={
+                        <SandboxTooltipContent
+                          enabled={sandbox()}
+                          network={config().experimental?.sandbox_restrict_network !== false}
+                        />
+                      }
+                      tooltipClass="prompt-sandbox-tooltip-content"
+                      onToggle={click(
+                        "sandbox_toggle",
+                        "configure_worktree_dialog",
+                        () => setSandbox(!sandbox()),
+                        () => ({ enabled: !sandbox() }),
+                      )}
+                    />
+                  </Show>
                   <Show when={canUseSpeech()}>
                     <SpeechToTextButton speech={speech} disabled={starting()} start={startSpeech} label={t} />
                   </Show>