fix(agent-manager): resolve live managed worktree sessions

.changeset/steady-agent-manager-ownership.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Recognize sessions discovered in managed Agent Manager worktrees during orchestration actions.

packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts

@@ -274,6 +274,7 @@ export class AgentManagerProvider implements Disposable {
       getPrs: () => this.prBridge.snapshot(),
       pushState: (ctx) => this.pushState(ctx),
       hasPanelSession: (id) => this.panelSessions.has(id),
+      routeSession: (id, dir) => this.panel?.sessions.setSessionDirectory(id, dir),
       closeSession: (id) => this.onCloseSession(id),
       postSessionClosed: (id, projectId) =>
         this.postToWebview({ type: "agentManager.sessionClosed", sessionId: id, projectId }),
@@ -300,7 +301,6 @@ export class AgentManagerProvider implements Disposable {
       (event) => this.onSessionLifecycle(event),
     )
   }
-
   /**
    * Keep each project's cached sidebar session list in sync with backend
    * session lifecycle events, so sessions created outside this panel (another

packages/kilo-vscode/src/agent-manager/WorktreeStateManager.ts

@@ -88,6 +88,7 @@ export interface ManagedSession {
 interface StateFile {
   worktrees: Record<string, Omit<Worktree, "id">>
   sessions: Record<string, Omit<ManagedSession, "id">>
+  closedSessions?: Record<string, string | null>
   sections?: Record<string, Omit<Section, "id">>
   tabOrder?: Record<string, string[]>
   worktreeOrder?: string[]
@@ -107,6 +108,7 @@ export interface StateLoadResult extends MigrationResult {
 import { KILO_DIR, migrateAgentManagerData, type MigrationResult } from "./constants"
 
 const STATE_FILE = "agent-manager.json"
+const CLOSED_LIMIT = 1_000
 
 let counter = 0
 
@@ -118,6 +120,7 @@ export class WorktreeStateManager {
   private readonly file: string
   private worktrees = new Map<string, Worktree>()
   private sessions = new Map<string, ManagedSession>()
+  private closed = new Map<string, string | null>()
   private sections = new Map<string, Section>()
   private tabOrder: Record<string, string[]> = {}
   private worktreeOrder: string[] = []
@@ -172,6 +175,10 @@ export class WorktreeStateManager {
     return this.sessions.get(id)
   }
 
+  isSessionClosed(id: string): boolean {
+    return this.closed.has(id)
+  }
+
   /** Returns the worktree directory for a session, or undefined for local sessions. */
   directoryFor(sessionId: string): string | undefined {
     const session = this.sessions.get(sessionId)
@@ -328,6 +335,10 @@ export class WorktreeStateManager {
       }
     }
 
+    for (const [session, worktree] of this.closed) {
+      if (worktree === id) this.closed.delete(session)
+    }
+
     // Clean up tab order for this worktree
     delete this.tabOrder[id]
 
