diff --git a/.changeset/fix-worktree-session-routing.md b/.changeset/fix-worktree-session-routing.md new file mode 100644 index 00000000000..d0d5fae9181 --- /dev/null +++ b/.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". diff --git a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts index a0b0ac9fd52..61b47c2c28d 100644 --- a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts +++ b/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, diff --git a/packages/kilo-vscode/src/agent-manager/__tests__/AgentManagerProvider.spec.ts b/packages/kilo-vscode/src/agent-manager/__tests__/AgentManagerProvider.spec.ts index 139962dd4c5..0c4dd458fcf 100644 --- a/packages/kilo-vscode/src/agent-manager/__tests__/AgentManagerProvider.spec.ts +++ b/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() diff --git a/packages/kilo-vscode/src/agent-manager/continue-in-worktree.ts b/packages/kilo-vscode/src/agent-manager/continue-in-worktree.ts index 30e0168ba97..e25cb355588 100644 --- a/packages/kilo-vscode/src/agent-manager/continue-in-worktree.ts +++ b/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})`) } diff --git a/packages/kilo-vscode/tests/unit/continue-in-worktree.test.ts b/packages/kilo-vscode/tests/unit/continue-in-worktree.test.ts index 9bbbb260b08..c2dab4a182a 100644 --- a/packages/kilo-vscode/tests/unit/continue-in-worktree.test.ts +++ b/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"]) }) }) }) diff --git a/packages/kilo-vscode/tests/unit/navigate.test.ts b/packages/kilo-vscode/tests/unit/navigate.test.ts index 5adbb3abece..61d7110eb60 100644 --- a/packages/kilo-vscode/tests/unit/navigate.test.ts +++ b/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 }, diff --git a/packages/kilo-vscode/webview-ui/agent-manager/AgentManagerApp.tsx b/packages/kilo-vscode/webview-ui/agent-manager/AgentManagerApp.tsx index ffdd0c1031b..d127020eaaa 100644 --- a/packages/kilo-vscode/webview-ui/agent-manager/AgentManagerApp.tsx +++ b/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(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([]) @@ -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) } diff --git a/packages/kilo-vscode/webview-ui/agent-manager/navigate.ts b/packages/kilo-vscode/webview-ui/agent-manager/navigate.ts index d611fc42d7e..b9098a99c4a 100644 --- a/packages/kilo-vscode/webview-ui/agent-manager/navigate.ts +++ b/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 } /**