diff --git a/.changeset/worktree-migration.md b/.changeset/worktree-migration.md new file mode 100644 index 00000000000..5b6c21c7069 --- /dev/null +++ b/.changeset/worktree-migration.md @@ -0,0 +1,5 @@ +--- +"kilo-code": minor +--- + +Migrate worktree creation from CLI to extension for parallel mode sessions diff --git a/apps/kilocode-docs/docs/advanced-usage/agent-manager.md b/apps/kilocode-docs/docs/advanced-usage/agent-manager.md index 77c78857753..afe80f2e069 100644 --- a/apps/kilocode-docs/docs/advanced-usage/agent-manager.md +++ b/apps/kilocode-docs/docs/advanced-usage/agent-manager.md @@ -42,19 +42,43 @@ You can continue a session later (local or remote): Parallel Mode runs the agent in an isolated Git worktree branch, keeping your main branch clean. -- Enable the “Parallel Mode” toggle before starting +- Enable the "Parallel Mode" toggle before starting - The extension prevents using Parallel Mode inside an existing worktree - Open the main repository (where .git is a directory) to use this feature -- While running, the Agent Manager parses and surfaces: - - Branch name created/used - - Worktree path - - A completion/merge instruction message when the agent finishes -- After completion - - Review the branch in your VCS UI - - Merge or cherry-pick the changes as desired - - Clean up the worktree when finished - -If you need to resume with Parallel Mode later, the extension re-attaches to the same session with the same branch context. + +### Worktree Location + +Worktrees are created in `.kilocode/worktrees/` within your project directory. This folder is automatically added to `.gitignore` to prevent accidental commits. + +``` +your-project/ +├── .kilocode/ +│ └── worktrees/ +│ └── feature-branch-1234567890/ # isolated working directory +├── .gitignore # auto-updated with .kilocode/worktrees/ +└── ... +``` + +### While Running + +The Agent Manager surfaces: + +- Branch name created/used +- Worktree path +- A completion/merge instruction message when the agent finishes + +### After Completion + +- The worktree is cleaned up automatically, but the branch is preserved +- Review the branch in your VCS UI +- Merge or cherry-pick the changes as desired + +### Resuming Sessions + +If you resume a Parallel Mode session later, the extension will: + +1. Reuse the existing worktree if it still exists +2. Or recreate it from the session's branch ## Remote sessions (Cloud) diff --git a/src/core/kilocode/agent-manager/AgentManagerProvider.ts b/src/core/kilocode/agent-manager/AgentManagerProvider.ts index d0ab9f85595..b508c46e4c4 100644 --- a/src/core/kilocode/agent-manager/AgentManagerProvider.ts +++ b/src/core/kilocode/agent-manager/AgentManagerProvider.ts @@ -11,6 +11,8 @@ import { isParallelModeCompletionMessage, parseParallelModeCompletionBranch, } from "./parallelModeParser" +import { WorktreeManager, WorktreeError } from "./WorktreeManager" +import { AgentTaskRunner, AgentTasks } from "./AgentTaskRunner" import { findKilocodeCli, type CliDiscoveryResult } from "./CliPathResolver" import { canInstallCli, getCliInstallCommand, getLocalCliInstallCommand, getLocalCliBinDir } from "./CliInstaller" import { CliProcessHandler, type CliProcessHandlerCallbacks } from "./CliProcessHandler" @@ -18,7 +20,7 @@ import type { StreamEvent, KilocodeStreamEvent, KilocodePayload, WelcomeStreamEv import { extractRawText, tryParsePayloadJson } from "./askErrorParser" import { RemoteSessionService } from "./RemoteSessionService" import { KilocodeEventProcessor } from "./KilocodeEventProcessor" -import type { RemoteSession } from "./types" +import type { RemoteSession, AgentSession } from "./types" import { getUri } from "../../webview/getUri" import { getNonce } from "../../webview/getNonce" import { getViteDevServerConfig } from "../../webview/getViteDevServerConfig" @@ -67,6 +69,8 @@ export class AgentManagerProvider implements vscode.Disposable { private processStartTimes: Map = new Map() // Track currently sending message per session (for one-at-a-time constraint) private sendingMessageMap: Map = new Map() + // Worktree manager for parallel mode sessions (lazy initialized) + private worktreeManager: WorktreeManager | undefined constructor( private readonly context: vscode.ExtensionContext, @@ -244,7 +248,7 @@ export class AgentManagerProvider implements vscode.Disposable { this.stopAgentSession(message.sessionId as string) break case "agentManager.finishWorktreeSession": - this.finishWorktreeSession(message.sessionId as string) + void this.finishWorktreeSession(message.sessionId as string) break case "agentManager.sendMessage": void this.sendMessage( @@ -428,6 +432,20 @@ export class AgentManagerProvider implements vscode.Disposable { } } + /** + * Get or create WorktreeManager for the current workspace + */ + private getWorktreeManager(): WorktreeManager { + if (!this.worktreeManager) { + const workspaceFolder = vscode.workspace.workspaceFolders?.[0]?.uri.fsPath + if (!workspaceFolder) { + throw new Error("No workspace folder open") + } + this.worktreeManager = new WorktreeManager(workspaceFolder, this.outputChannel) + } + return this.worktreeManager + } + /** * Start a new agent session using the kilocode CLI * @param prompt - The task prompt for the agent @@ -469,6 +487,18 @@ export class AgentManagerProvider implements vscode.Disposable { this.postMessage({ type: "agentManager.startSessionFailed" }) } + let effectiveWorkspace = workspaceFolder + let worktreeInfo: { branch: string; path: string; parentBranch: string } | undefined + + if (options?.parallelMode && workspaceFolder) { + worktreeInfo = await this.prepareWorktreeForSession(prompt, options.existingBranch) + if (!worktreeInfo) { + onSetupFailed() + return + } + effectiveWorkspace = worktreeInfo.path + } + await this.spawnCliWithCommonSetup( prompt, { @@ -476,6 +506,8 @@ export class AgentManagerProvider implements vscode.Disposable { label: options?.labelOverride, gitUrl, existingBranch: options?.existingBranch, + worktreeInfo, + effectiveWorkspace, }, onSetupFailed, ) @@ -486,6 +518,33 @@ export class AgentManagerProvider implements vscode.Disposable { return apiConfiguration } + /** + * Creates a worktree for parallel mode sessions. + * Returns worktree info on success, or undefined if creation failed (error already shown to user). + */ + private async prepareWorktreeForSession( + prompt: string, + existingBranch?: string, + ): Promise<{ branch: string; path: string; parentBranch: string } | undefined> { + try { + const manager = this.getWorktreeManager() + const worktreeInfo = await manager.createWorktree({ prompt, existingBranch }) + this.outputChannel.appendLine( + `[AgentManager] Created worktree: ${worktreeInfo.path} (branch: ${worktreeInfo.branch})`, + ) + return worktreeInfo + } catch (error) { + const errorMsg = error instanceof Error ? error.message : String(error) + this.outputChannel.appendLine(`[AgentManager] Failed to create worktree: ${errorMsg}`) + void vscode.window.showErrorMessage( + error instanceof WorktreeError + ? `Failed to create worktree: ${error.message}` + : `Failed to start parallel mode: ${errorMsg}`, + ) + return undefined + } + } + /** * Common helper to spawn a CLI process with standard setup. * Handles CLI path lookup, workspace folder validation, API config, and event callback wiring. @@ -499,6 +558,8 @@ export class AgentManagerProvider implements vscode.Disposable { gitUrl?: string existingBranch?: string sessionId?: string + worktreeInfo?: { branch: string; path: string; parentBranch: string } + effectiveWorkspace?: string }, onSetupFailed?: () => void, ): Promise { @@ -509,6 +570,9 @@ export class AgentManagerProvider implements vscode.Disposable { return false } + // Use effective workspace (worktree path) if provided, otherwise use workspace folder + const workspace = options.effectiveWorkspace || workspaceFolder + const cliDiscovery = await findKilocodeCli((msg) => this.outputChannel.appendLine(`[AgentManager] ${msg}`)) if (!cliDiscovery) { this.outputChannel.appendLine("ERROR: kilocode CLI not found") @@ -531,9 +595,15 @@ export class AgentManagerProvider implements vscode.Disposable { this.processHandler.spawnProcess( cliDiscovery.cliPath, - workspaceFolder, + workspace, prompt, - { ...options, apiConfiguration, shellPath: cliDiscovery.shellPath }, + { + ...options, + apiConfiguration, + shellPath: cliDiscovery.shellPath, + // Pass worktree info for session state tracking + worktreeInfo: options.worktreeInfo, + }, (sid, event) => { if (!this.processStartTimes.has(sid)) { this.processStartTimes.set(sid, processStartTime) @@ -809,21 +879,22 @@ export class AgentManagerProvider implements vscode.Disposable { } /** - * Finish a worktree (parallel mode) session by gracefully terminating the CLI process. - * The CLI's SIGTERM handler will run its normal dispose flow, including worktree commit/cleanup. - * We keep the process tracked so the exit handler can mark the session as done/error. + * Finish a worktree (parallel mode) session: + * 1. Stage all changes + * 2. Ask agent to generate commit message and commit + * 3. Fallback to programmatic commit if agent times out + * 4. Terminate CLI process + * 5. Clean up worktree (keep branch) */ - private finishWorktreeSession(sessionId: string): void { + private async finishWorktreeSession(sessionId: string): Promise { const session = this.registry.getSession(sessionId) if (!session?.parallelMode?.enabled) { - // Safety: "Finish to branch" must never apply to non-worktree sessions. this.outputChannel.appendLine( `[AgentManager] Ignoring finishWorktreeSession for non-worktree session: ${sessionId}`, ) return } - // Only allow finishing if session is still running if (session.status !== "running") { this.outputChannel.appendLine( `[AgentManager] Ignoring finishWorktreeSession for non-running session: ${sessionId} (status: ${session.status})`, @@ -831,11 +902,91 @@ export class AgentManagerProvider implements vscode.Disposable { return } - this.processHandler.terminateProcess(sessionId, "SIGTERM") - this.log(sessionId, "Finishing worktree session (commit + close)...") + const worktreePath = session.parallelMode.worktreePath + const branch = session.parallelMode.branch + + if (!worktreePath) { + this.outputChannel.appendLine(`[AgentManager] No worktree path for session: ${sessionId}`) + this.processHandler.terminateProcess(sessionId, "SIGTERM") + return + } + + this.log(sessionId, "Finishing worktree session...") + + try { + const manager = this.getWorktreeManager() + + // Stage all changes + const hasChanges = await manager.stageAllChanges(worktreePath) + + if (hasChanges) { + this.log(sessionId, "Asking agent to commit changes...") + + // Create task runner with sendMessage bound to this session + const taskRunner = new AgentTaskRunner(this.outputChannel, async (sid, message) => { + await this.sendMessageToStdin(sid, message) + }) + + // Ask agent to commit with a proper message + const commitTask = AgentTasks.createCommitTask(worktreePath, "chore: parallel mode task completion") + const result = await taskRunner.executeTask(sessionId, commitTask) + + if (result.completedByAgent) { + this.log(sessionId, "Agent committed changes successfully") + } else if (result.success) { + this.log(sessionId, "Used fallback commit message") + } else { + this.log(sessionId, `Commit failed: ${result.error}`) + } + } else { + this.log(sessionId, "No changes to commit") + } + + // Terminate CLI process + this.processHandler.terminateProcess(sessionId, "SIGTERM") + + // Clean up worktree (keep the branch for later merging) + await manager.removeWorktree(worktreePath) + + // Show completion message + this.showWorktreeCompletionMessage(branch) + } catch (error) { + const errorMsg = error instanceof Error ? error.message : String(error) + this.outputChannel.appendLine(`[AgentManager] Error finishing worktree session: ${errorMsg}`) + + // Still try to terminate the process + this.processHandler.terminateProcess(sessionId, "SIGTERM") + } + this.postStateToWebview() } + /** + * Send a message to a session's stdin (for agent instructions) + */ + private async sendMessageToStdin(sessionId: string, content: string): Promise { + const message = { + type: "askResponse", + askResponse: "messageResponse", + text: content, + } + await this.processHandler.writeToStdin(sessionId, message) + } + + /** + * Show completion message after finishing worktree session + */ + private showWorktreeCompletionMessage(branch?: string): void { + if (!branch) return + + const message = `Parallel mode complete! Changes committed to: ${branch}` + void vscode.window.showInformationMessage(message, "Copy Branch Name").then((selection) => { + if (selection === "Copy Branch Name") { + void vscode.env.clipboard.writeText(branch) + } + }) + } + /** * Send a follow-up message to a running agent session via stdin. */ @@ -960,6 +1111,23 @@ export class AgentManagerProvider implements vscode.Disposable { this.outputChannel.appendLine(`[AgentManager] Resuming session ${sessionId} with new prompt`) + // Handle parallel mode session resumption + if (session.parallelMode?.enabled && session.parallelMode.branch) { + const worktreeInfo = await this.prepareWorktreeForResume(session) + if (worktreeInfo) { + await this.spawnCliWithCommonSetup(content, { + sessionId, + parallelMode: true, + gitUrl: session.gitUrl, + worktreeInfo, + effectiveWorkspace: worktreeInfo.path, + }) + return + } + // If worktree preparation failed, fall through to non-parallel mode + this.outputChannel.appendLine(`[AgentManager] Failed to prepare worktree, resuming without parallel mode`) + } + await this.spawnCliWithCommonSetup(content, { sessionId, // This triggers --session= flag parallelMode: session.parallelMode?.enabled, @@ -967,6 +1135,47 @@ export class AgentManagerProvider implements vscode.Disposable { }) } + /** + * Prepare worktree for resuming a parallel mode session. + * Uses existing worktree if available, otherwise recreates it from the session's branch. + */ + private async prepareWorktreeForResume( + session: AgentSession, + ): Promise<{ branch: string; path: string; parentBranch: string } | undefined> { + if (!session.parallelMode?.branch) { + return undefined + } + + const existingPath = session.parallelMode.worktreePath + const branch = session.parallelMode.branch + const parentBranch = session.parallelMode.parentBranch || "main" + + // Check if existing worktree is still valid + if (existingPath && fs.existsSync(existingPath)) { + const gitFile = path.join(existingPath, ".git") + if (fs.existsSync(gitFile)) { + this.outputChannel.appendLine(`[AgentManager] Reusing existing worktree at: ${existingPath}`) + return { branch, path: existingPath, parentBranch } + } + } + + // Worktree doesn't exist - recreate it from the existing branch + this.outputChannel.appendLine(`[AgentManager] Recreating worktree for branch: ${branch}`) + try { + const manager = this.getWorktreeManager() + const worktreeInfo = await manager.createWorktree({ existingBranch: branch }) + + // Update session with new worktree path + this.registry.updateParallelModeInfo(session.sessionId, { worktreePath: worktreeInfo.path }) + + return worktreeInfo + } catch (error) { + const errorMsg = error instanceof Error ? error.message : String(error) + this.outputChannel.appendLine(`[AgentManager] Failed to recreate worktree: ${errorMsg}`) + return undefined + } + } + /** * Cancel/abort a running agent session via stdin. * Falls back to SIGTERM if stdin write fails. diff --git a/src/core/kilocode/agent-manager/AgentTaskRunner.ts b/src/core/kilocode/agent-manager/AgentTaskRunner.ts new file mode 100644 index 00000000000..fcb3e822f54 --- /dev/null +++ b/src/core/kilocode/agent-manager/AgentTaskRunner.ts @@ -0,0 +1,142 @@ +/** + * AgentTaskRunner - Sends instructions to agents and waits for completion + * + * Provides a reusable pattern for agent-driven tasks like: + * - Generating commit messages and committing + * - Creating pull requests + * - Running custom workflows + */ + +import * as vscode from "vscode" +import simpleGit from "simple-git" + +export interface AgentTask { + /** Human-readable name for logging */ + name: string + /** Instruction to send to the agent */ + instruction: string + /** Check if the task completed successfully */ + checkComplete: () => Promise + /** Timeout in milliseconds */ + timeoutMs: number + /** Optional fallback if agent doesn't complete in time */ + fallback?: () => Promise +} + +export interface AgentTaskResult { + success: boolean + completedByAgent: boolean + error?: string +} + +const DEFAULT_POLL_INTERVAL_MS = 1000 + +export class AgentTaskRunner { + constructor( + private readonly outputChannel: vscode.OutputChannel, + private readonly sendMessage: (sessionId: string, message: string) => Promise, + ) {} + + /** + * Execute a task by sending instruction to agent and waiting for completion + */ + async executeTask(sessionId: string, task: AgentTask): Promise { + this.log(`Starting task: ${task.name}`) + + try { + // Send instruction to agent + await this.sendMessage(sessionId, task.instruction) + this.log(`Sent instruction to agent`) + + // Poll for completion + const completed = await this.pollForCompletion(task.checkComplete, task.timeoutMs) + + if (completed) { + this.log(`Task completed by agent: ${task.name}`) + return { success: true, completedByAgent: true } + } + + // Agent didn't complete in time + if (task.fallback) { + this.log(`Agent timed out, running fallback for: ${task.name}`) + await task.fallback() + return { success: true, completedByAgent: false } + } + + this.log(`Agent timed out, no fallback for: ${task.name}`) + return { success: false, completedByAgent: false, error: "Agent timed out" } + } catch (error) { + const errorMsg = error instanceof Error ? error.message : String(error) + this.log(`Task failed: ${task.name} - ${errorMsg}`) + return { success: false, completedByAgent: false, error: errorMsg } + } + } + + private async pollForCompletion(checkComplete: () => Promise, timeoutMs: number): Promise { + const startTime = Date.now() + + while (Date.now() - startTime < timeoutMs) { + const complete = await checkComplete() + if (complete) { + return true + } + await this.sleep(DEFAULT_POLL_INTERVAL_MS) + } + + return false + } + + private sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)) + } + + private log(message: string): void { + this.outputChannel.appendLine(`[AgentTaskRunner] ${message}`) + } +} + +/** + * Pre-built task definitions for common operations + */ +export const AgentTasks = { + /** + * Create a commit task that asks the agent to generate a proper commit message + */ + createCommitTask(worktreePath: string, fallbackMessage: string): AgentTask { + const git = simpleGit(worktreePath) + + return { + name: "commit-changes", + instruction: `Inspect the git diff and commit all staged changes with a proper conventional commit message (e.g., 'feat:', 'fix:', 'chore:', etc.). Use execute_command to run 'git diff --staged', then commit with an appropriate message using 'git commit -m "your-message"'.`, + timeoutMs: 60_000, + checkComplete: async () => { + try { + const stagedDiff = await git.diff(["--staged"]) + // If no staged changes, commit was successful + return !stagedDiff.trim() + } catch { + return false + } + }, + fallback: async () => { + await git.commit(fallbackMessage) + }, + } + }, + + /** + * Create a PR task that asks the agent to create a pull request + * (placeholder for future implementation) + */ + createPullRequestTask(worktreePath: string, baseBranch: string): AgentTask { + return { + name: "create-pull-request", + instruction: `Create a pull request from the current branch to ${baseBranch}. First, push the current branch to origin if not already pushed. Then use the GitHub CLI (gh) or API to create a PR with an appropriate title and description based on the commits.`, + timeoutMs: 60_000, + checkComplete: async () => { + // TODO: Check if PR was created via gh pr list or API + return false + }, + } + }, +} diff --git a/src/core/kilocode/agent-manager/CliArgsBuilder.ts b/src/core/kilocode/agent-manager/CliArgsBuilder.ts index f37299c2ec3..f5817ede85b 100644 --- a/src/core/kilocode/agent-manager/CliArgsBuilder.ts +++ b/src/core/kilocode/agent-manager/CliArgsBuilder.ts @@ -1,7 +1,5 @@ export interface BuildCliArgsOptions { - parallelMode?: boolean sessionId?: string - existingBranch?: string } /** @@ -15,15 +13,6 @@ export function buildCliArgs(workspace: string, prompt: string, options?: BuildC // --yolo: auto-approve tool uses (file reads, writes, commands, etc.) const args = ["--json-io", "--yolo", `--workspace=${workspace}`] - if (options?.parallelMode) { - args.push("--parallel") - - // Add existing branch flag if specified (resume on existing branch) - if (options.existingBranch) { - args.push(`--existing-branch=${options.existingBranch}`) - } - } - if (options?.sessionId) { args.push(`--session=${options.sessionId}`) } diff --git a/src/core/kilocode/agent-manager/CliOutputParser.ts b/src/core/kilocode/agent-manager/CliOutputParser.ts index d5ea2b1b186..c27f19cf0df 100644 --- a/src/core/kilocode/agent-manager/CliOutputParser.ts +++ b/src/core/kilocode/agent-manager/CliOutputParser.ts @@ -254,7 +254,11 @@ function toStreamEvent(parsed: Record): StreamEvent | null { } // Detect session_title_generated event from CLI (format: { event: "session_title_generated", sessionId: "...", title: "...", timestamp: ... }) - if (parsed.event === "session_title_generated" && typeof parsed.sessionId === "string" && typeof parsed.title === "string") { + if ( + parsed.event === "session_title_generated" && + typeof parsed.sessionId === "string" && + typeof parsed.title === "string" + ) { return { streamEventType: "session_title_generated", sessionId: parsed.sessionId as string, diff --git a/src/core/kilocode/agent-manager/CliProcessHandler.ts b/src/core/kilocode/agent-manager/CliProcessHandler.ts index 9889c696b9e..374b08e4c2f 100644 --- a/src/core/kilocode/agent-manager/CliProcessHandler.ts +++ b/src/core/kilocode/agent-manager/CliProcessHandler.ts @@ -42,6 +42,8 @@ interface PendingProcessInfo { desiredLabel?: string worktreeBranch?: string // Captured from welcome event before session_created worktreePath?: string // Captured from welcome event before session_created + /** Worktree info if created by extension (for parallel mode) - has full details upfront */ + worktreeInfo?: { branch: string; path: string; parentBranch: string } provisionalSessionId?: string // Used to show streaming content before session_created sawApiReqStarted?: boolean // Track if api_req_started arrived before session_created gitUrl?: string @@ -147,6 +149,8 @@ export class CliProcessHandler { existingBranch?: string /** Shell PATH from login shell - ensures CLI can access tools like git on macOS */ shellPath?: string + /** Worktree info if created by extension (for parallel mode) */ + worktreeInfo?: { branch: string; path: string; parentBranch: string } } | undefined, onCliEvent: (sessionId: string, event: StreamEvent) => void, @@ -181,10 +185,10 @@ export class CliProcessHandler { } // Build CLI command + // Note: Worktree/parallel mode is handled by AgentManagerProvider creating the worktree + // and passing the worktree path as the workspace. CLI is unaware of worktrees. const cliArgs = buildCliArgs(workspace, prompt, { - parallelMode: options?.parallelMode, sessionId: options?.sessionId, - existingBranch: options?.existingBranch, }) const env = this.buildEnvWithApiConfiguration(options?.apiConfiguration, options?.shellPath) @@ -242,6 +246,7 @@ export class CliProcessHandler { desiredSessionId: options?.sessionId, desiredLabel: options?.label, gitUrl: options?.gitUrl, + worktreeInfo: options?.worktreeInfo, stderrBuffer: [], stdoutBuffer: [], timeoutId: setTimeout(() => this.handlePendingTimeout(), PENDING_SESSION_TIMEOUT_MS), @@ -544,6 +549,7 @@ export class CliProcessHandler { realSessionId: string, worktreeBranch: string | undefined, worktreePath: string | undefined, + worktreeInfo: { branch: string; path: string; parentBranch: string } | undefined, parallelMode: boolean | undefined, ): void { this.debugLog(`Upgrading provisional session ${provisionalSessionId} -> ${realSessionId}`) @@ -558,7 +564,15 @@ export class CliProcessHandler { this.callbacks.onSessionRenamed?.(provisionalSessionId, realSessionId) - if (parallelMode && (worktreeBranch || worktreePath)) { + // Apply worktree info if we have it (extension created the worktree) + if (worktreeInfo && parallelMode) { + this.registry.updateParallelModeInfo(realSessionId, { + branch: worktreeInfo.branch, + worktreePath: worktreeInfo.path, + parentBranch: worktreeInfo.parentBranch, + }) + } else if (parallelMode && (worktreeBranch || worktreePath)) { + // Fallback: use branch/path from CLI welcome event this.registry.updateParallelModeInfo(realSessionId, { branch: worktreeBranch, worktreePath, @@ -622,6 +636,7 @@ export class CliProcessHandler { parallelMode, worktreeBranch, worktreePath, + worktreeInfo, desiredSessionId, desiredLabel, sawApiReqStarted, @@ -645,6 +660,7 @@ export class CliProcessHandler { sessionId, worktreeBranch, worktreePath, + worktreeInfo, parallelMode, ) return @@ -672,10 +688,18 @@ export class CliProcessHandler { this.debugLog(`Session created with CLI sessionId: ${event.sessionId}, mapped to: ${session.sessionId}`) - // Apply worktree branch if captured from welcome event - if (worktreeBranch && parallelMode) { + // Apply worktree info if we have it (extension created the worktree) + if (worktreeInfo && parallelMode) { + this.registry.updateParallelModeInfo(session.sessionId, { + branch: worktreeInfo.branch, + worktreePath: worktreeInfo.path, + parentBranch: worktreeInfo.parentBranch, + }) + this.debugLog(`Applied worktree info: ${worktreeInfo.path} (branch: ${worktreeInfo.branch})`) + } else if (worktreeBranch && parallelMode) { + // Fallback: use branch from CLI welcome event this.registry.updateParallelModeInfo(session.sessionId, { branch: worktreeBranch }) - this.debugLog(`Applied worktree branch: ${worktreeBranch}`) + this.debugLog(`Applied worktree branch from CLI: ${worktreeBranch}`) } const resolvedWorktreePath = diff --git a/src/core/kilocode/agent-manager/WorktreeManager.ts b/src/core/kilocode/agent-manager/WorktreeManager.ts new file mode 100644 index 00000000000..890b2a01811 --- /dev/null +++ b/src/core/kilocode/agent-manager/WorktreeManager.ts @@ -0,0 +1,374 @@ +/** + * WorktreeManager - Manages git worktrees for agent sessions + * + * Handles creation, discovery, commit, and cleanup of worktrees + * stored in {projectRoot}/.kilocode/worktrees/ + */ + +import * as vscode from "vscode" +import * as path from "path" +import * as fs from "fs" +import simpleGit, { SimpleGit } from "simple-git" + +export interface WorktreeInfo { + branch: string + path: string + parentBranch: string + createdAt: number +} + +export interface CreateWorktreeResult { + branch: string + path: string + parentBranch: string +} + +export interface CommitResult { + success: boolean + skipped?: boolean + reason?: string + error?: string +} + +export class WorktreeError extends Error { + constructor( + public readonly code: string, + message: string, + ) { + super(message) + this.name = "WorktreeError" + } +} + +/** + * Generate a valid git branch name from a prompt. + * Exported for testing. + */ +export function generateBranchName(prompt: string): string { + const sanitized = prompt + .slice(0, 50) + .toLowerCase() + .replace(/[^a-z0-9]+/g, "-") + .replace(/^-+|-+$/g, "") + .replace(/-+/g, "-") + + const timestamp = Date.now() + return `${sanitized || "kilo"}-${timestamp}` +} + +export class WorktreeManager { + private readonly projectRoot: string + private readonly worktreesDir: string + private readonly git: SimpleGit + private readonly outputChannel: vscode.OutputChannel + + constructor(projectRoot: string, outputChannel: vscode.OutputChannel) { + this.projectRoot = projectRoot + this.worktreesDir = path.join(projectRoot, ".kilocode", "worktrees") + this.git = simpleGit(projectRoot) + this.outputChannel = outputChannel + } + + /** + * Create a new worktree for an agent session + */ + async createWorktree(params: { prompt?: string; existingBranch?: string }): Promise { + const isRepo = await this.git.checkIsRepo() + if (!isRepo) { + throw new WorktreeError("NOT_GIT_REPO", "Workspace is not a git repository") + } + + await this.ensureWorktreesDir() + await this.ensureGitignore() + + const parentBranch = await this.getCurrentBranch() + + let branch: string + if (params.existingBranch) { + const exists = await this.branchExists(params.existingBranch) + if (!exists) { + throw new WorktreeError("BRANCH_NOT_FOUND", `Branch "${params.existingBranch}" does not exist`) + } + branch = params.existingBranch + } else { + branch = generateBranchName(params.prompt || "agent-task") + } + + let worktreePath = path.join(this.worktreesDir, branch) + + if (fs.existsSync(worktreePath)) { + this.log(`Worktree directory exists, removing: ${worktreePath}`) + await fs.promises.rm(worktreePath, { recursive: true, force: true }) + } + + try { + const args = params.existingBranch + ? ["worktree", "add", worktreePath, branch] + : ["worktree", "add", "-b", branch, worktreePath] + + await this.git.raw(args) + this.log(`Created worktree at: ${worktreePath} (branch: ${branch})`) + } catch (error) { + const errorMsg = error instanceof Error ? error.message : String(error) + + if (errorMsg.includes("already exists")) { + const newBranch = `${branch}-${Date.now()}` + worktreePath = path.join(this.worktreesDir, newBranch) + this.log(`Branch exists, retrying with: ${newBranch}`) + + await this.git.raw(["worktree", "add", "-b", newBranch, worktreePath]) + branch = newBranch + this.log(`Created worktree at: ${worktreePath} (branch: ${branch})`) + } else { + throw new WorktreeError("WORKTREE_CREATE_FAILED", `Failed to create worktree: ${errorMsg}`) + } + } + + return { branch, path: worktreePath, parentBranch } + } + + /** + * Stage all changes in a worktree. + * Returns true if there are staged changes after staging. + */ + async stageAllChanges(worktreePath: string): Promise { + const git = simpleGit(worktreePath) + + const status = await git.status() + if (status.isClean()) { + this.log("No changes to stage") + return false + } + + await git.add("-A") + + // Verify we have staged changes + const stagedDiff = await git.diff(["--staged"]) + const hasChanges = !!stagedDiff.trim() + + this.log(hasChanges ? "Changes staged successfully" : "No changes after staging") + return hasChanges + } + + /** + * Check if a worktree has staged changes + */ + async hasStagedChanges(worktreePath: string): Promise { + try { + const git = simpleGit(worktreePath) + const stagedDiff = await git.diff(["--staged"]) + return !!stagedDiff.trim() + } catch { + return false + } + } + + /** + * Commit staged changes with a message (fallback for when agent doesn't commit) + */ + async commitStagedChanges(worktreePath: string, message: string): Promise { + try { + const git = simpleGit(worktreePath) + await git.commit(message) + this.log(`Committed changes: ${message}`) + return { success: true, skipped: false } + } catch (error) { + const errorMsg = error instanceof Error ? error.message : String(error) + this.log(`Commit failed: ${errorMsg}`) + return { success: false, error: errorMsg } + } + } + + /** + * Commit all changes in a worktree (stages + commits in one step) + * Use this for programmatic commits. For agent-driven commits, use stageAllChanges + AgentTaskRunner. + */ + async commitChanges(worktreePath: string, message?: string): Promise { + try { + const hasChanges = await this.stageAllChanges(worktreePath) + if (!hasChanges) { + return { success: true, skipped: true, reason: "no_changes" } + } + + return this.commitStagedChanges(worktreePath, message || "chore: parallel mode task completion") + } catch (error) { + const errorMsg = error instanceof Error ? error.message : String(error) + this.log(`Commit failed: ${errorMsg}`) + return { success: false, error: errorMsg } + } + } + + /** + * Remove a worktree (keeps the branch) + */ + async removeWorktree(worktreePath: string): Promise { + try { + await this.git.raw(["worktree", "remove", worktreePath]) + this.log(`Removed worktree: ${worktreePath}`) + } catch (error) { + const errorMsg = error instanceof Error ? error.message : String(error) + this.log(`Warning: Failed to remove worktree: ${errorMsg}, trying force removal`) + + try { + await this.git.raw(["worktree", "remove", "--force", worktreePath]) + this.log(`Force removed worktree: ${worktreePath}`) + } catch (forceError) { + const forceErrorMsg = forceError instanceof Error ? forceError.message : String(forceError) + this.log(`Failed to force remove worktree: ${forceErrorMsg}`) + } + } + } + + /** + * Discover existing worktrees in .kilocode/worktrees/ + */ + async discoverWorktrees(): Promise { + if (!fs.existsSync(this.worktreesDir)) { + return [] + } + + const entries = await fs.promises.readdir(this.worktreesDir, { withFileTypes: true }) + const results = await Promise.all( + entries + .filter((entry) => entry.isDirectory()) + .map((entry) => this.getWorktreeInfo(path.join(this.worktreesDir, entry.name))), + ) + + return results.filter((info): info is WorktreeInfo => info !== undefined) + } + + /** + * Get info for a single worktree directory. + * Returns undefined if the directory is not a valid worktree or cannot be read. + */ + private async getWorktreeInfo(wtPath: string): Promise { + const isWorktree = await this.isValidWorktree(wtPath) + if (!isWorktree) { + return undefined + } + + try { + const git = simpleGit(wtPath) + const [branch, stat, parentBranch] = await Promise.all([ + git.revparse(["--abbrev-ref", "HEAD"]), + fs.promises.stat(wtPath), + this.getDefaultBranch(), + ]) + + return { + branch: branch.trim(), + path: wtPath, + parentBranch, + createdAt: stat.birthtimeMs, + } + } catch (error) { + this.log(`Failed to get info for worktree ${wtPath}: ${error}`) + return undefined + } + } + + /** + * Get diff between worktree HEAD and parent branch + */ + async getWorktreeDiff(worktreePath: string, parentBranch: string): Promise { + const git = simpleGit(worktreePath) + return git.diff([`${parentBranch}...HEAD`]) + } + + /** + * Ensure .kilocode/worktrees/ directory exists + */ + private async ensureWorktreesDir(): Promise { + if (!fs.existsSync(this.worktreesDir)) { + await fs.promises.mkdir(this.worktreesDir, { recursive: true }) + this.log(`Created worktrees directory: ${this.worktreesDir}`) + } + } + + /** + * Ensure .kilocode/worktrees/ is in .gitignore + */ + async ensureGitignore(): Promise { + const gitignorePath = path.join(this.projectRoot, ".gitignore") + const entry = ".kilocode/worktrees/" + + let content = "" + if (fs.existsSync(gitignorePath)) { + content = await fs.promises.readFile(gitignorePath, "utf-8") + if (content.includes(entry)) return + } + + const addition = content.endsWith("\n") || content === "" ? "" : "\n" + const gitignoreEntry = `${addition}\n# Kilo Code agent worktrees\n${entry}\n` + + await fs.promises.appendFile(gitignorePath, gitignoreEntry) + this.log("Added .kilocode/worktrees/ to .gitignore") + } + + /** + * Check if directory is a valid git worktree + */ + private async isValidWorktree(dirPath: string): Promise { + const gitFile = path.join(dirPath, ".git") + + if (!fs.existsSync(gitFile)) return false + + try { + const stat = await fs.promises.stat(gitFile) + return stat.isFile() + } catch { + return false + } + } + + /** + * Get current branch name + */ + private async getCurrentBranch(): Promise { + const branch = await this.git.revparse(["--abbrev-ref", "HEAD"]) + return branch.trim() + } + + /** + * Check if a branch exists + */ + private async branchExists(branchName: string): Promise { + try { + const branches = await this.git.branch() + return branches.all.includes(branchName) || branches.all.includes(`remotes/origin/${branchName}`) + } catch { + return false + } + } + + /** + * Get the default branch for this repository. + * Tries to detect from remote HEAD, falls back to main/master detection. + */ + private async getDefaultBranch(): Promise { + try { + // Try to get default branch from origin/HEAD + const remoteHead = await this.git.raw(["symbolic-ref", "refs/remotes/origin/HEAD"]) + const match = remoteHead.trim().match(/refs\/remotes\/origin\/(.+)$/) + if (match) { + return match[1] + } + } catch { + // Remote HEAD not available, fall back to branch detection + } + + try { + const branches = await this.git.branch() + if (branches.all.includes("main")) return "main" + if (branches.all.includes("master")) return "master" + } catch { + // Ignore branch detection errors + } + + return "main" + } + + private log(message: string): void { + this.outputChannel.appendLine(`[WorktreeManager] ${message}`) + } +} diff --git a/src/core/kilocode/agent-manager/__tests__/AgentManagerProvider.spec.ts b/src/core/kilocode/agent-manager/__tests__/AgentManagerProvider.spec.ts index 5e4444edfe3..dbc4234e8fb 100644 --- a/src/core/kilocode/agent-manager/__tests__/AgentManagerProvider.spec.ts +++ b/src/core/kilocode/agent-manager/__tests__/AgentManagerProvider.spec.ts @@ -82,6 +82,29 @@ describe("AgentManagerProvider CLI spawning", () => { getRemoteUrl: vi.fn().mockResolvedValue(undefined), })) + // Mock WorktreeManager for parallel mode tests + vi.doMock("../WorktreeManager", () => ({ + WorktreeManager: vi.fn().mockImplementation(() => ({ + createWorktree: vi.fn().mockResolvedValue({ + branch: "test-branch-123", + path: "/tmp/workspace/.kilocode/worktrees/test-branch-123", + parentBranch: "main", + }), + commitChanges: vi.fn().mockResolvedValue({ success: true }), + removeWorktree: vi.fn().mockResolvedValue(undefined), + discoverWorktrees: vi.fn().mockResolvedValue([]), + ensureGitignore: vi.fn().mockResolvedValue(undefined), + })), + WorktreeError: class WorktreeError extends Error { + constructor( + public code: string, + message: string, + ) { + super(message) + } + }, + })) + class TestProc extends EventEmitter { stdout = new EventEmitter() stderr = new EventEmitter() diff --git a/src/core/kilocode/agent-manager/__tests__/AgentTaskRunner.test.ts b/src/core/kilocode/agent-manager/__tests__/AgentTaskRunner.test.ts new file mode 100644 index 00000000000..5ddefe3f78f --- /dev/null +++ b/src/core/kilocode/agent-manager/__tests__/AgentTaskRunner.test.ts @@ -0,0 +1,170 @@ +import { describe, it, expect, vi, beforeEach } from "vitest" +import { AgentTaskRunner, AgentTasks, type AgentTask } from "../AgentTaskRunner" + +// Mock simple-git +const mockGit = { + diff: vi.fn(), + commit: vi.fn(), +} + +vi.mock("simple-git", () => ({ + default: vi.fn(() => mockGit), +})) + +// Mock vscode +vi.mock("vscode", () => ({ + default: {}, +})) + +describe("AgentTaskRunner", () => { + const mockOutputChannel = { + appendLine: vi.fn(), + } + let sendMessage: ReturnType + let runner: AgentTaskRunner + + beforeEach(() => { + vi.clearAllMocks() + sendMessage = vi.fn().mockResolvedValue(undefined) + runner = new AgentTaskRunner(mockOutputChannel as any, sendMessage) + }) + + describe("executeTask", () => { + it("sends instruction and returns success when task completes", async () => { + const task: AgentTask = { + name: "test-task", + instruction: "Do something", + timeoutMs: 100, + checkComplete: vi.fn().mockResolvedValue(true), + } + + const result = await runner.executeTask("session-1", task) + + expect(sendMessage).toHaveBeenCalledWith("session-1", "Do something") + expect(result.success).toBe(true) + expect(result.completedByAgent).toBe(true) + }) + + it("runs fallback when agent times out", async () => { + const fallback = vi.fn().mockResolvedValue(undefined) + const task: AgentTask = { + name: "test-task", + instruction: "Do something", + timeoutMs: 50, + checkComplete: vi.fn().mockResolvedValue(false), + fallback, + } + + const result = await runner.executeTask("session-1", task) + + expect(result.success).toBe(true) + expect(result.completedByAgent).toBe(false) + expect(fallback).toHaveBeenCalled() + }) + + it("returns failure when agent times out and no fallback", async () => { + const task: AgentTask = { + name: "test-task", + instruction: "Do something", + timeoutMs: 50, + checkComplete: vi.fn().mockResolvedValue(false), + } + + const result = await runner.executeTask("session-1", task) + + expect(result.success).toBe(false) + expect(result.completedByAgent).toBe(false) + expect(result.error).toBe("Agent timed out") + }) + + it("returns failure when sendMessage throws", async () => { + sendMessage.mockRejectedValue(new Error("Connection lost")) + const task: AgentTask = { + name: "test-task", + instruction: "Do something", + timeoutMs: 100, + checkComplete: vi.fn().mockResolvedValue(true), + } + + const result = await runner.executeTask("session-1", task) + + expect(result.success).toBe(false) + expect(result.error).toBe("Connection lost") + }) + + it("polls until completion within timeout", async () => { + let callCount = 0 + const task: AgentTask = { + name: "test-task", + instruction: "Do something", + timeoutMs: 5000, + checkComplete: vi.fn().mockImplementation(async () => { + callCount++ + return callCount >= 3 // Complete on 3rd check + }), + } + + const result = await runner.executeTask("session-1", task) + + expect(result.success).toBe(true) + expect(result.completedByAgent).toBe(true) + expect(callCount).toBe(3) + }) + }) +}) + +describe("AgentTasks", () => { + beforeEach(() => { + vi.clearAllMocks() + }) + + describe("createCommitTask", () => { + it("creates task with correct instruction", () => { + const task = AgentTasks.createCommitTask("/path/to/worktree", "fallback message") + + expect(task.name).toBe("commit-changes") + expect(task.instruction).toContain("conventional commit message") + expect(task.instruction).toContain("git diff --staged") + expect(task.timeoutMs).toBe(60_000) + }) + + it("checkComplete returns true when no staged changes", async () => { + mockGit.diff.mockResolvedValue("") + const task = AgentTasks.createCommitTask("/path/to/worktree", "fallback message") + + const complete = await task.checkComplete() + + expect(complete).toBe(true) + expect(mockGit.diff).toHaveBeenCalledWith(["--staged"]) + }) + + it("checkComplete returns false when staged changes exist", async () => { + mockGit.diff.mockResolvedValue("some diff content") + const task = AgentTasks.createCommitTask("/path/to/worktree", "fallback message") + + const complete = await task.checkComplete() + + expect(complete).toBe(false) + }) + + it("fallback commits with provided message", async () => { + mockGit.commit.mockResolvedValue(undefined) + const task = AgentTasks.createCommitTask("/path/to/worktree", "chore: fallback") + + await task.fallback!() + + expect(mockGit.commit).toHaveBeenCalledWith("chore: fallback") + }) + }) + + describe("createPullRequestTask", () => { + it("creates task with correct instruction", () => { + const task = AgentTasks.createPullRequestTask("/path/to/worktree", "main") + + expect(task.name).toBe("create-pull-request") + expect(task.instruction).toContain("pull request") + expect(task.instruction).toContain("main") + expect(task.timeoutMs).toBe(60_000) + }) + }) +}) diff --git a/src/core/kilocode/agent-manager/__tests__/CliArgsBuilder.spec.ts b/src/core/kilocode/agent-manager/__tests__/CliArgsBuilder.spec.ts index 4ef8eae8736..49079a3661f 100644 --- a/src/core/kilocode/agent-manager/__tests__/CliArgsBuilder.spec.ts +++ b/src/core/kilocode/agent-manager/__tests__/CliArgsBuilder.spec.ts @@ -43,10 +43,11 @@ describe("buildCliArgs", () => { expect(args[3]).toBe(prompt) }) - it("includes --parallel flag when parallelMode is true", () => { - const args = buildCliArgs("/workspace", "prompt", { parallelMode: true }) + it("does not include --parallel flag (worktree handled by extension)", () => { + // CLI is now worktree-agnostic - extension creates worktree and passes path as workspace + const args = buildCliArgs("/workspace", "prompt") - expect(args).toContain("--parallel") + expect(args).not.toContain("--parallel") }) it("includes --session flag when sessionId is provided", () => { @@ -55,20 +56,12 @@ describe("buildCliArgs", () => { expect(args).toContain("--session=abc123") }) - it("combines all options correctly", () => { + it("combines session option correctly", () => { const args = buildCliArgs("/workspace", "prompt", { - parallelMode: true, sessionId: "session-id", }) - expect(args).toEqual([ - "--json-io", - "--yolo", - "--workspace=/workspace", - "--parallel", - "--session=session-id", - "prompt", - ]) + expect(args).toEqual(["--json-io", "--yolo", "--workspace=/workspace", "--session=session-id", "prompt"]) }) it("uses --yolo for auto-approval of tool uses", () => { diff --git a/src/core/kilocode/agent-manager/__tests__/CliArgsBuilder.test.ts b/src/core/kilocode/agent-manager/__tests__/CliArgsBuilder.test.ts index 40a056910b7..2d23c6fb437 100644 --- a/src/core/kilocode/agent-manager/__tests__/CliArgsBuilder.test.ts +++ b/src/core/kilocode/agent-manager/__tests__/CliArgsBuilder.test.ts @@ -9,16 +9,11 @@ describe("CliArgsBuilder", () => { const args = buildCliArgs(workspace, prompt) expect(args).toContain("--json-io") + expect(args).toContain("--yolo") expect(args).toContain(`--workspace=${workspace}`) expect(args).toContain(prompt) }) - it("adds --parallel flag when parallelMode is true", () => { - const args = buildCliArgs(workspace, prompt, { parallelMode: true }) - - expect(args).toContain("--parallel") - }) - it("adds --session flag when sessionId is provided", () => { const args = buildCliArgs(workspace, prompt, { sessionId: "abc123" }) @@ -32,34 +27,11 @@ describe("CliArgsBuilder", () => { expect(args).toContain("--session=abc123") }) - describe("existingBranch", () => { - it("adds --existing-branch flag when parallelMode and existingBranch are set", () => { - const args = buildCliArgs(workspace, prompt, { - parallelMode: true, - existingBranch: "feature/my-branch", - }) - - expect(args).toContain("--parallel") - expect(args).toContain("--existing-branch=feature/my-branch") - }) - - it("ignores existingBranch when parallelMode is false", () => { - const args = buildCliArgs(workspace, prompt, { - parallelMode: false, - existingBranch: "feature/my-branch", - }) - - expect(args).not.toContain("--parallel") - expect(args.some((arg) => arg.includes("--existing-branch"))).toBe(false) - }) - - it("handles branch names with special characters", () => { - const args = buildCliArgs(workspace, prompt, { - parallelMode: true, - existingBranch: "feature/add-user-auth", - }) + it("does not include --parallel flag (worktree is handled by extension)", () => { + // CLI is now worktree-agnostic - extension creates worktree and passes path as workspace + const args = buildCliArgs(workspace, prompt) - expect(args).toContain("--existing-branch=feature/add-user-auth") - }) + expect(args).not.toContain("--parallel") + expect(args.some((arg) => arg.includes("--existing-branch"))).toBe(false) }) }) diff --git a/src/core/kilocode/agent-manager/__tests__/WorktreeManager.test.ts b/src/core/kilocode/agent-manager/__tests__/WorktreeManager.test.ts new file mode 100644 index 00000000000..b9a2bf088e4 --- /dev/null +++ b/src/core/kilocode/agent-manager/__tests__/WorktreeManager.test.ts @@ -0,0 +1,323 @@ +import { describe, it, expect, vi, beforeEach } from "vitest" +import * as path from "path" +import * as fs from "fs" + +// Helper to check paths cross-platform (handles both / and \ separators) +const containsPathSegments = (fullPath: string, ...segments: string[]): boolean => { + const normalized = fullPath.replace(/\\/g, "/") + return segments.every((seg) => normalized.includes(seg)) +} + +// Mock simple-git +const mockGit = { + checkIsRepo: vi.fn(), + raw: vi.fn(), + revparse: vi.fn(), + branch: vi.fn(), + status: vi.fn(), + add: vi.fn(), + commit: vi.fn(), + diff: vi.fn(), +} + +vi.mock("simple-git", () => ({ + default: vi.fn(() => mockGit), +})) + +// Mock fs using importOriginal +vi.mock("fs", async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + default: { + ...actual, + existsSync: vi.fn(), + promises: { + ...actual.promises, + mkdir: vi.fn(), + readFile: vi.fn(), + appendFile: vi.fn(), + rm: vi.fn(), + readdir: vi.fn(), + stat: vi.fn(), + }, + }, + existsSync: vi.fn(), + promises: { + mkdir: vi.fn(), + readFile: vi.fn(), + appendFile: vi.fn(), + rm: vi.fn(), + readdir: vi.fn(), + stat: vi.fn(), + }, + } +}) + +// Mock vscode +vi.mock("vscode", () => ({ + default: {}, +})) + +import { WorktreeManager, WorktreeError, generateBranchName } from "../WorktreeManager" + +describe("WorktreeManager", () => { + const projectRoot = "/test/project" + const mockOutputChannel = { + appendLine: vi.fn(), + } + + let manager: WorktreeManager + + beforeEach(() => { + vi.clearAllMocks() + manager = new WorktreeManager(projectRoot, mockOutputChannel as any) + }) + + describe("createWorktree", () => { + it("throws if workspace is not a git repo", async () => { + mockGit.checkIsRepo.mockResolvedValue(false) + + await expect(manager.createWorktree({ prompt: "test" })).rejects.toThrow(WorktreeError) + await expect(manager.createWorktree({ prompt: "test" })).rejects.toThrow("not a git repository") + }) + + it("creates worktree with generated branch name from prompt", async () => { + mockGit.checkIsRepo.mockResolvedValue(true) + mockGit.revparse.mockResolvedValue("main") + mockGit.raw.mockResolvedValue("") + vi.mocked(fs.existsSync).mockReturnValue(false) + vi.mocked(fs.promises.mkdir).mockResolvedValue(undefined) + vi.mocked(fs.promises.readFile).mockResolvedValue("") + vi.mocked(fs.promises.appendFile).mockResolvedValue(undefined) + + const result = await manager.createWorktree({ prompt: "Add authentication" }) + + expect(result.branch).toMatch(/^add-authentication-\d+$/) + expect(containsPathSegments(result.path, ".kilocode", "worktrees")).toBe(true) + expect(result.parentBranch).toBe("main") + }) + + it("uses existing branch when provided", async () => { + mockGit.checkIsRepo.mockResolvedValue(true) + mockGit.revparse.mockResolvedValue("main") + mockGit.branch.mockResolvedValue({ all: ["feature/existing-branch"] }) + mockGit.raw.mockResolvedValue("") + vi.mocked(fs.existsSync).mockReturnValue(false) + vi.mocked(fs.promises.mkdir).mockResolvedValue(undefined) + vi.mocked(fs.promises.readFile).mockResolvedValue("") + vi.mocked(fs.promises.appendFile).mockResolvedValue(undefined) + + const result = await manager.createWorktree({ existingBranch: "feature/existing-branch" }) + + expect(result.branch).toBe("feature/existing-branch") + // Check the git raw call - path may use / or \ depending on OS + expect(mockGit.raw).toHaveBeenCalledWith( + expect.arrayContaining(["worktree", "add", "feature/existing-branch"]), + ) + }) + + it("throws if existing branch does not exist", async () => { + mockGit.checkIsRepo.mockResolvedValue(true) + mockGit.revparse.mockResolvedValue("main") + mockGit.branch.mockResolvedValue({ all: [] }) + vi.mocked(fs.existsSync).mockReturnValue(false) + vi.mocked(fs.promises.mkdir).mockResolvedValue(undefined) + vi.mocked(fs.promises.readFile).mockResolvedValue("") + + await expect(manager.createWorktree({ existingBranch: "nonexistent" })).rejects.toThrow(WorktreeError) + await expect(manager.createWorktree({ existingBranch: "nonexistent" })).rejects.toThrow("does not exist") + }) + + it("removes existing worktree directory before creating", async () => { + mockGit.checkIsRepo.mockResolvedValue(true) + mockGit.revparse.mockResolvedValue("main") + mockGit.raw.mockResolvedValue("") + vi.mocked(fs.existsSync).mockImplementation((p) => { + // Return true for worktree path check, false for gitignore (cross-platform) + const normalized = String(p).replace(/\\/g, "/") + return normalized.includes(".kilocode/worktrees") + }) + vi.mocked(fs.promises.mkdir).mockResolvedValue(undefined) + vi.mocked(fs.promises.rm).mockResolvedValue(undefined) + vi.mocked(fs.promises.readFile).mockResolvedValue(".kilocode/worktrees/") + vi.mocked(fs.promises.appendFile).mockResolvedValue(undefined) + + await manager.createWorktree({ prompt: "test" }) + + // Verify rm was called - path separators vary by OS + expect(fs.promises.rm).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ recursive: true, force: true }), + ) + // Verify the path contains the expected segments + const rmCall = vi.mocked(fs.promises.rm).mock.calls[0] + expect(containsPathSegments(String(rmCall[0]), ".kilocode", "worktrees")).toBe(true) + }) + }) + + describe("stageAllChanges", () => { + it("returns false when no changes", async () => { + mockGit.status.mockResolvedValue({ isClean: () => true }) + + const result = await manager.stageAllChanges("/worktree/path") + + expect(result).toBe(false) + }) + + it("stages changes and returns true", async () => { + mockGit.status.mockResolvedValue({ isClean: () => false }) + mockGit.add.mockResolvedValue(undefined) + mockGit.diff.mockResolvedValue("some diff content") + + const result = await manager.stageAllChanges("/worktree/path") + + expect(result).toBe(true) + expect(mockGit.add).toHaveBeenCalledWith("-A") + }) + }) + + describe("commitChanges", () => { + it("skips commit when no changes", async () => { + mockGit.status.mockResolvedValue({ isClean: () => true }) + + const result = await manager.commitChanges("/worktree/path") + + expect(result.success).toBe(true) + expect(result.skipped).toBe(true) + expect(result.reason).toBe("no_changes") + }) + + it("commits all changes with default message", async () => { + mockGit.status.mockResolvedValue({ isClean: () => false }) + mockGit.add.mockResolvedValue(undefined) + mockGit.diff.mockResolvedValue("some diff content") + mockGit.commit.mockResolvedValue(undefined) + + const result = await manager.commitChanges("/worktree/path") + + expect(result.success).toBe(true) + expect(result.skipped).toBe(false) + expect(mockGit.add).toHaveBeenCalledWith("-A") + expect(mockGit.commit).toHaveBeenCalledWith("chore: parallel mode task completion") + }) + + it("uses custom commit message when provided", async () => { + mockGit.status.mockResolvedValue({ isClean: () => false }) + mockGit.add.mockResolvedValue(undefined) + mockGit.diff.mockResolvedValue("some diff content") + mockGit.commit.mockResolvedValue(undefined) + + await manager.commitChanges("/worktree/path", "feat: add auth") + + expect(mockGit.commit).toHaveBeenCalledWith("feat: add auth") + }) + + it("returns error on commit failure", async () => { + mockGit.status.mockResolvedValue({ isClean: () => false }) + mockGit.add.mockResolvedValue(undefined) + mockGit.diff.mockResolvedValue("some diff content") + mockGit.commit.mockRejectedValue(new Error("commit failed")) + + const result = await manager.commitChanges("/worktree/path") + + expect(result.success).toBe(false) + expect(result.error).toBe("commit failed") + }) + }) + + describe("removeWorktree", () => { + it("removes worktree normally", async () => { + mockGit.raw.mockResolvedValue("") + + await manager.removeWorktree("/worktree/path") + + expect(mockGit.raw).toHaveBeenCalledWith(["worktree", "remove", "/worktree/path"]) + }) + + it("force removes worktree on failure", async () => { + mockGit.raw.mockRejectedValueOnce(new Error("worktree has changes")).mockResolvedValueOnce("") + + await manager.removeWorktree("/worktree/path") + + expect(mockGit.raw).toHaveBeenCalledWith(["worktree", "remove", "--force", "/worktree/path"]) + }) + }) + + describe("discoverWorktrees", () => { + it("returns empty array when worktrees dir does not exist", async () => { + vi.mocked(fs.existsSync).mockReturnValue(false) + + const result = await manager.discoverWorktrees() + + expect(result).toEqual([]) + }) + + it("discovers valid worktrees", async () => { + vi.mocked(fs.existsSync).mockReturnValue(true) + vi.mocked(fs.promises.readdir).mockResolvedValue([ + { name: "feature-branch", isDirectory: () => true }, + ] as any) + vi.mocked(fs.promises.stat).mockResolvedValue({ + isFile: () => true, + birthtimeMs: 1234567890, + } as any) + mockGit.revparse.mockResolvedValue("feature-branch\n") + mockGit.branch.mockResolvedValue({ all: ["main"] }) + + const result = await manager.discoverWorktrees() + + expect(result).toHaveLength(1) + expect(result[0].branch).toBe("feature-branch") + }) + }) + + describe("ensureGitignore", () => { + it("adds gitignore entry when not present", async () => { + vi.mocked(fs.existsSync).mockReturnValue(true) + vi.mocked(fs.promises.readFile).mockResolvedValue("node_modules/\n") + vi.mocked(fs.promises.appendFile).mockResolvedValue(undefined) + + await manager.ensureGitignore() + + expect(fs.promises.appendFile).toHaveBeenCalledWith( + path.join(projectRoot, ".gitignore"), + expect.stringContaining(".kilocode/worktrees/"), + ) + }) + + it("skips adding entry when already present", async () => { + vi.mocked(fs.existsSync).mockReturnValue(true) + vi.mocked(fs.promises.readFile).mockResolvedValue(".kilocode/worktrees/\n") + vi.mocked(fs.promises.appendFile).mockResolvedValue(undefined) + + await manager.ensureGitignore() + + expect(fs.promises.appendFile).not.toHaveBeenCalled() + }) + }) +}) + +describe("generateBranchName", () => { + it("sanitizes prompt to valid branch name", () => { + const result = generateBranchName("Add User Authentication!!") + expect(result).toMatch(/^add-user-authentication-\d+$/) + }) + + it("truncates long prompts", () => { + const longPrompt = "A".repeat(100) + const result = generateBranchName(longPrompt) + // Branch name should be truncated (50 chars max from prompt) + expect(result.length).toBeLessThan(70) + }) + + it("handles special characters", () => { + const result = generateBranchName("Fix bug #123 & add feature!") + expect(result).toMatch(/^fix-bug-123-add-feature-\d+$/) + }) + + it("handles empty prompt", () => { + const result = generateBranchName("") + expect(result).toMatch(/^kilo-\d+$/) + }) +}) diff --git a/src/core/kilocode/agent-manager/types.ts b/src/core/kilocode/agent-manager/types.ts index 6de77b55d7d..0dccd17d006 100644 --- a/src/core/kilocode/agent-manager/types.ts +++ b/src/core/kilocode/agent-manager/types.ts @@ -13,7 +13,8 @@ export type SessionSource = "local" | "remote" export interface ParallelModeInfo { enabled: boolean branch?: string // e.g., "add-authentication-1702734891234" - worktreePath?: string // e.g., "/tmp/kilocode-worktree-add-auth..." + worktreePath?: string // e.g., ".kilocode/worktrees/add-auth..." + parentBranch?: string // e.g., "main" - the branch worktree was created from completionMessage?: string // Merge instructions from CLI on completion }