diff --git a/packages/kilo-vscode/src/agent-manager/worktree-importer.ts b/packages/kilo-vscode/src/agent-manager/worktree-importer.ts index d840f67eb55..f589a7c4ebe 100644 --- a/packages/kilo-vscode/src/agent-manager/worktree-importer.ts +++ b/packages/kilo-vscode/src/agent-manager/worktree-importer.ts @@ -56,63 +56,14 @@ export class WorktreeImporter { } async branch(branch: string): Promise { - 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 { + await this.run({ url }) + } + + private async run(target: { branch: string } | { url: string }): Promise { 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 } 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 53536201df3..ef69f9f3171 100644 --- a/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts +++ b/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}`) + } }) })