@@ -339,6 +350,7 @@ export class WorktreeStateManager {
   }
 
   addSession(sessionId: string, worktreeId: string | null): ManagedSession {
+    this.closed.delete(sessionId)
     const session: ManagedSession = { id: sessionId, worktreeId, createdAt: new Date().toISOString() }
     this.sessions.set(sessionId, session)
     const worktree = worktreeId ? this.worktrees.get(worktreeId) : undefined
@@ -370,6 +382,13 @@ export class WorktreeStateManager {
     void this.save()
   }
 
+  closeSession(id: string, worktreeId: string | null): void {
+    this.closed.delete(id)
+    this.closed.set(id, worktreeId)
+    if (this.closed.size > CLOSED_LIMIT) this.closed.delete(this.closed.keys().next().value!)
+    void this.save()
+  }
+
   removeSession(id: string): void {
     this.sessions.delete(id)
 
@@ -709,6 +728,7 @@ export class WorktreeStateManager {
     const data = JSON.parse(content) as StateFile
     this.worktrees.clear()
     this.sessions.clear()
+    this.closed.clear()
     this.sections.clear()
     this.tabOrder = {}
     this.worktreeOrder = []
@@ -737,6 +757,7 @@ export class WorktreeStateManager {
       }
       this.sessions.set(id, session)
     }
+    this.restoreClosed(data.closedSessions)
     for (const [id, sec] of Object.entries(data.sections ?? {})) {
       this.sections.set(id, { id, ...sec })
     }
@@ -762,6 +783,13 @@ export class WorktreeStateManager {
     }
   }
 
+  private restoreClosed(value: StateFile["closedSessions"]): void {
+    if (!value || typeof value !== "object" || Array.isArray(value)) return
+    for (const [id, ref] of Object.entries(value)) {
+      if (ref === null || (typeof ref === "string" && this.worktrees.has(ref))) this.closed.set(id, ref)
+    }
+  }
+
   /** Remove worktrees whose directories no longer exist on disk and prune orphaned sessions. */
   async validate(root: string): Promise<void> {
     let changed = false
@@ -840,6 +868,7 @@ export class WorktreeStateManager {
       const { id: _, ...rest } = s
       data.sessions[id] = rest
     }
+    if (this.closed.size > 0) data.closedSessions = Object.fromEntries(this.closed)
     if (this.sections.size > 0) {
       data.sections = {}
       for (const [id, sec] of this.sections) {

packages/kilo-vscode/src/agent-manager/orchestration-bridge.ts

@@ -4,7 +4,7 @@ import type { SSEPayload } from "../services/cli-backend/sdk-sse-adapter"
 import { sameDirectory } from "../kilo-provider-utils"
 import type { LocalStats, WorktreeStats } from "./GitStatsPoller"
 import type { PRStatus } from "./types"
-import type { WorktreeStateManager } from "./WorktreeStateManager"
+import type { ManagedSession, WorktreeStateManager } from "./WorktreeStateManager"
 import {
   OrchestrationError,
   answer,
@@ -50,6 +50,7 @@ interface Options {
   stats(directory?: string): Promise<{ worktrees: WorktreeStats[]; local?: LocalStats }>
   prs(directory?: string): Map<string, PRStatus>
   push(directory?: string): void
+  resolve?(sessionID: string, directory?: string): ManagedSession | undefined
   managed(sessionID: string, directory?: string): boolean
   close(sessionID: string, directory?: string): Promise<void>
   directories?(): string[]
@@ -287,6 +288,7 @@ export class AgentManagerOrchestrationBridge {
           text: request.prompt,
           messageID: request.id,
           signal: active.controller.signal,
+          managed: this.options.resolve?.(request.targetSessionID, origin.directory),
         })
         if (this.disposed || active.cancelled) return
         return { result: { operation: "prompt", sessionID: request.targetSessionID, delivered: true } }
@@ -295,7 +297,12 @@ export class AgentManagerOrchestrationBridge {
         return await this.resolveQuestion(client, root, state, request, origin, active)
       }
       if (request.operation === "move") {
-        move({ state, sessionID: request.targetSessionID, sectionID: request.sectionID })
+        move({
+          state,
+          sessionID: request.targetSessionID,
+          sectionID: request.sectionID,
+          managed: this.options.resolve?.(request.targetSessionID, origin.directory),
+        })
         this.options.push(origin.directory)
         if (this.disposed || active.cancelled) return
         return {
@@ -338,6 +345,7 @@ export class AgentManagerOrchestrationBridge {
       sessionID: request.targetSessionID,
       questionID: request.questionID,
       answers: request.answers,
+      managed: this.options.resolve?.(request.targetSessionID, origin.directory),
     })
     if (this.disposed || active.cancelled) return
     return {

packages/kilo-vscode/src/agent-manager/orchestration-domain.ts

@@ -316,6 +316,7 @@ interface Target {
   root: string
   state: WorktreeStateManager
   sessionID: string
+  managed?: ManagedSession
 }
 
 interface Located {
@@ -326,8 +327,8 @@ interface Located {
 // Verify the target is a live managed session of this workspace and return its authoritative
 // directory plus display name, so error messages can echo exact IDs back to the caller.
 async function locate(input: Target): Promise<Located> {
-  const managed = input.state.getSession(input.sessionID)
-  if (!managed)
+  const managed = input.state.getSession(input.sessionID) ?? input.managed
+  if (!managed || managed.id !== input.sessionID)
     throw new OrchestrationError("unknown_session", "The session is not managed by this Agent Manager workspace")
   const dir = directory(input.root, input.state, managed)
   if (
@@ -398,6 +399,7 @@ export async function prompt(input: {
   messageID: string
   signal?: AbortSignal
   idleTimeoutMs?: number
+  managed?: ManagedSession
 }): Promise<void> {
   if (input.signal?.aborted) return
   const target = await locate(input)
@@ -424,6 +426,7 @@ export async function answer(input: {
   sessionID: string
   questionID?: string
   answers: string[][]
+  managed?: ManagedSession
 }): Promise<{ questionID: string }> {
   const dir = (await locate(input)).dir
   const listed = await input.client.question.list({ directory: dir })
@@ -487,9 +490,14 @@ async function waitForIdle(
   return waitForIdle(client, directory, sessionID, signal, timeout, start)
 }
 
-export function move(input: { state: WorktreeStateManager; sessionID: string; sectionID: string | null }): void {
-  const session = input.state.getSession(input.sessionID)
-  if (!session)
+export function move(input: {
+  state: WorktreeStateManager
+  sessionID: string
+  sectionID: string | null
+  managed?: ManagedSession
+}): void {
+  const session = input.state.getSession(input.sessionID) ?? input.managed
+  if (!session || session.id !== input.sessionID)
     throw new OrchestrationError("unknown_session", "The session is not managed by this Agent Manager workspace")
   if (!session.worktreeId) {
     if (input.sectionID === null) return

packages/kilo-vscode/src/agent-manager/orchestration-setup.ts

@@ -20,6 +20,7 @@ export interface OrchestrationBridgeDeps {
   getPrs: () => Map<string, PRStatus>
   pushState: (ctx?: ProjectContext) => void
   hasPanelSession: (id: string) => boolean
+  routeSession: (id: string, directory: string) => void
   closeSession: (id: string) => Promise<unknown>
   postSessionClosed: (id: string, projectId?: string) => void
   log: (...args: unknown[]) => void
@@ -45,15 +46,37 @@ export function createOrchestrationBridge(deps: OrchestrationBridgeDeps): AgentM
       const ctx = dir ? deps.contexts.byDirectory(dir) : undefined
       deps.pushState(ctx)
     },
+    resolve: (id, dir) => {
+      const ctx = dir ? deps.contexts.byDirectory(dir) : undefined
+      const state = ctx?.peekState()
+      if (state?.isSessionClosed(id)) return undefined
+      const stored = state?.getSession(id)
+      if (stored) return stored
+      if (!ctx) return undefined
+      const live = ctx.sessions().find((session) => session.id === id)
+      if (!live?.worktreeId || !state?.getWorktree(live.worktreeId)) return undefined
+      return { id, worktreeId: live.worktreeId, createdAt: live.createdAt }
+    },
     managed: (id, dir) => {
       const ctx = dir ? deps.contexts.byDirectory(dir) : undefined
-      if (ctx) return ctx.hasLiveSession(id) || !!ctx.peekState()?.getSession(id)
+      if (ctx) {
+        const state = ctx.peekState()
+        return !state?.isSessionClosed(id) && (!!state?.getSession(id) || ctx.hasLiveSession(id))
+      }
       return deps.hasPanelSession(id) || !!deps.getState()?.getSession(id)
     },
     close: async (id, dir) => {
       const ctx = dir ? deps.contexts.byDirectory(dir) : undefined
       if (ctx) {
+        const state = ctx.peekState()
+        const stored = state?.getSession(id)
+        const live = ctx.sessions().find((session) => session.id === id)
+        const wt = live?.worktreeId ? state?.getWorktree(live.worktreeId) : undefined
+        if (wt && !stored) deps.routeSession(id, wt.path)
         await deps.projectScope.run(ctx, () => deps.closeSession(id))
+        state?.closeSession(id, wt?.id ?? stored?.worktreeId ?? null)
+        await state?.flush()
+        ctx.removeLiveSession(id)
       } else {
         await deps.closeSession(id)
       }

packages/kilo-vscode/src/agent-manager/state-recovery.ts

@@ -22,7 +22,7 @@ export function restoreWorktrees(state: WorktreeStateManager, infos: WorktreeInf
       })
 
     if (!existing) result.worktrees++
-    if (!info.sessionId) continue
+    if (!info.sessionId || state.isSessionClosed(info.sessionId)) continue
 
     const session = state.getSession(info.sessionId)
     if (!session) {