fix(agent-manager): prevent PTYs escaping worktree cleanup

.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.

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

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")

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<void> {
-    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<void> {
     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),

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,13 +104,56 @@ export class ScriptTerminalManager {
   private readonly entries = new Map<string, Entry>()
   private readonly terminals = new Map<string, Entry>()
   private readonly ptys = new Map<string, Entry>()
+  private readonly creates = new Map<string, Set<Promise<unknown>>>()
+  private readonly blocked = new Map<string, number>()
 
   constructor(private readonly deps: ScriptTerminalDeps) {}
 
   async start(
     kind: ScriptTerminalKind,
     config: ScriptTerminalConfig,
     done: (exit: ScriptTerminalExit) => void,
+  ): Promise<RunHandle> {
+    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<Promise<unknown>>()
+    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<void> {
+    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<RunHandle> {
     const id = key(kind, config.worktreeId, config.projectId)
     const prior = this.entries.get(id)

packages/kilo-vscode/src/agent-manager/discard-worktree.ts

@@ -9,28 +9,29 @@ export async function discardWorktree(
   branch: string,
   sessionId?: string,
 ): Promise<void> {
+  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()
 }

packages/kilo-vscode/src/agent-manager/provider-lifecycle.ts

@@ -42,7 +42,7 @@ export interface LifecycleHost {
   capture: (event: string, props: Record<string, unknown>) => void
   autoName: () => { enabled: boolean }
   client: () => KiloClient
-  removePtys: (directory: string) => Promise<void>
+  acquirePtyCleanup: (directory: string) => Promise<() => void>
   metadata: (client: KiloClient, dir: string) => Promise<Record<string, unknown>>
   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

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
   }