Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/fix-worktree-session-routing.md
Original file line number Diff line number Diff line change
@@ -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".
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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})`)
}
Expand Down
4 changes: 2 additions & 2 deletions packages/kilo-vscode/tests/unit/continue-in-worktree.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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", () => {
Expand All @@ -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"])
})
})
})
12 changes: 12 additions & 0 deletions packages/kilo-vscode/tests/unit/navigate.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 },
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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[]>([])
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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)
}
Expand Down
11 changes: 9 additions & 2 deletions packages/kilo-vscode/webview-ui/agent-manager/navigate.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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) {
Expand All @@ -119,7 +126,7 @@ export function restoreLocalSessions(
).map((item) => item.id)
}

return missing.length > 0 ? merged : undefined
return changed ? merged : undefined
}

/**
Expand Down
Loading