fix(vscode): route worktree sessions to correct tab (#9208)

Ensure Agent Manager pushes worktree session state before generic sessionCreated messages and reconciles local tabs against worktree ownership so forked, continued, and newly-created worktree sessions do not appear in Local. Fixes #8983.

.changeset/fix-worktree-session-routing.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Fixed Agent Manager incorrectly duplicating worktree sessions into the Local tab when forking a session or using "Continue in Worktree".

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

@@ -776,8 +776,10 @@ export class AgentManagerProvider implements Disposable {
     const state = this.getStateManager()!
     state.addSession(session.id, created.worktree.id)
     this.registerWorktreeSession(session.id, created.result.path)
-    this.panel?.sessions.registerSession(session)
+    // Push state before registerSession so the webview's sessionCreated handler
+    // sees the worktree mapping and routes the session to the worktree tab.
     this.notifyWorktreeReady(session.id, created.result, created.worktree.id)
+    this.panel?.sessions.registerSession(session)
     this.host.capture("Agent Manager Session Started", {
       source: PLATFORM,
       sessionId: session.id,

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

@@ -152,6 +152,25 @@ describe("AgentManagerProvider worktree creation", () => {
     expect(manager.panel!.sessions.registerSession).toHaveBeenCalledWith(session)
   })
 
+  // Regression for #8983: notifyWorktreeReady must push agentManager.state before
+  // registerSession posts sessionCreated. Reverse order makes the webview route the
+  // new worktree session into the Local tab.
+  it("pushes worktree state before registering the session", async () => {
+    const manager = createHarness()
+    manager.createWorktreeOnDisk.mockResolvedValue({
+      worktree: { id: "wt-1" },
+      result: { path: "/repo/.kilo/worktrees/wt-1", branch: "feature/wt-1", parentBranch: "main" },
+    })
+    manager.createSessionInWorktree.mockResolvedValue({ id: "session-1" })
+    manager.getStateManager.mockReturnValue({ addSession: vi.fn() })
+
+    await manager.onCreateWorktree()
+
+    const notify = manager.notifyWorktreeReady.mock.invocationCallOrder[0]!
+    const register = manager.panel!.sessions.registerSession.mock.invocationCallOrder[0]!
+    expect(notify).toBeLessThan(register)
+  })
+
   it("waits for state initialization before creating a worktree", async () => {
     const manager = createHarness()
     const ready = deferred()

packages/kilo-vscode/src/agent-manager/continue-in-worktree.ts

@@ -99,8 +99,12 @@ export function registerSession(
   const state = ctx.getStateManager()
   if (state) state.addSession(session.id, worktreeId)
   ctx.registerWorktreeSession(session.id, result.path)
-  ctx.registerSession(session)
+  // Push state before registerSession so the webview knows this is a worktree
+  // session before receiving the sessionCreated message. Without this ordering,
+  // the sessionCreated handler would add the session to the local tab because
+  // managedSessions (and thus worktreeSessionIds) hadn't been updated yet.
   ctx.notifyReady(session.id, result, worktreeId)
+  ctx.registerSession(session)
   ctx.capture("Continue in Worktree", { source: PLATFORM, sessionId: session.id, worktreeId })
   ctx.log(`Continued sidebar session ${sourceId} → worktree ${worktreeId} (session ${session.id})`)
 }

packages/kilo-vscode/tests/unit/continue-in-worktree.test.ts

@@ -126,7 +126,7 @@ describe("continue-in-worktree steps", () => {
         capture: () => calls.push("capture"),
       })
       registerSession(c, session("s1"), result("/tmp/wt"), "wt1", "src-session")
-      expect(calls).toEqual(["addSession", "registerWorktreeSession", "registerSession", "notifyReady", "capture"])
+      expect(calls).toEqual(["addSession", "registerWorktreeSession", "notifyReady", "registerSession", "capture"])
     })
 
     it("works without state manager", () => {
@@ -139,7 +139,7 @@ describe("continue-in-worktree steps", () => {
         capture: () => calls.push("capture"),
       })
       registerSession(c, session("s1"), result("/tmp/wt"), "wt1", "src-session")
-      expect(calls).toEqual(["registerWorktreeSession", "registerSession", "notifyReady", "capture"])
+      expect(calls).toEqual(["registerWorktreeSession", "notifyReady", "registerSession", "capture"])
     })
   })
 })

packages/kilo-vscode/tests/unit/navigate.test.ts

@@ -225,6 +225,18 @@ describe("restoreLocalSessions", () => {
     expect(result).toEqual(["s2"])
   })
 
