feat(vscode): add git worktree support for agent manager sessions (#418)

* feat(vscode): add git worktree support for agent manager sessions

* test(vscode): add WorktreeManager unit and integration tests

* refactor(vscode): address review — extract branch-name, shared SessionMode type, remove premature delete handling

* refactor(vscode): pass directory through handler chain, fix listener leak, add cleanup regression tests

- Remove messageDirectory class field; pass dir as parameter to all handler methods
- Only trust directory override when onBeforeMessage interceptor is attached
- Add proper type for CreateWorktreeSessionRequest, remove `as any` cast
- Fix onMessage listener leak with onCleanup in AgentManagerApp
- Use removeWorktree() for pre-creation cleanup instead of inline fs.rm
- Add regression tests for orphaned directory removal and cleanup before re-creation

* refactor(vscode): bind directory to session not message, fix duplicate import

- KiloProvider: add sessionDirectories map and setSessionDirectory()
- getWorkspaceDirectory() resolves by sessionId lookup instead of per-message override
- Remove dir param from all 18 handler methods — they resolve directory from their sessionId
- AgentManagerProvider: call setSessionDirectory() at session creation and recovery
- Remove directory injection from message interceptor (interceptor only routes custom messages)
- Fix duplicate ExtensionMessage import in AgentManagerApp.tsx
- Forward agent param in worktree session creation

* fix: update stale docstring in AgentManagerProvider

* fix: path guard on rm, clean up sessionDirectories on delete/dispose, track recovered sessions for SSE, log metadata write failures

* fix: recover worktree sessions on restart by merging worktree directories into loadSessions

- handleLoadSessions fetches sessions from all registered worktree directories
  and merges them into the response (deduped by session ID)
- recoverWorktrees result is awaited by the interceptor before loadSessions
  passes through, ensuring sessionDirectories are populated first
- Fixes: worktree sessions disappearing after VS Code restart

* fix: non-blocking worktree recovery — refresh session list after discovery instead of blocking loadSessions

* fix: don't retry with new branch when existingBranch was requested

Author: marius-kilocode

bun.lock

@@ -325,8 +325,9 @@
     "ignore": "^7.0.3",
     "js-tiktoken": "^1.0.18",
     "lru-cache": "^11.0.2",
-    "quick-lru": "^7.0.0",
     "openai": "^4.85.4",
+    "quick-lru": "^7.0.0",
+    "simple-git": "3.31.1",
     "solid-js": "^1.9.11",
     "uri-js": "^4.4.1",
     "web-tree-sitter": "^0.24.7",

packages/kilo-vscode/package.json

@@ -23,10 +23,16 @@ export class KiloProvider implements vscode.WebviewViewProvider {
   private cachedConfigMessage: unknown = null
 
   private trackedSessionIds: Set<string> = new Set()
+  /** Per-session directory overrides (e.g., worktree paths registered by AgentManagerProvider). */
+  private sessionDirectories = new Map<string, string>()
   private unsubscribeEvent: (() => void) | null = null
   private unsubscribeState: (() => void) | null = null
   private webviewMessageDisposable: vscode.Disposable | null = null
 
+  /** Optional interceptor called before the standard message handler.
+   *  Return null to consume the message, or return a (possibly transformed) message. */
+  private onBeforeMessage: ((msg: Record<string, unknown>) => Promise<Record<string, unknown> | null>) | null = null
+
   constructor(
     private readonly extensionUri: vscode.Uri,
     private readonly connectionService: KiloConnectionService,
@@ -144,13 +150,58 @@ export class KiloProvider implements vscode.WebviewViewProvider {
     this.initializeConnection()
   }
 
+  /**
+   * Register a session created externally (e.g., worktree sessions from AgentManagerProvider).
+   * Sets currentSession, adds to trackedSessionIds, and notifies the webview.
+   */
+  public registerSession(session: SessionInfo): void {
+    this.currentSession = session
+    this.trackedSessionIds.add(session.id)
+    this.postMessage({
+      type: "sessionCreated",
+      session: this.sessionToWebview(session),
+    })
+  }
+
+  /**
+   * Add a session ID to the tracked set without changing currentSession.
+   * Used to re-register worktree sessions after clearSession wipes the set.
+   */
+  public trackSession(sessionId: string): void {
+    this.trackedSessionIds.add(sessionId)
+  }
+
+  /**
+   * Register a directory override for a session (e.g., worktree path).
+   * When set, all operations for this session use this directory instead of the workspace root.
+   */
+  public setSessionDirectory(sessionId: string, directory: string): void {
+    this.sessionDirectories.set(sessionId, directory)
+  }
+
+  /**
+   * Re-fetch and send the full session list to the webview.
+   * Called by AgentManagerProvider after worktree recovery completes.
+   */
+  public refreshSessions(): void {
+    void this.handleLoadSessions()
+  }
+
   /**
    * Attach to a webview that already has its own HTML set.
    * Sets up message handling and connection without overriding HTML content.
+   *
+   * @param options.onBeforeMessage - Optional interceptor called before the standard handler.
+   *   Return null to consume the message (stop propagation), or return the message
+   *   (possibly transformed) to continue with standard handling.
    */
-  public attachToWebview(webview: vscode.Webview): void {
+  public attachToWebview(
+    webview: vscode.Webview,
+    options?: { onBeforeMessage?: (msg: Record<string, unknown>) => Promise<Record<string, unknown> | null> },
+  ): void {
     this.isWebviewReady = false
     this.webview = webview
+    this.onBeforeMessage = options?.onBeforeMessage ?? null
     this.setupWebviewMessageHandler(webview)
     this.initializeConnection()
   }
@@ -162,6 +213,18 @@ export class KiloProvider implements vscode.WebviewViewProvider {
   private setupWebviewMessageHandler(webview: vscode.Webview): void {
     this.webviewMessageDisposable?.dispose()
     this.webviewMessageDisposable = webview.onDidReceiveMessage(async (message) => {
+      // Run interceptor if attached (e.g., AgentManagerProvider worktree logic)
+      if (this.onBeforeMessage) {
+        try {
+          const result = await this.onBeforeMessage(message)
+          if (result === null) return // consumed by interceptor
+          message = result
+        } catch (error) {
+          console.error("[Kilo New] KiloProvider: interceptor error:", error)
+          return
+        }
+      }
+
       switch (message.type) {
         case "webviewReady":
           console.log("[Kilo New] KiloProvider: ✅ webviewReady received")
@@ -472,7 +535,7 @@ export class KiloProvider implements vscode.WebviewViewProvider {
     }
 
     try {
-      const workspaceDir = this.getWorkspaceDirectory()
+      const workspaceDir = this.getWorkspaceDirectory(sessionID)
       const messagesData = await this.httpClient.getMessages(sessionID, workspaceDir)
 
       // Update currentSession so fallback logic in handleSendMessage/handleAbort
@@ -532,6 +595,26 @@ export class KiloProvider implements vscode.WebviewViewProvider {
       const workspaceDir = this.getWorkspaceDirectory()
       const sessions = await this.httpClient.listSessions(workspaceDir)
 
+      // Also fetch sessions from worktree directories so they appear in the list
+      const worktreeDirs = new Set(this.sessionDirectories.values())
+      const extra = await Promise.all(
+        [...worktreeDirs].map((dir) =>
+          this.httpClient!.listSessions(dir).catch((err) => {
+            console.error(`[Kilo New] KiloProvider: Failed to list sessions for ${dir}:`, err)
+            return [] as SessionInfo[]
+          }),
+        ),
+      )
+      const seen = new Set(sessions.map((s) => s.id))
+      for (const batch of extra) {
+        for (const s of batch) {
+          if (!seen.has(s.id)) {
+            sessions.push(s)
+            seen.add(s.id)
+          }
+        }
+      }
+
       this.postMessage({
         type: "sessionsLoaded",
         sessions: sessions.map((s) => this.sessionToWebview(s)),
@@ -555,9 +638,10 @@ export class KiloProvider implements vscode.WebviewViewProvider {
     }
 
     try {
-      const workspaceDir = this.getWorkspaceDirectory()
+      const workspaceDir = this.getWorkspaceDirectory(sessionID)
       await this.httpClient.deleteSession(sessionID, workspaceDir)
       this.trackedSessionIds.delete(sessionID)
+      this.sessionDirectories.delete(sessionID)
       if (this.currentSession?.id === sessionID) {
         this.currentSession = null
       }
@@ -581,7 +665,7 @@ export class KiloProvider implements vscode.WebviewViewProvider {
     }
 
     try {
-      const workspaceDir = this.getWorkspaceDirectory()
+      const workspaceDir = this.getWorkspaceDirectory(sessionID)
       const updated = await this.httpClient.updateSession(sessionID, { title }, workspaceDir)
       if (this.currentSession?.id === sessionID) {
         this.currentSession = updated
@@ -774,7 +858,7 @@ export class KiloProvider implements vscode.WebviewViewProvider {
     }
 
     try {
-      const workspaceDir = this.getWorkspaceDirectory()
+      const workspaceDir = this.getWorkspaceDirectory(sessionID || this.currentSession?.id)
 
       // Create session if needed
       if (!sessionID && !this.currentSession) {
@@ -842,7 +926,7 @@ export class KiloProvider implements vscode.WebviewViewProvider {
     }
 
     try {
-      const workspaceDir = this.getWorkspaceDirectory()
+      const workspaceDir = this.getWorkspaceDirectory(targetSessionID)
       await this.httpClient.abortSession(targetSessionID, workspaceDir)
     } catch (error) {
       console.error("[Kilo New] KiloProvider: Failed to abort session:", error)
@@ -877,7 +961,7 @@ export class KiloProvider implements vscode.WebviewViewProvider {
     }
 
     try {
-      const workspaceDir = this.getWorkspaceDirectory()
+      const workspaceDir = this.getWorkspaceDirectory(target)
       await this.httpClient.summarize(target, providerID, modelID, workspaceDir)
     } catch (error) {
       console.error("[Kilo New] KiloProvider: Failed to compact session:", error)
@@ -907,7 +991,7 @@ export class KiloProvider implements vscode.WebviewViewProvider {
     }
 
     try {
-      const workspaceDir = this.getWorkspaceDirectory()
+      const workspaceDir = this.getWorkspaceDirectory(targetSessionID)
       await this.httpClient.respondToPermission(targetSessionID, permissionId, response, workspaceDir)
     } catch (error) {
       console.error("[Kilo New] KiloProvider: Failed to respond to permission:", error)
@@ -924,7 +1008,7 @@ export class KiloProvider implements vscode.WebviewViewProvider {
     }
 
     try {
-      await this.httpClient.replyToQuestion(requestID, answers, this.getWorkspaceDirectory())
+      await this.httpClient.replyToQuestion(requestID, answers, this.getWorkspaceDirectory(this.currentSession?.id))
     } catch (error) {
       console.error("[Kilo New] KiloProvider: Failed to reply to question:", error)
       this.postMessage({ type: "questionError", requestID })
@@ -941,7 +1025,7 @@ export class KiloProvider implements vscode.WebviewViewProvider {
     }
 
     try {
-      await this.httpClient.rejectQuestion(requestID, this.getWorkspaceDirectory())
+      await this.httpClient.rejectQuestion(requestID, this.getWorkspaceDirectory(this.currentSession?.id))
     } catch (error) {
       console.error("[Kilo New] KiloProvider: Failed to reject question:", error)
       this.postMessage({ type: "questionError", requestID })
@@ -1329,9 +1413,14 @@ export class KiloProvider implements vscode.WebviewViewProvider {
   }
 
   /**
-   * Get the workspace directory.
+   * Get the workspace directory for a session.
+   * Checks session directory overrides first (e.g., worktree paths), then falls back to workspace root.
    */
-  private getWorkspaceDirectory(): string {
+  private getWorkspaceDirectory(sessionId?: string): string {
+    if (sessionId) {
+      const dir = this.sessionDirectories.get(sessionId)
+      if (dir) return dir
+    }
     const workspaceFolders = vscode.workspace.workspaceFolders
     if (workspaceFolders && workspaceFolders.length > 0) {
       return workspaceFolders[0].uri.fsPath
@@ -1359,5 +1448,6 @@ export class KiloProvider implements vscode.WebviewViewProvider {
     this.unsubscribeState?.()
     this.webviewMessageDisposable?.dispose()
     this.trackedSessionIds.clear()
+    this.sessionDirectories.clear()
   }
 }