fix(agent-manager): parallelize worktree terminal cleanup

.changeset/fast-worktree-terminal-cleanup.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Speed up worktree deletion when multiple terminals need to close, while preserving conversation history and process cleanup checks.

packages/kilo-vscode/src/agent-manager/pty-cleanup.ts

@@ -24,13 +24,18 @@ export async function removePtys(
   const result = await client.v2.pty.list({ location: { directory } })
   if (result.error) throw result.error
   const failed: unknown[] = []
-  for (const pty of result.data?.data ?? []) {
-    try {
-      const removed = await client.v2.pty.remove({ ptyID: pty.id, location: { directory } })
-      if (removed.error) failed.push(removed.error)
-    } catch (error) {
-      failed.push(error)
-    }
+  const ptys = result.data?.data ?? []
+  for (let index = 0; index < ptys.length; index += 4) {
+    await Promise.all(
+      ptys.slice(index, index + 4).map(async (pty) => {
+        try {
+          const removed = await client.v2.pty.remove({ ptyID: pty.id, location: { directory } })
+          if (removed.error) failed.push(removed.error)
+        } catch (error) {
+          failed.push(error)
+        }
+      }),
+    )
   }
   if (failed.length > 0) throw new AggregateError(failed, `Failed to remove PTYs in ${directory}`)
 }

packages/kilo-vscode/tests/unit/pty-cleanup.test.ts

@@ -9,26 +9,99 @@ import type { SessionTerminalManager } from "../../src/agent-manager/SessionTerm
 import type { TerminalRouter } from "../../src/agent-manager/terminal-routing"
 
 describe("Agent Manager PTY cleanup", () => {
-  it("removes every listed PTY even when one removal fails", async () => {
-    const removed: string[] = []
+  it("removes PTYs in batches of four and waits for every removal", async () => {
+    const directory = "/worktree"
+    const ptys = Array.from({ length: 9 }, (_, i) => ({
+      id: `pty-${i}`,
+      started: Promise.withResolvers<void>(),
+      finish: Promise.withResolvers<void>(),
+    }))
+    const calls: { ptyID: string; location: { directory: string } }[] = []
+    const active = new Set<string>()
+    const done: string[] = []
+    let peak = 0
     const client = {
       v2: {
         pty: {
           list: async (input: { location: { directory: string } }) => {
-            expect(input).toEqual({ location: { directory: "/worktree" } })
-            return { data: { data: [{ id: "pty-a" }, { id: "pty-b" }] } }
+            expect(input).toEqual({ location: { directory } })
+            return { data: { data: ptys.map((pty) => ({ id: pty.id })) } }
           },
           remove: async (input: { ptyID: string; location: { directory: string } }) => {
-            removed.push(input.ptyID)
-            if (input.ptyID === "pty-a") return { error: new Error("offline") }
+            const pty = ptys.find((pty) => pty.id === input.ptyID)!
+            calls.push(input)
+            active.add(pty.id)
+            peak = Math.max(peak, active.size)
+            pty.started.resolve()
+            await pty.finish.promise
+            active.delete(pty.id)
+            done.push(pty.id)
             return { data: undefined }
           },
         },
       },
     } as unknown as KiloClient
+    const task = removePtys(async (dir) => {
+      expect(dir).toBe(directory)
+      return client
+    }, directory).then(() => {
+      expect(done).toHaveLength(ptys.length)
+    })
+
+    for (const batch of [ptys.slice(0, 4), ptys.slice(4, 8), ptys.slice(8)]) {
+      await Promise.race([Promise.all(batch.map((pty) => pty.started.promise)), task])
+      expect([...active]).toEqual(batch.map((pty) => pty.id))
+      for (const pty of batch) {
+        pty.finish.resolve()
+        await pty.finish.promise
+      }
+    }
+
+    await task
+    expect(peak).toBe(4)
+    expect(active.size).toBe(0)
+    expect(calls).toEqual(ptys.map((pty) => ({ ptyID: pty.id, location: { directory } })))
+  })
+
+  it("awaits later batches before aggregating SDK errors and thrown rejections", async () => {
+    const ids = ["pty-a", "pty-b", "pty-c", "pty-d", "pty-e"]
+    const calls: string[] = []
+    const returned = new Error("offline")
+    const thrown = new Error("connection lost")
+    const started = Promise.withResolvers<void>()
+    const finish = Promise.withResolvers<void>()
+    let done = false
+    const client = {
+      v2: {
+        pty: {
+          list: async () => ({ data: { data: ids.map((id) => ({ id })) } }),
+          remove: async (input: { ptyID: string; location: { directory: string } }) => {
+            calls.push(input.ptyID)
+            if (input.ptyID === "pty-a") return { error: returned }
+            if (input.ptyID === "pty-b") throw thrown
+            if (input.ptyID === "pty-e") {
+              started.resolve()
+              await finish.promise
+              done = true
+            }
+            return { data: undefined }
+          },
+        },
+      },
+    } as unknown as KiloClient
+    const task = removePtys(async () => client, "/worktree").finally(() => {
+      expect(done).toBe(true)
+    })
+
+    await Promise.race([started.promise, task])
+    expect(calls).toEqual(ids)
+    finish.resolve()
 
-    await expect(removePtys(async () => client, "/worktree")).rejects.toThrow("Failed to remove PTYs")
-    expect(removed).toEqual(["pty-a", "pty-b"])
+    const error = await task.catch((err: unknown) => err)
+    if (!(error instanceof AggregateError)) throw new Error("Expected AggregateError", { cause: error })
+    expect(error.message).toBe("Failed to remove PTYs in /worktree")
+    expect(error.errors).toHaveLength(2)
+    expect(error.errors).toEqual(expect.arrayContaining([returned, thrown]))
   })
 
   it("propagates a list failure so callers can isolate it from disk cleanup", async () => {