+  it("evicts worktree-bound sessions already in current local state", () => {
+    // Regression: sessionCreated (SSE) can race ahead of agentManager.state and
+    // wrongly add a worktree session to localSessionIDs. On the next state push
+    // the worktree mapping arrives and the session must be evicted from local.
+    const sessions = [
+      { id: "s1", worktreeId: null },
+      { id: "s2", worktreeId: "wt-1" },
+    ]
+    const result = restoreLocalSessions(sessions, ["s1", "s2"], undefined, isPending, identity)
+    expect(result).toEqual(["s1"])
+  })
+
   it("applies tab order on restore", () => {
     const sessions = [
       { id: "s1", worktreeId: null },

packages/kilo-vscode/webview-ui/agent-manager/AgentManagerApp.tsx

@@ -339,6 +339,9 @@ const AgentManagerContent: Component = () => {
   // Recover persisted local session IDs from webview state
   const persisted = vscode.getState<{ localSessionIDs?: string[]; sidebarWidth?: number }>()
   const [localSessionIDs, setLocalSessionIDs] = createSignal<string[]>(persisted?.localSessionIDs ?? [])
+  /** Remove a session ID from the local tab (no-op if absent). */
+  const evictLocal = (sid: string) =>
+    setLocalSessionIDs((prev) => (prev.includes(sid) ? prev.filter((id) => id !== sid) : prev))
   const [sidebarWidth, setSidebarWidth] = createSignal(persisted?.sidebarWidth ?? DEFAULT_SIDEBAR_WIDTH)
   const [sessionsCollapsed, setSessionsCollapsed] = createSignal(false)
   const [sections, setSections] = createSignal<SectionState[]>([])
@@ -1167,14 +1170,7 @@ const AgentManagerContent: Component = () => {
         const ev = msg as AgentManagerWorktreeSetupMessage
         if (ev.status === "ready" || ev.status === "error") {
           const error = ev.status === "error"
-          // Remove from busy map
-          if (ev.worktreeId) {
-            setBusyWorktrees((prev) => {
-              const next = new Map(prev)
-              next.delete(ev.worktreeId!)
-              return next
-            })
-          }
+          if (ev.worktreeId) setBusyWorktrees((prev) => new Map([...prev].filter(([k]) => k !== ev.worktreeId)))
           setSetup({
             active: true,
             message: ev.message,
@@ -1186,9 +1182,9 @@ const AgentManagerContent: Component = () => {
           globalThis.setTimeout(() => setSetup({ active: false, message: "" }), error ? 3000 : 500)
           if (!error && ev.sessionId) {
             session.selectSession(ev.sessionId)
-            // Auto-switch sidebar to the worktree containing this session
             const ms = managedSessions().find((s) => s.id === ev.sessionId)
             if (ms?.worktreeId) setSelection(ms.worktreeId)
+            evictLocal(ev.sessionId)
           }
         } else {
           // Track this worktree as setting up and auto-select it in the sidebar
@@ -1223,6 +1219,10 @@ const AgentManagerContent: Component = () => {
             return [...prev, ev.sessionId]
           })
           vscode.postMessage({ type: "agentManager.persistSession", sessionId: ev.sessionId })
+        } else {
+          saveTabMemory()
+          setSelection(ev.worktreeId)
+          evictLocal(ev.sessionId)
         }
         session.selectSession(ev.sessionId)
       }

packages/kilo-vscode/webview-ui/agent-manager/navigate.ts

@@ -95,6 +95,11 @@ export function restoreLocalSessions(
   applyOrder: (items: { id: string }[], order: string[]) => { id: string }[],
 ): string[] | undefined {
   const locals = sessions.filter((s) => !s.worktreeId).map((s) => s.id)
+  // Sessions assigned to a worktree must never appear in the local tab. A race
+  // where sessionCreated (SSE) arrives before agentManager.state can incorrectly
+  // add a worktree session to localSessionIDs; evict them here on every state push.
+  const worktree = new Set(sessions.filter((s) => s.worktreeId).map((s) => s.id))
+  const evict = (ids: string[]) => (worktree.size > 0 ? ids.filter((id) => !worktree.has(id)) : ids)
   const real = current.filter((id) => !isPending(id))
 
   // First restore: current has no real sessions but disk has some
@@ -109,7 +114,9 @@ export function restoreLocalSessions(
   // Merge any disk-persisted sessions missing from current (e.g. vscode.setState
   // debounce didn't fire before close, but persistSession already wrote to disk)
   const missing = locals.filter((id) => !current.includes(id))
-  const merged = missing.length > 0 ? [...current, ...missing] : current
+  const base = missing.length > 0 ? [...current, ...missing] : current
+  const merged = evict(base)
+  const changed = missing.length > 0 || merged.length !== base.length
 
   // Apply tab order if present
   if (tabOrder && merged.length > 0) {
@@ -119,7 +126,7 @@ export function restoreLocalSessions(
     ).map((item) => item.id)
   }
 
-  return missing.length > 0 ? merged : undefined
+  return changed ? merged : undefined
 }
 
 /**