diff --git a/.changeset/worktree-pty-cleanup.md b/.changeset/worktree-pty-cleanup.md new file mode 100644 index 00000000000..6801652757d --- /dev/null +++ b/.changeset/worktree-pty-cleanup.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Prevent Agent Manager terminal processes from escaping cleanup while a worktree is deleted. diff --git a/packages/core/src/kilocode/pty/termination.ts b/packages/core/src/kilocode/pty/termination.ts index 92c0c8b9bbe..33b29936d20 100644 --- a/packages/core/src/kilocode/pty/termination.ts +++ b/packages/core/src/kilocode/pty/termination.ts @@ -151,7 +151,7 @@ export async function terminate(proc: Process, input: Runtime = runtime): Promis windowsHide: true, timeout: SPAWN_TIMEOUT_MS, }) - if (!killed && !state.exited) direct(proc) + if ((!killed || input.alive(proc.pid)) && !state.exited) direct(proc) if (!state.exited) await input.sleep(GRACE_MS) await verify(proc, state.exited, input) return diff --git a/packages/core/test/kilocode/pty-termination.test.ts b/packages/core/test/kilocode/pty-termination.test.ts index dc8997936b3..d3de5e63545 100644 --- a/packages/core/test/kilocode/pty-termination.test.ts +++ b/packages/core/test/kilocode/pty-termination.test.ts @@ -15,6 +15,7 @@ function runtime( platform: NodeJS.Platform, input: { taskkill?: boolean + taskkillLeavesAlive?: boolean signal?: "throw" tree?: Array<{ pid: number; parent: number }> } = {}, @@ -32,7 +33,7 @@ function runtime( taskkill: async (file, args, opts) => { tasks.push({ file, args, opts }) const result = input.taskkill ?? true - if (result) alive = false + if (result && !input.taskkillLeavesAlive) alive = false return result }, tree: async () => input.tree ?? [], @@ -46,7 +47,7 @@ function runtime( sleeps.push(ms) }, } - return { value, tasks, signals, sleeps } + return { value, tasks, signals, sleeps, dead: () => (alive = false) } } describe("pty process-tree termination", () => { @@ -68,6 +69,20 @@ describe("pty process-tree termination", () => { expect(input.sleeps).toEqual([200]) }) + test("falls back when Windows taskkill reports success but the PTY remains alive", async () => { + const input = runtime("win32", { taskkillLeavesAlive: true }) + const item = fake(42) + item.proc.kill = (signal) => { + item.calls.push(signal) + input.dead() + } + + await KiloPtyTermination.terminate(item.proc, input.value) + + expect(item.calls).toEqual([undefined]) + expect(input.sleeps).toEqual([200]) + }) + test("signals POSIX process groups before escalating", async () => { const item = fake(42) const input = runtime("linux") diff --git a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts index 3c4172f44f1..36a36205097 100644 --- a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts +++ b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts @@ -36,7 +36,7 @@ import { SessionTerminalManager } from "./SessionTerminalManager" import { createTerminalHost } from "./terminal-host" import { TerminalRouter } from "./terminal-routing" import { discardWorktree as discard } from "./discard-worktree" -import { removePtys as cleanupPtys } from "./pty-cleanup" +import { acquirePtyCleanup } from "./pty-cleanup" import { executeVscodeTask } from "./task-runner" import { runWorktreeSetupScript } from "./setup-script-task" import { RunController } from "./run/controller" @@ -178,7 +178,7 @@ export class AgentManagerProvider implements Disposable { push: () => this.pushState(), setup: (dir, branch, id) => this.runSetupScriptForWorktree(dir, branch, id), session: (dir, branch, id) => this.createSessionInWorktree(dir, branch, id), - removePtys: (directory) => this.removePtys(directory), + acquirePtyCleanup: (directory) => this.acquirePtyCleanup(directory), register: (sid, dir) => this.registerWorktreeSession(sid, dir), ready: (sid, result, id) => this.notifyWorktreeReady(sid, result, id), log: (...args) => this.log(...args), @@ -1014,23 +1014,20 @@ export class AgentManagerProvider implements Disposable { } } - private async removePtys(directory: string): Promise { - const release = await this.terminalRouter.blockDirectory(directory) - try { - await this.terminalRouter.closeDirectory(directory) - await cleanupPtys((dir) => this.connectionService.getClientAsync(dir), directory) - } finally { - release() - } + private async acquirePtyCleanup(directory: string): Promise<() => void> { + return acquirePtyCleanup({ + directory, + terminals: this.terminalRouter, + scripts: this.scripts.manager, + getClient: (dir) => this.connectionService.getClientAsync(dir), + }) } private async discardWorktree(id: string, dir: string, branch: string, sessionId?: string): Promise { const ctx = this.context if (!ctx) return - // The helper clears PTYs before ctx.worktreeManager().removeWorktree(dir, branch). return discard(ctx, this.lifecycleHost, id, dir, branch, sessionId) } - /** Send worktreeSetup.ready + pushState after worktree creation. */ private notifyWorktreeReady(sessionId: string, result: CreateWorktreeResult, worktreeId?: string): void { this.pushState() @@ -1094,10 +1091,14 @@ export class AgentManagerProvider implements Disposable { return true }, cleanupWorktree: async (wid, dir) => { - await this.removePtys(dir) - this.getStateManager()?.removeWorktree(wid) - await this.getWorktreeManager()?.removeWorktree(dir) - this.pushState() + const releasePtyCleanup = await this.acquirePtyCleanup(dir) + try { + await this.getWorktreeManager()?.removeWorktree(dir) + this.getStateManager()?.removeWorktree(wid) + this.pushState() + } finally { + releasePtyCleanup() + } }, setup: (dir, branch, id) => this.runSetupScriptForWorktree(dir, branch, id), createSessionInWorktree: (dir, branch, id, source) => this.createSessionInWorktree(dir, branch, id, source), @@ -1467,7 +1468,7 @@ export class AgentManagerProvider implements Disposable { capture: (event, props) => this.host.capture(event, props), autoName: () => this.host.autoBranchNaming(), client: () => this.connectionService.getClient(), - removePtys: (directory) => this.removePtys(directory), + acquirePtyCleanup: (directory) => this.acquirePtyCleanup(directory), metadata: (client, dir) => sandboxSessionMetadata(this.connectionService.sandboxPreference, client, dir), post: (msg) => this.postToWebview(msg), log: (...args) => this.log(...args), diff --git a/packages/kilo-vscode/src/agent-manager/ScriptTerminalManager.ts b/packages/kilo-vscode/src/agent-manager/ScriptTerminalManager.ts index b1c438cb120..c68355861d5 100644 --- a/packages/kilo-vscode/src/agent-manager/ScriptTerminalManager.ts +++ b/packages/kilo-vscode/src/agent-manager/ScriptTerminalManager.ts @@ -1,4 +1,5 @@ import type { KiloClient } from "@kilocode/sdk/v2/client" +import path from "node:path" import type { TerminalFont } from "./terminal-font" import type { RunHandle } from "./run/manager" @@ -86,6 +87,11 @@ function key(kind: ScriptTerminalKind, worktreeId: string, projectId?: string): return `${projectId ?? "single"}:${kind}:${worktreeId}` } +function directoryKey(directory: string) { + const value = path.resolve(directory) + return process.platform === "win32" ? value.toLowerCase() : value +} + function terminalId(): string { return `script:${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 8)}` } @@ -98,6 +104,8 @@ export class ScriptTerminalManager { private readonly entries = new Map() private readonly terminals = new Map() private readonly ptys = new Map() + private readonly creates = new Map>>() + private readonly blocked = new Map() constructor(private readonly deps: ScriptTerminalDeps) {} @@ -105,6 +113,47 @@ export class ScriptTerminalManager { kind: ScriptTerminalKind, config: ScriptTerminalConfig, done: (exit: ScriptTerminalExit) => void, + ): Promise { + const directory = directoryKey(config.cwd) + if (this.blocked.has(directory)) throw new Error(`PTY directory is being removed: ${config.cwd}`) + const task = this.startImpl(kind, config, done) + const creates = this.creates.get(directory) ?? new Set>() + if (!this.creates.has(directory)) this.creates.set(directory, creates) + creates.add(task) + try { + return await task + } finally { + creates.delete(task) + if (creates.size === 0) this.creates.delete(directory) + } + } + + async blockDirectory(directory: string) { + const target = directoryKey(directory) + this.blocked.set(target, (this.blocked.get(target) ?? 0) + 1) + const creates = this.creates.get(target) + if (creates) await Promise.allSettled([...creates]) + let released = false + return () => { + if (released) return + released = true + const count = this.blocked.get(target) + if (!count || count === 1) this.blocked.delete(target) + else this.blocked.set(target, count - 1) + } + } + + async closeDirectory(directory: string): Promise { + const target = directoryKey(directory) + const entries = [...this.entries.values()].filter((entry) => directoryKey(entry.cwd) === target) + const results = await Promise.all(entries.map((entry) => this.close(entry.terminalId, true))) + if (results.some((result) => !result)) throw new Error(`Failed to close script terminals in ${directory}`) + } + + private async startImpl( + kind: ScriptTerminalKind, + config: ScriptTerminalConfig, + done: (exit: ScriptTerminalExit) => void, ): Promise { const id = key(kind, config.worktreeId, config.projectId) const prior = this.entries.get(id) diff --git a/packages/kilo-vscode/src/agent-manager/discard-worktree.ts b/packages/kilo-vscode/src/agent-manager/discard-worktree.ts index 8d600d51a36..f37354b2241 100644 --- a/packages/kilo-vscode/src/agent-manager/discard-worktree.ts +++ b/packages/kilo-vscode/src/agent-manager/discard-worktree.ts @@ -9,28 +9,29 @@ export async function discardWorktree( branch: string, sessionId?: string, ): Promise { + let releasePtyCleanup: () => void try { - await host.removePtys(dir) + releasePtyCleanup = await host.acquirePtyCleanup(dir) } catch (error) { host.log(`Failed to remove PTYs after worktree setup failed:`, error) return } - if (sessionId) { - try { - await host.client().session.delete({ sessionID: sessionId, directory: dir }, { throwOnError: true }) - } catch (error) { - host.log(`Failed to delete session ${sessionId} after worktree setup failed:`, error) - return - } - } - try { + if (sessionId) { + try { + await host.client().session.delete({ sessionID: sessionId, directory: dir }, { throwOnError: true }) + } catch (error) { + host.log(`Failed to delete session ${sessionId} after worktree setup failed:`, error) + } + } await ctx.worktreeManager().removeWorktree(dir, branch) + ctx.peekState()?.removeWorktree(id) + host.push() } catch (error) { host.log(`Failed to remove worktree ${id} after setup failed:`, error) return + } finally { + releasePtyCleanup() } - ctx.peekState()?.removeWorktree(id) - host.push() } diff --git a/packages/kilo-vscode/src/agent-manager/provider-lifecycle.ts b/packages/kilo-vscode/src/agent-manager/provider-lifecycle.ts index ca6c79e102c..3d53cf77a73 100644 --- a/packages/kilo-vscode/src/agent-manager/provider-lifecycle.ts +++ b/packages/kilo-vscode/src/agent-manager/provider-lifecycle.ts @@ -42,7 +42,7 @@ export interface LifecycleHost { capture: (event: string, props: Record) => void autoName: () => { enabled: boolean } client: () => KiloClient - removePtys: (directory: string) => Promise + acquirePtyCleanup: (directory: string) => Promise<() => void> metadata: (client: KiloClient, dir: string) => Promise> post: (message: AgentManagerOutMessage) => void log: (...args: unknown[]) => void @@ -64,15 +64,22 @@ export async function createLifecycleWorktree( const session = await host.createSession(created.result.path, created.result.branch, created.worktree.id) if (!session) { + let releasePtyCleanup: () => void try { - await host.removePtys(created.result.path) + releasePtyCleanup = await host.acquirePtyCleanup(created.result.path) } catch (error) { host.log("Failed to remove worktree PTYs:", error) return null } - await ctx.worktreeManager().removeWorktree(created.result.path, created.result.branch) - ctx.peekState()?.removeWorktree(created.worktree.id) - host.push() + try { + await ctx.worktreeManager().removeWorktree(created.result.path, created.result.branch) + ctx.peekState()?.removeWorktree(created.worktree.id) + host.push() + } catch (error) { + host.log("Failed to remove worktree after session creation failed:", error) + } finally { + releasePtyCleanup() + } return null } @@ -117,21 +124,26 @@ export async function deleteLifecycleWorktree( return null } const branch = worktree.branchOwned === false ? undefined : (worktree.originalBranch ?? worktree.branch) + let releasePtyCleanup: () => void try { - await host.removePtys(worktree.path) - await ctx.worktreeManager().removeWorktree(worktree.path, branch) + releasePtyCleanup = await host.acquirePtyCleanup(worktree.path) } catch (error) { host.log(`Failed to remove worktree from disk: ${error}`) host.unskipStats(worktreeId) return null } - const orphaned = state.removeWorktree(worktreeId) - host.removePR(worktreeId) - host.forgetName(worktreeId) - host.stopDiffs(worktree.path, orphaned) - for (const s of orphaned) host.sessions.clearDirectory(s.id) - host.push() - host.log(`Deleted worktree ${worktreeId}${branch ? ` (${branch})` : ""}`) + try { + await ctx.worktreeManager().removeWorktree(worktree.path, branch) + const orphaned = state.removeWorktree(worktreeId) + host.removePR(worktreeId) + host.forgetName(worktreeId) + host.stopDiffs(worktree.path, orphaned) + for (const s of orphaned) host.sessions.clearDirectory(s.id) + host.push() + host.log(`Deleted worktree ${worktreeId}${branch ? ` (${branch})` : ""}`) + } finally { + releasePtyCleanup() + } return null } @@ -161,7 +173,8 @@ export async function removeStaleLifecycleWorktree( return null } try { - await host.removePtys(worktree.path) + const releasePtyCleanup = await host.acquirePtyCleanup(worktree.path) + releasePtyCleanup() } catch (error) { host.log(`Failed to remove stale worktree PTYs: ${error}`) return null diff --git a/packages/kilo-vscode/src/agent-manager/provider-multi-version.ts b/packages/kilo-vscode/src/agent-manager/provider-multi-version.ts index 693fa329d0d..2897d1c33b9 100644 --- a/packages/kilo-vscode/src/agent-manager/provider-multi-version.ts +++ b/packages/kilo-vscode/src/agent-manager/provider-multi-version.ts @@ -183,15 +183,22 @@ async function provisionVersion( const session = await host.createSession(wt.result.path, wt.result.branch, wt.worktree.id) if (!session) { + let releasePtyCleanup: () => void try { - await host.removePtys(wt.result.path) + releasePtyCleanup = await host.acquirePtyCleanup(wt.result.path) } catch (error) { host.log("Failed to remove worktree PTYs:", error) return null } - await ctx.worktreeManager().removeWorktree(wt.result.path, wt.result.branch) - ctx.peekState()?.removeWorktree(wt.worktree.id) - host.push() + try { + await ctx.worktreeManager().removeWorktree(wt.result.path, wt.result.branch) + ctx.peekState()?.removeWorktree(wt.worktree.id) + host.push() + } catch (error) { + host.log("Failed to remove worktree after session creation failed:", error) + } finally { + releasePtyCleanup() + } host.log(`Failed to create session for version ${spec.index + 1}`) return null } diff --git a/packages/kilo-vscode/src/agent-manager/pty-cleanup.ts b/packages/kilo-vscode/src/agent-manager/pty-cleanup.ts index 381a2198c00..c72fd162521 100644 --- a/packages/kilo-vscode/src/agent-manager/pty-cleanup.ts +++ b/packages/kilo-vscode/src/agent-manager/pty-cleanup.ts @@ -1,4 +1,6 @@ import type { KiloClient } from "@kilocode/sdk/v2/client" +import type { ScriptTerminalManager } from "./ScriptTerminalManager" +import type { TerminalRouter } from "./terminal-routing" export async function removePtys( getClient: (directory: string) => Promise, @@ -18,3 +20,29 @@ export async function removePtys( } if (failed.length > 0) throw new AggregateError(failed, `Failed to remove PTYs in ${directory}`) } + +export async function acquirePtyCleanup(input: { + directory: string + terminals: TerminalRouter + scripts: ScriptTerminalManager + getClient: (directory: string) => Promise +}) { + const releases = await Promise.all([ + input.terminals.blockDirectory(input.directory), + input.scripts.blockDirectory(input.directory), + ]) + try { + await input.terminals.closeDirectory(input.directory) + await input.scripts.closeDirectory(input.directory) + await removePtys(input.getClient, input.directory) + let released = false + return () => { + if (released) return + released = true + for (const release of releases) release() + } + } catch (error) { + for (const release of releases) release() + throw error + } +} diff --git a/packages/kilo-vscode/src/agent-manager/worktree-importer.ts b/packages/kilo-vscode/src/agent-manager/worktree-importer.ts index 3cfd7b4f849..7efdab1c808 100644 --- a/packages/kilo-vscode/src/agent-manager/worktree-importer.ts +++ b/packages/kilo-vscode/src/agent-manager/worktree-importer.ts @@ -12,7 +12,7 @@ export interface WorktreeImporterHost { setup(path: string, branch?: string, worktreeId?: string): Promise session(path: string, branch: string, worktreeId?: string): Promise register(sessionId: string, directory: string): void - removePtys?(directory: string): Promise + acquirePtyCleanup?(directory: string): Promise<() => void> ready(sessionId: string, result: CreateWorktreeResult, worktreeId?: string): void log(...args: unknown[]): void } @@ -110,17 +110,22 @@ export class WorktreeImporter { this.host.post({ type: "agentManager.importResult", projectId, success: true, message: success }) this.host.log(`${log} as worktree ${worktree.id}`) } catch (error) { - if (this.host.removePtys) { + let releasePtyCleanup: (() => void) | undefined + if (this.host.acquirePtyCleanup) { try { - await this.host.removePtys(result.path) + releasePtyCleanup = await this.host.acquirePtyCleanup(result.path) } catch (cleanup) { this.host.log("Failed to remove worktree PTYs:", cleanup) throw error } } - state.removeWorktree(worktree.id) - await manager.removeWorktree(result.path) - this.host.push() + try { + await manager.removeWorktree(result.path) + state.removeWorktree(worktree.id) + this.host.push() + } finally { + releasePtyCleanup?.() + } throw error } } catch (error) { diff --git a/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts b/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts index 5f0f1df1d55..8a5cb893d33 100644 --- a/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts +++ b/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts @@ -771,7 +771,7 @@ describe("Agent Manager Provider — onMessage routing", () => { `${creating}|create|add:false|push|${setup}|setup|session|state-session|register|ready|${success}|log`, ) expect(await run(kind, "setup")).toBe( - `${creating}|create|add:false|push|${setup}|setup|state-remove|disk|push|setup failed|setup failed|${creating}|create|add:false|push|${setup}|setup|state-remove|disk|push|setup failed|setup failed`, + `${creating}|create|add:false|push|${setup}|setup|disk|state-remove|push|setup failed|setup failed|${creating}|create|add:false|push|${setup}|setup|disk|state-remove|push|setup failed|setup failed`, ) const duplicate = branch ? 'Branch "topic" is already checked out in another worktree' diff --git a/packages/kilo-vscode/tests/unit/pty-cleanup.test.ts b/packages/kilo-vscode/tests/unit/pty-cleanup.test.ts index d20c71342fe..6515a5eb9a9 100644 --- a/packages/kilo-vscode/tests/unit/pty-cleanup.test.ts +++ b/packages/kilo-vscode/tests/unit/pty-cleanup.test.ts @@ -44,7 +44,7 @@ describe("Agent Manager PTY cleanup", () => { } as unknown as ProjectContext const host = { push: () => calls.push("push"), - removePtys: async () => { + acquirePtyCleanup: async () => { calls.push("pty") throw new Error("backend offline") }, @@ -54,4 +54,46 @@ describe("Agent Manager PTY cleanup", () => { await discardWorktree(ctx, host, "wt-1", "/worktree", "branch") expect(calls).toEqual(["pty", "log"]) }) + + it("keeps the cleanup gate until disk deletion completes", async () => { + const calls: string[] = [] + const release = () => calls.push("release") + const ctx = { + peekState: () => ({ removeWorktree: () => calls.push("state") }), + worktreeManager: () => ({ removeWorktree: async () => calls.push("disk") }), + } as unknown as ProjectContext + const host = { + push: () => calls.push("push"), + acquirePtyCleanup: async () => release, + client: () => ({ session: { delete: async () => undefined } }) as unknown as KiloClient, + log: () => undefined, + } as unknown as LifecycleHost + + await discardWorktree(ctx, host, "wt-1", "/worktree", "branch") + expect(calls).toEqual(["disk", "state", "push", "release"]) + }) + + it("continues disk cleanup when session deletion fails", async () => { + const calls: string[] = [] + const ctx = { + peekState: () => ({ removeWorktree: () => calls.push("state") }), + worktreeManager: () => ({ removeWorktree: async () => calls.push("disk") }), + } as unknown as ProjectContext + const host = { + push: () => calls.push("push"), + acquirePtyCleanup: async () => () => calls.push("release"), + client: () => + ({ + session: { + delete: async () => { + throw new Error("session offline") + }, + }, + }) as unknown as KiloClient, + log: () => calls.push("log"), + } as unknown as LifecycleHost + + await discardWorktree(ctx, host, "wt-1", "/worktree", "branch", "session-1") + expect(calls).toEqual(["log", "disk", "state", "push", "release"]) + }) }) diff --git a/packages/kilo-vscode/tests/unit/sandbox-bootstrap.test.ts b/packages/kilo-vscode/tests/unit/sandbox-bootstrap.test.ts index 7d7a02172d8..7ebbb61e72a 100644 --- a/packages/kilo-vscode/tests/unit/sandbox-bootstrap.test.ts +++ b/packages/kilo-vscode/tests/unit/sandbox-bootstrap.test.ts @@ -91,6 +91,7 @@ describe("ensureSandbox", () => { describe("Agent Manager sandbox startup", () => { const provider = readFileSync(join(__dirname, "..", "..", "src", "agent-manager", "AgentManagerProvider.ts"), "utf8") + const discard = readFileSync(join(__dirname, "..", "..", "src", "agent-manager", "discard-worktree.ts"), "utf8") const flow = readFileSync(join(__dirname, "..", "..", "src", "agent-manager", "provider-multi-version.ts"), "utf8") const dialog = readFileSync( join(__dirname, "..", "..", "webview-ui", "agent-manager", "NewWorktreeDialog.tsx"), @@ -129,7 +130,7 @@ describe("Agent Manager sandbox startup", () => { 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(discard).toContain("removeWorktree(dir, branch)") expect(flow).toContain("wt.result.path, wt.result.branch, sessionId") }) diff --git a/packages/kilo-vscode/tests/unit/script-terminal-manager.test.ts b/packages/kilo-vscode/tests/unit/script-terminal-manager.test.ts index c9f97d8b6ad..064711bb2b5 100644 --- a/packages/kilo-vscode/tests/unit/script-terminal-manager.test.ts +++ b/packages/kilo-vscode/tests/unit/script-terminal-manager.test.ts @@ -117,6 +117,23 @@ const config = { } describe("ScriptTerminalManager", () => { + it("blocks a script PTY create while its worktree is being removed", async () => { + const gate = deferred() + const ctx = harness({ create: async () => gate.promise }) + const started = ctx.manager.start("setup", config, () => undefined) + await wait() + + const blocked = ctx.manager.blockDirectory(config.cwd) + await expect(ctx.manager.start("run", config, () => undefined)).rejects.toThrow("directory is being removed") + gate.resolve({ data: { location: { directory: config.cwd }, data: info() } }) + await started + const release = await blocked + await ctx.manager.closeDirectory(config.cwd) + release() + + expect(ctx.calls.remove).toHaveLength(1) + }) + it("creates a Run PTY with explicit command settings and a safe snapshot", async () => { const ctx = harness() const done: unknown[] = []