refactor(agent-manager): consolidate import transaction (#12651)

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

@@ -56,63 +56,14 @@ export class WorktreeImporter {
   }
 
   async branch(branch: string): Promise<void> {
-    const manager = this.host.manager()
-    const state = this.host.state()
-    if (!manager || !state) {
-      this.host.post({ type: "agentManager.importResult", success: false, message: "Not a git repository" })
-      return
-    }
-    if (this.busy()) return
-
-    this.importing = true
-    try {
-      this.host.post({
-        type: "agentManager.worktreeSetup",
-        status: "creating",
-        message: "Creating worktree from branch...",
-      })
-      const result = await manager.createWorktree({ existingBranch: branch })
-      const worktree = state.addWorktree({
-        branch: result.branch,
-        path: result.path,
-        parentBranch: result.parentBranch,
-        remote: result.remote,
-        branchOwned: false,
-      })
-      this.host.push()
-
-      try {
-        this.host.post({
-          type: "agentManager.worktreeSetup",
-          status: "creating",
-          message: "Running setup script...",
-          branch: result.branch,
-          worktreeId: worktree.id,
-        })
-        await this.host.setup(result.path, result.branch, worktree.id)
-
-        const session = await this.host.session(result.path, result.branch, worktree.id)
-        if (!session) throw new Error("Failed to create session")
-
-        state.addSession(session.id, worktree.id)
-        this.host.register(session.id, result.path)
-        this.host.ready(session.id, result, worktree.id)
-        this.host.post({ type: "agentManager.importResult", success: true, message: `Opened branch ${branch}` })
-        this.host.log(`Imported branch ${branch} as worktree ${worktree.id}`)
-      } catch (error) {
-        state.removeWorktree(worktree.id)
-        await manager.removeWorktree(result.path)
-        this.host.push()
-        throw error
-      }
-    } catch (error) {
-      this.importError(error, `Branch "${branch}" is already checked out in another worktree`)
-    } finally {
-      this.importing = false
-    }
+    await this.run({ branch })
   }
 
   async pr(url: string): Promise<void> {
+    await this.run({ url })
+  }
+
+  private async run(target: { branch: string } | { url: string }): Promise<void> {
     const manager = this.host.manager()
     const state = this.host.state()
     if (!manager || !state) {
@@ -120,11 +71,21 @@ export class WorktreeImporter {
       return
     }
     if (this.busy()) return
-
     this.importing = true
+    const branch = "branch" in target
+    const creating = branch ? "Creating worktree from branch..." : "Resolving PR..."
+    const setup = branch ? "Running setup script..." : "Setting up worktree..."
+    const duplicate = branch
+      ? `Branch "${target.branch}" is already checked out in another worktree`
+      : "This PR's branch is already checked out in another worktree"
     try {
-      this.host.post({ type: "agentManager.worktreeSetup", status: "creating", message: "Resolving PR..." })
-      const result = await manager.createFromPR(url)
+      const progress = { type: "agentManager.worktreeSetup", status: "creating" } as const
+      this.host.post({ ...progress, message: creating })
+      const result = branch
+        ? await manager.createWorktree({ existingBranch: target.branch })
+        : await manager.createFromPR(target.url)
+      const success = branch ? `Opened branch ${target.branch}` : `Opened PR branch ${result.branch}`
+      const log = branch ? `Imported branch ${target.branch}` : `Imported PR ${target.url}`
       const worktree = state.addWorktree({
         branch: result.branch,
         path: result.path,
@@ -133,29 +94,16 @@ export class WorktreeImporter {
         branchOwned: false,
       })
       this.host.push()
-
       try {
-        this.host.post({
-          type: "agentManager.worktreeSetup",
-          status: "creating",
-          message: "Setting up worktree...",
-          branch: result.branch,
-          worktreeId: worktree.id,
-        })
+        this.host.post({ ...progress, message: setup, branch: result.branch, worktreeId: worktree.id })
         await this.host.setup(result.path, result.branch, worktree.id)
-
         const session = await this.host.session(result.path, result.branch, worktree.id)
         if (!session) throw new Error("Failed to create session")
-
         state.addSession(session.id, worktree.id)
         this.host.register(session.id, result.path)
         this.host.ready(session.id, result, worktree.id)
-        this.host.post({
-          type: "agentManager.importResult",
-          success: true,
-          message: `Opened PR branch ${result.branch}`,
-        })
-        this.host.log(`Imported PR ${url} as worktree ${worktree.id}`)
+        this.host.post({ type: "agentManager.importResult", success: true, message: success })
+        this.host.log(`${log} as worktree ${worktree.id}`)
       } catch (error) {
         state.removeWorktree(worktree.id)
         await manager.removeWorktree(result.path)
@@ -163,7 +111,7 @@ export class WorktreeImporter {
         throw error
       }
     } catch (error) {
-      this.importError(error, "This PR's branch is already checked out in another worktree")
+      this.importError(error, duplicate)
     } finally {
       this.importing = false
     }

packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts

@@ -11,6 +11,7 @@ import { describe, it, expect } from "bun:test"
 import fs from "node:fs"
 import path from "node:path"
 import { Project, SyntaxKind } from "ts-morph"
+import { WorktreeImporter } from "../../src/agent-manager/worktree-importer"
 
 const ROOT = path.resolve(import.meta.dir, "../..")
 const KILO_PROVIDER_FILE = path.join(ROOT, "src/KiloProvider.ts")
@@ -578,11 +579,61 @@ describe("Agent Manager Provider — onMessage routing", () => {
 
   it("worktree import behavior lives in the cohesive importer", () => {
     const text = importer()
-    const providerText = body("onImportMessage")
-    expect(text).toContain("class WorktreeImporter")
-    expect(text).toContain("createFromPR")
-    expect(text).toContain("createWorktree")
-    expect(providerText).toContain("this.importer")
+    for (const value of ["createFromPR", "createWorktree", "this.busy()"]) expect(text).toContain(value)
+    expect(body("onImportMessage")).toContain("this.importer")
+  })
+
+  it("preserves branch and PR import ordering and rollback", async () => {
+    const run = async (kind: "branch" | "pr", fail?: "setup" | "duplicate") => {
+      const events: string[] = []
+      const create = async () => {
+        events.push("create")
+        if (fail === "duplicate") throw new Error("already checked out")
+        return { branch: "topic", path: "/repo/topic", parentBranch: "main" }
+      }
+      const importer = new WorktreeImporter({
+        manager: () =>
+          ({ createWorktree: create, createFromPR: create, removeWorktree: async () => events.push("disk") }) as never,
+        state: () =>
+          ({
+            addWorktree: (input: { branchOwned: boolean }) => ({
+              id: events.push(`add:${input.branchOwned}`) ? "worktree" : "",
+            }),
+            addSession: () => events.push("state-session"),
+            removeWorktree: () => events.push("state-remove"),
+          }) as never,
+        post: (msg) => events.push("message" in msg ? String(msg.message) : msg.type),
+        push: () => events.push("push"),
+        setup: async () => {
+          events.push("setup")
+          if (fail === "setup") throw new Error("setup failed")
+        },
+        session: async () => (events.push("session"), { id: "session" }) as never,
+        register: () => events.push("register"),
+        ready: () => events.push("ready"),
+        log: () => events.push("log"),
+      })
+      const action = () => (kind === "branch" ? importer.branch("topic") : importer.pr("https://example.test/pull/1"))
+      await action()
+      if (fail === "setup") await action()
+      return events.join("|")
+    }
+    for (const kind of ["branch", "pr"] as const) {
+      const branch = kind === "branch"
+      const creating = branch ? "Creating worktree from branch..." : "Resolving PR..."
+      const setup = branch ? "Running setup script..." : "Setting up worktree..."
+      const success = branch ? "Opened branch topic" : "Opened PR branch topic"
+      expect(await run(kind)).toBe(
+        `${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`,
+      )
+      const duplicate = branch
+        ? 'Branch "topic" is already checked out in another worktree'
+        : "This PR's branch is already checked out in another worktree"
+      expect(await run(kind, "duplicate")).toBe(`${creating}|create|${duplicate}|${duplicate}`)
+    }
   })
 })