diff --git a/.changeset/agent-manager-worktree-sandbox.md b/.changeset/agent-manager-worktree-sandbox.md new file mode 100644 index 0000000000..75c4a0cab1 --- /dev/null +++ b/.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 diff --git a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts index 3e2b59e592..4f230c179b 100644 --- a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts +++ b/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 { + 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) diff --git a/packages/kilo-vscode/src/agent-manager/sandbox-bootstrap.ts b/packages/kilo-vscode/src/agent-manager/sandbox-bootstrap.ts new file mode 100644 index 0000000000..13e81e6ebc --- /dev/null +++ b/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 { + 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) +} diff --git a/packages/kilo-vscode/src/agent-manager/types.ts b/packages/kilo-vscode/src/agent-manager/types.ts index e48b9fa7a4..b4c7ec2b9e 100644 --- a/packages/kilo-vscode/src/agent-manager/types.ts +++ b/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 { 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 0450a6996a..973a4e375b 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 @@ -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,7 +64,7 @@ describe("PromptInput sandbox toggle", () => { expect(src).toContain("") 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)") @@ -70,8 +72,11 @@ describe("PromptInput sandbox toggle", () => { 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(" { expect(src).toContain( "const sandboxNetworkEnabled = () => config().experimental?.sandbox_restrict_network !== false", ) - expect(src).toContain('') expect(src).toContain("") - expect(src).toContain('') - expect(src).toContain('') - 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('') + expect(button).toContain('') + expect(button).toContain('') + 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: {") }) }) diff --git a/packages/kilo-vscode/tests/unit/sandbox-bootstrap.test.ts b/packages/kilo-vscode/tests/unit/sandbox-bootstrap.test.ts new file mode 100644 index 0000000000..136dc3da29 --- /dev/null +++ b/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("") + }) + + test("places the sandbox toggle with prompt actions instead of model selectors", () => { + const selectors = dialog.indexOf('
') + const actions = dialog.indexOf('
', selectors) + const sandbox = dialog.indexOf(" 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) => 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(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
+ + + } + tooltipClass="prompt-sandbox-tooltip-content" + onToggle={click( + "sandbox_toggle", + "configure_worktree_dialog", + () => setSandbox(!sandbox()), + () => ({ enabled: !sandbox() }), + )} + /> + 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 27e22d195e..5b97d843be 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx +++ b/packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx @@ -20,6 +20,7 @@ import { useConfig } from "../../context/config" import { useProvider } from "../../context/provider" import { ModelSelector } from "../shared/ModelSelector" import { ModeSwitcher } from "../shared/ModeSwitcher" +import { SandboxButtonBase, SandboxTooltipContent } from "../shared/SandboxButton" import { SpeechToTextButton } from "../speech-to-text/SpeechToTextButton" import { canUseSpeechToText, selectedSpeechToTextModel } from "../speech-to-text/availability" import { ThinkingSelector } from "../shared/ThinkingSelector" @@ -83,49 +84,6 @@ interface PromptInputProps { pendingSessionID?: string } -export const SandboxTooltipContent: Component<{ enabled: boolean; network: boolean }> = (props) => { - const language = useLanguage() - - return ( -
-
- {language.t(props.enabled ? "prompt.action.sandbox.status.enabled" : "prompt.action.sandbox.status.disabled")} -
-
- - {language.t("prompt.action.sandbox.filesystem")} - - {language.t( - props.enabled ? "prompt.action.sandbox.filesystem.restricted" : "prompt.action.sandbox.unrestricted", - )} - -
-
- - {language.t("prompt.action.sandbox.network")} - - {language.t( - props.enabled && props.network - ? "prompt.action.sandbox.network.blocked" - : props.enabled - ? "prompt.action.sandbox.network.allowed" - : "prompt.action.sandbox.unrestricted", - )} - -
-
- {language.t( - props.enabled - ? "prompt.action.sandbox.description.enabled" - : props.network - ? "prompt.action.sandbox.description.disabled" - : "prompt.action.sandbox.description.disabledNetworkAllowed", - )} -
-
- ) -} - export const PromptInput: Component = (props) => { const session = useSession() const server = useServer() @@ -1268,33 +1226,15 @@ export const PromptInput: Component = (props) => { - - ) - } - contentClass="prompt-sandbox-tooltip-content" - placement="top" - > - - + } + tooltipClass="prompt-sandbox-tooltip-content" + onToggle={toggleSandbox} + /> + + ) +} diff --git a/packages/kilo-vscode/webview-ui/src/stories/prompt-input.stories.tsx b/packages/kilo-vscode/webview-ui/src/stories/prompt-input.stories.tsx index 438157fa6c..cdded50649 100644 --- a/packages/kilo-vscode/webview-ui/src/stories/prompt-input.stories.tsx +++ b/packages/kilo-vscode/webview-ui/src/stories/prompt-input.stories.tsx @@ -15,7 +15,8 @@ import type { Meta, StoryObj } from "storybook-solidjs-vite" import { type ParentComponent } from "solid-js" import { StoryProviders, mockSessionValue } from "./StoryProviders" import { SessionContext } from "../context/session" -import { PromptInput, SandboxTooltipContent } from "../components/chat/PromptInput" +import { PromptInput } from "../components/chat/PromptInput" +import { SandboxTooltipContent } from "../components/shared/SandboxButton" import { Button } from "@kilocode/kilo-ui/button" import { Icon } from "@kilocode/kilo-ui/icon" import { Tooltip } from "@kilocode/kilo-ui/tooltip" diff --git a/packages/kilo-vscode/webview-ui/src/types/messages/webview-messages.ts b/packages/kilo-vscode/webview-ui/src/types/messages/webview-messages.ts index 773257d646..2b12b3ed3c 100644 --- a/packages/kilo-vscode/webview-ui/src/types/messages/webview-messages.ts +++ b/packages/kilo-vscode/webview-ui/src/types/messages/webview-messages.ts @@ -685,6 +685,9 @@ export interface CreateMultiVersionRequest { // Overrides `versions`, `providerID`, and `modelID`. variant?: string modelAllocations?: ModelAllocation[] + // When set, start each created worktree session with the sandbox override + // reconciled to this state. Only sent when sandbox controls are available. + sandbox?: boolean } // Persist tab order for a context (worktree ID or "local")