From da6dcfbf59b2b2f249c8157721cb50a45ad9143e Mon Sep 17 00:00:00 2001
From: Marius <marius@kilocode.ai>
Date: Thu, 9 Apr 2026 18:39:16 +0200
Subject: [PATCH] fix(agent-manager): worktree setup overlay stuck when using
 custom base branch (#7542)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

onCreateMultiVersion was missing worktreeId in calls to
runSetupScriptForWorktree, createSessionInWorktree, and
notifyWorktreeReady. The webview tracks setup state by worktreeId, so
without it the busy overlay could never be cleared — blocking all
interaction with the worktree permanently.

Also declares the kilo-worktree-setup task type in package.json and
catches executeTask failures gracefully to prevent VS Code from surfacing
its own error notification for the undeclared task type.
---
 packages/kilo-vscode/package.json                     | 11 +++++++++++
 .../src/agent-manager/AgentManagerProvider.ts         |  7 ++++---
 packages/kilo-vscode/src/agent-manager/task-runner.ts | 11 ++++++++++-
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/packages/kilo-vscode/package.json b/packages/kilo-vscode/package.json
index d8c72aec7c8..8fcd0ce742a 100644
--- a/packages/kilo-vscode/package.json
+++ b/packages/kilo-vscode/package.json
@@ -50,6 +50,17 @@
   ],
   "main": "./dist/extension.js",
   "contributes": {
+    "taskDefinitions": [
+      {
+        "type": "kilo-worktree-setup",
+        "properties": {
+          "script": {
+            "type": "string",
+            "description": "The setup script command to execute"
+          }
+        }
+      }
+    ],
     "viewsContainers": {
       "activitybar": [
         {
diff --git a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts
index 1c6f285a539..db6f79add15 100644
--- a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts
+++ b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts
@@ -927,9 +927,9 @@ export class AgentManagerProvider implements Disposable {
         continue
       }
 
-      await this.runSetupScriptForWorktree(wt.result.path, wt.result.branch)
+      await this.runSetupScriptForWorktree(wt.result.path, wt.result.branch, wt.worktree.id)
 
-      const session = await this.createSessionInWorktree(wt.result.path, wt.result.branch)
+      const session = await this.createSessionInWorktree(wt.result.path, wt.result.branch, wt.worktree.id)
       if (!session) {
         const state = this.getStateManager()
         const manager = this.getWorktreeManager()
@@ -942,7 +942,7 @@ export class AgentManagerProvider implements Disposable {
       const state = this.getStateManager()!
       state.addSession(session.id, wt.worktree.id)
       this.registerWorktreeSession(session.id, wt.result.path)
-      this.notifyWorktreeReady(session.id, wt.result)
+      this.notifyWorktreeReady(session.id, wt.result, wt.worktree.id)
 
       // Set the per-version model immediately so the UI selector reflects
       // the correct model as soon as the worktree appears, before Phase 2.
@@ -1420,6 +1420,7 @@ export class AgentManagerProvider implements Disposable {
         status: "error",
         message: `Setup script failed: ${msg}`,
         branch,
+        worktreeId,
       })
     }
   }
diff --git a/packages/kilo-vscode/src/agent-manager/task-runner.ts b/packages/kilo-vscode/src/agent-manager/task-runner.ts
index 0c9e7fa7c46..1329939e274 100644
--- a/packages/kilo-vscode/src/agent-manager/task-runner.ts
+++ b/packages/kilo-vscode/src/agent-manager/task-runner.ts
@@ -28,7 +28,16 @@ export async function executeVscodeTask(config: SetupTaskConfig): Promise<number
     showReuseMessage: false,
   }
 
-  const execution = await vscode.tasks.executeTask(task)
+  let execution: vscode.TaskExecution
+  try {
+    execution = await vscode.tasks.executeTask(task)
+  } catch {
+    // Task type may not be registered in certain VS Code environments
+    // (e.g. remote, codespaces, or if package.json contribution is not loaded yet).
+    // Return undefined so SetupScriptRunner treats it as a non-fatal skip
+    // rather than VS Code surfacing its own error notification.
+    return undefined
+  }
 
   return new Promise((resolve, reject) => {
     let done = false