diff --git a/.changeset/move-worktrees-safely.md b/.changeset/move-worktrees-safely.md new file mode 100644 index 00000000000..e77ccea53c2 --- /dev/null +++ b/.changeset/move-worktrees-safely.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Clean up incomplete worktrees when moving a session fails to transfer its Git changes. diff --git a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts index 72cfd6d4cd1..096a4514ac3 100644 --- a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts +++ b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts @@ -1746,6 +1746,18 @@ export class AgentManagerProvider implements Disposable { getClient: () => this.connectionService.getClient(), createWorktreeOnDisk: (opts) => this.createWorktreeOnDisk(opts), runSetupScript: (p, b, id) => this.runSetupScriptForWorktree(p, b, id), + cleanupWorktree: async (id) => { + await this.onDeleteWorktree(id) + }, + notifyError: (error, result, id) => { + this.postToWebview({ + type: "agentManager.worktreeSetup", + status: "error", + message: error, + branch: result.branch, + worktreeId: id, + }) + }, getStateManager: () => this.getStateManager(), registerWorktreeSession: (sid, dir) => this.registerWorktreeSession(sid, dir), registerSession: (session) => this.panel?.sessions.registerSession(session), diff --git a/packages/kilo-vscode/src/agent-manager/continue-in-worktree.ts b/packages/kilo-vscode/src/agent-manager/continue-in-worktree.ts index cd49f2d2064..ef4fcdd9226 100644 --- a/packages/kilo-vscode/src/agent-manager/continue-in-worktree.ts +++ b/packages/kilo-vscode/src/agent-manager/continue-in-worktree.ts @@ -14,6 +14,8 @@ export interface ContinueContext { result: CreateWorktreeResult } | null> runSetupScript: (path: string, branch: string, worktreeId: string) => Promise + cleanupWorktree: (worktreeId: string) => Promise + notifyError: (error: string, result: CreateWorktreeResult, worktreeId: string) => void getStateManager: () => WorktreeStateManager | undefined registerWorktreeSession: (sessionId: string, directory: string) => void registerSession: (session: Session) => void @@ -72,6 +74,23 @@ export async function transferState( return { ok: true, value: undefined } } +async function rollback( + ctx: ContinueContext, + prepared: { worktreeId: string; result: CreateWorktreeResult }, + error: string, + progress: (status: string, detail?: string, error?: string) => void, +): Promise { + await ctx.cleanupWorktree(prepared.worktreeId).catch((err) => { + ctx.log("Failed to clean up worktree after continue error:", getErrorMessage(err)) + }) + try { + ctx.notifyError(error, prepared.result, prepared.worktreeId) + } catch (err) { + ctx.log("Failed to notify Agent Manager about continue error:", getErrorMessage(err)) + } + progress("error", undefined, error) +} + /** Fork the session into the worktree directory. */ export async function forkSession(ctx: ContinueContext, sessionId: string, dir: string): Promise> { let client: KiloClient @@ -136,7 +155,7 @@ export async function continueInWorktree( progress("transferring", "Transferring changes...") const transferred = await transferState(ctx, captured.value, prepared.value.result.path) - if (!transferred.ok) return progress("error", undefined, transferred.error) + if (!transferred.ok) return rollback(ctx, prepared.value, transferred.error, progress) progress("forking", "Starting session...") const forked = await forkSession(ctx, sessionId, prepared.value.result.path) diff --git a/packages/kilo-vscode/tests/unit/continue-in-worktree.test.ts b/packages/kilo-vscode/tests/unit/continue-in-worktree.test.ts index 4b72acdcac0..8cf641d90a2 100644 --- a/packages/kilo-vscode/tests/unit/continue-in-worktree.test.ts +++ b/packages/kilo-vscode/tests/unit/continue-in-worktree.test.ts @@ -1,18 +1,49 @@ -import { describe, expect, it, mock } from "bun:test" +import { afterEach, describe, expect, it, mock } from "bun:test" +import * as fs from "node:fs/promises" +import * as os from "node:os" +import * as path from "node:path" +import simpleGit from "simple-git" import { abortSession, - captureState, + continueInWorktree, forkSession, registerSession, type ContinueContext, - type StepResult, } from "../../src/agent-manager/continue-in-worktree" -import type { CreateWorktreeResult } from "../../src/agent-manager/WorktreeManager" +import { WorktreeManager, type CreateWorktreeResult } from "../../src/agent-manager/WorktreeManager" import { forkText } from "../../src/agent-manager/fork-handoff" import type { Session } from "@kilocode/sdk/v2/client" const noop = () => {} const log = noop as (...args: unknown[]) => void +const dirs: string[] = [] + +afterEach(async () => { + await Promise.all(dirs.splice(0, dirs.length).map((dir) => fs.rm(dir, { recursive: true, force: true }))) +}) + +async function repo(): Promise { + const dir = await fs.mkdtemp(path.join(os.tmpdir(), "continue-worktree-")) + dirs.push(dir) + const git = simpleGit(dir) + await git.init(["--initial-branch=main"]) + await git.addConfig("user.email", "test@test.com") + await git.addConfig("user.name", "Test") + await fs.writeFile(path.join(dir, "state.txt"), "base\n") + await git.add("state.txt") + await git.commit("initial") + return dir +} + +function client(fork = mock(async () => ({ data: session("forked") }))) { + return { + session: { + abort: mock(async () => ({})), + fork, + promptAsync: mock(async () => ({})), + }, + } as never +} function session(id: string): Session { return { @@ -36,6 +67,8 @@ function ctx(overrides: Partial = {}): ContinueContext { }, createWorktreeOnDisk: async () => null, runSetupScript: async () => {}, + cleanupWorktree: async () => {}, + notifyError: noop, getStateManager: () => undefined, registerWorktreeSession: noop, registerSession: noop, @@ -154,3 +187,47 @@ describe("continue-in-worktree steps", () => { }) }) }) + +describe("continueInWorktree", () => { + it("rolls back the created worktree when Git transfer fails", async () => { + const root = await repo() + const git = simpleGit(root) + await fs.writeFile(path.join(root, "state.txt"), "local dirty\n") + + const manager = new WorktreeManager(root, noop) + const setup = mock(async () => {}) + const cleanup = mock(async () => { + if (created) await manager.removeWorktree(created.path, created.branch) + }) + const notify = mock((_error: string, _result: CreateWorktreeResult, _worktreeId: string) => {}) + const progress: Array<{ status: string; error?: string }> = [] + let created: CreateWorktreeResult | undefined + const c = ctx({ + root, + getClient: () => client(), + createWorktreeOnDisk: async (opts) => { + const value = await manager.createWorktree(opts) + created = value + const target = simpleGit(value.path) + await fs.writeFile(path.join(value.path, "state.txt"), "conflict\n") + await target.add("state.txt") + await target.commit("conflict") + return { worktree: { id: "wt-1" }, result: value } + }, + runSetupScript: setup, + cleanupWorktree: cleanup, + notifyError: notify, + }) + + await continueInWorktree(c, "source", (status, _detail, error) => progress.push({ status, error })) + + expect(progress.at(-1)?.status).toBe("error") + expect(progress.at(-1)?.error).toContain("Unstaged patch failed") + expect(setup).toHaveBeenCalledTimes(1) + expect(cleanup).toHaveBeenCalledTimes(1) + expect(notify).toHaveBeenCalledWith(expect.stringContaining("Unstaged patch failed"), created!, "wt-1") + expect(created).toBeDefined() + await expect(fs.stat(created!.path)).rejects.toThrow() + expect((await git.branchLocal()).all).not.toContain(created!.branch) + }) +})