diff --git a/.changeset/steady-agent-manager-ownership.md b/.changeset/steady-agent-manager-ownership.md new file mode 100644 index 0000000000..2f7cb3be4c --- /dev/null +++ b/.changeset/steady-agent-manager-ownership.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Recognize sessions discovered in managed Agent Manager worktrees during orchestration actions. diff --git a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts index f25e48235c..3b21aa9175 100644 --- a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts +++ b/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 diff --git a/packages/kilo-vscode/src/agent-manager/WorktreeStateManager.ts b/packages/kilo-vscode/src/agent-manager/WorktreeStateManager.ts index d924129626..85c29c10ce 100644 --- a/packages/kilo-vscode/src/agent-manager/WorktreeStateManager.ts +++ b/packages/kilo-vscode/src/agent-manager/WorktreeStateManager.ts @@ -88,6 +88,7 @@ export interface ManagedSession { interface StateFile { worktrees: Record> sessions: Record> + closedSessions?: Record sections?: Record> tabOrder?: Record 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() private sessions = new Map() + private closed = new Map() private sections = new Map() private tabOrder: Record = {} 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 { 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) { diff --git a/packages/kilo-vscode/src/agent-manager/orchestration-bridge.ts b/packages/kilo-vscode/src/agent-manager/orchestration-bridge.ts index 3912fa4cf1..e057bbe1af 100644 --- a/packages/kilo-vscode/src/agent-manager/orchestration-bridge.ts +++ b/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 push(directory?: string): void + resolve?(sessionID: string, directory?: string): ManagedSession | undefined managed(sessionID: string, directory?: string): boolean close(sessionID: string, directory?: string): Promise 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 { diff --git a/packages/kilo-vscode/src/agent-manager/orchestration-domain.ts b/packages/kilo-vscode/src/agent-manager/orchestration-domain.ts index 0f6a1f3f5f..b356293197 100644 --- a/packages/kilo-vscode/src/agent-manager/orchestration-domain.ts +++ b/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 { - 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 { 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 diff --git a/packages/kilo-vscode/src/agent-manager/orchestration-setup.ts b/packages/kilo-vscode/src/agent-manager/orchestration-setup.ts index 507d0e79d0..349d1a8747 100644 --- a/packages/kilo-vscode/src/agent-manager/orchestration-setup.ts +++ b/packages/kilo-vscode/src/agent-manager/orchestration-setup.ts @@ -20,6 +20,7 @@ export interface OrchestrationBridgeDeps { getPrs: () => Map pushState: (ctx?: ProjectContext) => void hasPanelSession: (id: string) => boolean + routeSession: (id: string, directory: string) => void closeSession: (id: string) => Promise 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) } diff --git a/packages/kilo-vscode/src/agent-manager/state-recovery.ts b/packages/kilo-vscode/src/agent-manager/state-recovery.ts index b39cbc0f52..e94541301d 100644 --- a/packages/kilo-vscode/src/agent-manager/state-recovery.ts +++ b/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) { diff --git a/packages/kilo-vscode/tests/unit/agent-manager-orchestration-bridge.test.ts b/packages/kilo-vscode/tests/unit/agent-manager-orchestration-bridge.test.ts index a543211b41..93c724715d 100644 --- a/packages/kilo-vscode/tests/unit/agent-manager-orchestration-bridge.test.ts +++ b/packages/kilo-vscode/tests/unit/agent-manager-orchestration-bridge.test.ts @@ -4,6 +4,9 @@ import * as os from "os" import * as path from "path" import type { AgentManagerRequest, Session } from "@kilocode/sdk/v2/client" import { AgentManagerOrchestrationBridge } from "../../src/agent-manager/orchestration-bridge" +import { createOrchestrationBridge } from "../../src/agent-manager/orchestration-setup" +import { ProjectContexts } from "../../src/agent-manager/project/contexts" +import { ProjectScope } from "../../src/agent-manager/project/scope" import { WorktreeStateManager } from "../../src/agent-manager/WorktreeStateManager" import type { SSEPayload } from "../../src/services/cli-backend/sdk-sse-adapter" @@ -113,6 +116,7 @@ describe("AgentManagerOrchestrationBridge", () => { }, prs: (dir) => (overrides?.prs ? overrides.prs(dir) : new Map()), push: (dir) => (overrides?.push ? overrides.push(dir) : push()), + resolve: (id, dir) => overrides?.resolve?.(id, dir), managed: (id, dir) => (overrides?.managed ? overrides.managed(id, dir) : managed.has(id)), close: async (id, dir) => (overrides?.close ? overrides.close(id, dir) : close(id, dir)), log: () => undefined, @@ -126,6 +130,7 @@ describe("AgentManagerOrchestrationBridge", () => { bridge, client, close, + connection, handlers, lists, managed, @@ -308,6 +313,186 @@ describe("AgentManagerOrchestrationBridge", () => { test.bridge.dispose() }) + it("routes prompt, answer, move, and stop for a live-only managed worktree session", async () => { + const wt = state.getWorktrees()[0]! + const live = { id: "ses_live", worktreeId: wt.id, createdAt: "" } + const section = state.addSection("Review", null) + const contexts = new ProjectContexts({ + workspaceRoot: () => root, + registry: { list: () => [], get: () => undefined }, + enabled: () => false, + deps: { log: () => undefined, state: () => state }, + }) + const ctx = contexts.active()! + ctx.stateManager() + ctx.upsertSession({ + ...live, + parentID: null, + title: "Live", + updatedAt: "", + revert: null, + summary: null, + }) + const routes = new Map() + const test = harness() + test.bridge.dispose() + const close = mock(async (id: string) => { + expect(routes.get(id)).toBe(dir) + routes.delete(id) + }) + const bridge = createOrchestrationBridge({ + connectionService: test.connection as never, + contexts, + projectScope: new ProjectScope(), + getRoot: () => ctx.root, + getState: () => state, + getStateReady: () => Promise.resolve(), + initStateReady: () => Promise.resolve(), + getStats: async () => ({ worktrees: [] }), + getPrs: () => new Map(), + pushState: () => undefined, + hasPanelSession: () => false, + routeSession: (id, path) => void routes.set(id, path), + closeSession: close, + postSessionClosed: () => undefined, + log: () => undefined, + }) + const send = (request: AgentManagerRequest) => test.request(request, ctx.root) + + send({ + id: "amr_live_prompt", + sessionID: "ses_caller", + operation: "prompt", + targetSessionID: live.id, + prompt: "Continue", + }) + await waitFor(() => test.replies.length === 1) + expect(test.promptAsync).toHaveBeenCalledWith(expect.objectContaining({ sessionID: live.id, directory: dir }), { + throwOnError: true, + }) + ;(test.client.question.list as ReturnType).mockImplementation(async () => ({ + data: [ + { + id: "que_live", + sessionID: live.id, + questions: [{ header: "Approve", question: "Proceed?", options: [{ label: "Yes", description: "go" }] }], + }, + ], + })) + send({ + id: "amr_live_answer", + sessionID: "ses_caller", + operation: "answer", + targetSessionID: live.id, + answers: [["Yes"]], + }) + await waitFor(() => test.replies.length === 2) + expect(test.questionReply).toHaveBeenCalledWith( + { requestID: "que_live", answers: [["Yes"]], directory: dir }, + { throwOnError: true }, + ) + + send({ + id: "amr_live_move", + sessionID: "ses_caller", + operation: "move", + targetSessionID: live.id, + sectionID: section.id, + }) + await waitFor(() => test.replies.length === 3) + expect(state.getWorktree(wt.id)?.sectionId).toBe(section.id) + + send({ + id: "amr_live_stop", + sessionID: "ses_caller", + operation: "stop", + targetSessionID: live.id, + }) + await waitFor(() => test.replies.length === 4) + expect(close).toHaveBeenCalledWith(live.id) + expect(ctx.hasLiveSession(live.id)).toBe(false) + expect(state.getSession(live.id)).toBeUndefined() + + ctx.upsertSession({ + ...live, + parentID: null, + title: "Live", + updatedAt: "", + revert: null, + summary: null, + }) + send({ + id: "amr_live_closed", + sessionID: "ses_caller", + operation: "prompt", + targetSessionID: live.id, + prompt: "Do not reopen", + }) + await waitFor(() => test.rejections.length === 1) + expect(test.rejections[0]).toMatchObject({ error: { code: "unknown_session" } }) + bridge.dispose() + }) + + it("rejects a stopped live-only session after its project state is restored", async () => { + const wt = state.getWorktrees()[0]! + state.closeSession("ses_stopped", wt.id) + await state.flush() + const restored = new WorktreeStateManager(root, () => undefined) + await restored.load() + const contexts = new ProjectContexts({ + workspaceRoot: () => root, + registry: { list: () => [], get: () => undefined }, + enabled: () => false, + deps: { log: () => undefined, state: () => restored }, + }) + const ctx = contexts.active()! + ctx.stateManager() + ctx.upsertSession({ + id: "ses_stopped", + worktreeId: wt.id, + parentID: null, + title: "Stopped", + createdAt: "", + updatedAt: "", + revert: null, + summary: null, + }) + const test = harness() + test.bridge.dispose() + const bridge = createOrchestrationBridge({ + connectionService: test.connection as never, + contexts, + projectScope: new ProjectScope(), + getRoot: () => ctx.root, + getState: () => restored, + getStateReady: () => Promise.resolve(), + initStateReady: () => Promise.resolve(), + getStats: async () => ({ worktrees: [] }), + getPrs: () => new Map(), + pushState: () => undefined, + hasPanelSession: () => false, + routeSession: () => undefined, + closeSession: async () => undefined, + postSessionClosed: () => undefined, + log: () => undefined, + }) + + test.request( + { + id: "amr_restored_stopped", + sessionID: "ses_caller", + operation: "prompt", + targetSessionID: "ses_stopped", + prompt: "Do not reopen", + }, + ctx.root, + ) + await waitFor(() => test.rejections.length === 1) + expect(test.rejections[0]).toMatchObject({ error: { code: "unknown_session" } }) + expect(test.promptAsync).not.toHaveBeenCalled() + bridge.dispose() + }) + it("answers a managed session's pending question through the backend reply route", async () => { const test = harness() ;(test.client.question.list as ReturnType).mockImplementation(async () => ({ @@ -451,6 +636,53 @@ describe("AgentManagerOrchestrationBridge", () => { test.bridge.dispose() }) + it("keeps live-only secondary worktree sessions scoped to their owning project", async () => { + const secondary = fs.mkdtempSync(path.join(os.tmpdir(), "am-orchestration-secondary-live-")) + const worktree = path.join(secondary, "worktree") + fs.mkdirSync(path.join(secondary, ".kilo"), { recursive: true }) + fs.mkdirSync(worktree) + const other = new WorktreeStateManager(secondary, () => undefined) + const wt = other.addWorktree({ branch: "fix/secondary-live", path: worktree, parentBranch: "main" }) + const live = { id: "ses_secondary_live", worktreeId: wt.id, createdAt: "" } + const test = harness({ + root: (origin) => (origin === secondary ? secondary : root), + ready: async (origin) => (origin === secondary ? other : state), + state: (origin) => (origin === secondary ? other : state), + resolve: (id, origin) => (id === live.id && origin === secondary ? live : undefined), + }) + + test.request( + { + id: "amr_secondary_live", + sessionID: "ses_caller", + operation: "prompt", + targetSessionID: live.id, + prompt: "Continue", + }, + secondary, + ) + await waitFor(() => test.replies.length === 1) + expect(test.promptAsync).toHaveBeenCalledWith( + expect.objectContaining({ sessionID: live.id, directory: worktree }), + { throwOnError: true }, + ) + expect(other.getSession(live.id)).toBeUndefined() + + test.request({ + id: "amr_secondary_foreign", + sessionID: "ses_caller", + operation: "prompt", + targetSessionID: live.id, + prompt: "Cross project", + }) + await waitFor(() => test.rejections.length === 1) + expect(test.rejections[0]).toMatchObject({ error: { code: "unknown_session" } }) + + test.bridge.dispose() + await other.flush() + fs.rmSync(secondary, { recursive: true, force: true }) + }) + it("handles requests for secondary project directories in multi-project mode", async () => { const secondaryRoot = fs.mkdtempSync(path.join(os.tmpdir(), "am-orchestration-secondary-")) fs.mkdirSync(path.join(secondaryRoot, ".kilo"), { recursive: true }) diff --git a/packages/kilo-vscode/tests/unit/agent-manager-orchestration-domain.test.ts b/packages/kilo-vscode/tests/unit/agent-manager-orchestration-domain.test.ts index b8f29e480d..e2dd882902 100644 --- a/packages/kilo-vscode/tests/unit/agent-manager-orchestration-domain.test.ts +++ b/packages/kilo-vscode/tests/unit/agent-manager-orchestration-domain.test.ts @@ -3,8 +3,10 @@ import * as fs from "fs" import * as os from "os" import * as path from "path" import type { KiloClient, QuestionRequest, Session } from "@kilocode/sdk/v2/client" -import { OrchestrationError, answer, overview, prompt } from "../../src/agent-manager/orchestration-domain" +import { OrchestrationError, answer, move, overview, prompt } from "../../src/agent-manager/orchestration-domain" import { WorktreeStateManager } from "../../src/agent-manager/WorktreeStateManager" +import { ProjectContext } from "../../src/agent-manager/project/context" +import { collectProjectSessions } from "../../src/agent-manager/project/init" import type { PRStatus as AgentManagerPRStatus } from "../../src/agent-manager/types" const noQuestions: QuestionRequest[] = [] @@ -210,6 +212,96 @@ describe("Agent Manager orchestration domain", () => { ) }) + it("prompts, answers, and moves a session discovered in a managed worktree", async () => { + const wt = state.addWorktree({ branch: "fix/discovered", path: worktree, parentBranch: "main" }) + const section = state.addSection("Review", null) + const session = { + id: "ses_discovered", + slug: "discovered", + projectID: "prj-test", + directory: worktree, + title: "Discovered", + version: "1", + time: { created: 1, updated: 1 }, + } satisfies Session + const ctx = new ProjectContext("prj-test", root, true, { log: () => undefined, state: () => state }) + ctx.stateManager() + const views = await collectProjectSessions(ctx, { + listSessions: async (dir) => (dir === worktree ? [session] : []), + setSessionDirectory: () => undefined, + }) + expect(views).toEqual([expect.objectContaining({ id: session.id, worktreeId: wt.id })]) + ctx.setSessions(views) + expect(state.getSession(session.id)).toBeUndefined() + const managed = { id: session.id, worktreeId: views[0]!.worktreeId, createdAt: views[0]!.createdAt } + + const questions: QuestionRequest[] = [] + const delivered = mock(async () => ({ data: undefined })) + const replied = mock(async () => ({ data: true })) + const client = { + session: { + get: mock(async () => ({ data: session })), + status: mock(async () => ({ data: {} })), + promptAsync: delivered, + }, + permission: { list: mock(async () => ({ data: [] })) }, + question: { list: mock(async () => ({ data: questions })), reply: replied }, + } as unknown as KiloClient + + await prompt({ client, root, state, sessionID: session.id, text: "Continue", messageID: "amr_discovered", managed }) + expect(delivered).toHaveBeenCalledWith(expect.objectContaining({ sessionID: session.id, directory: worktree }), { + throwOnError: true, + }) + + questions.push({ + id: "que_discovered", + sessionID: session.id, + questions: [{ header: "Approve", question: "Proceed?", options: [{ label: "Yes", description: "Continue" }] }], + }) + await answer({ client, root, state, sessionID: session.id, answers: [["Yes"]], managed }) + expect(replied).toHaveBeenCalledWith( + { requestID: "que_discovered", answers: [["Yes"]], directory: worktree }, + { throwOnError: true }, + ) + + move({ state, sessionID: session.id, sectionID: section.id, managed }) + expect(state.getWorktree(wt.id)?.sectionId).toBe(section.id) + expect(state.getSession(session.id)).toBeUndefined() + }) + + it("recognizes a worktree session received through a live lifecycle event", async () => { + const wt = state.addWorktree({ branch: "fix/live", path: worktree, parentBranch: "main" }) + const ctx = new ProjectContext("prj-test", root, true, { log: () => undefined, state: () => state }) + ctx.stateManager() + ctx.upsertSession({ + id: "ses_live", + parentID: null, + title: "Live", + createdAt: "", + updatedAt: "", + revert: null, + summary: null, + worktreeId: wt.id, + }) + + expect(ctx.hasLiveSession("ses_live")).toBe(true) + expect(state.getSession("ses_live")).toBeUndefined() + const managed = { id: "ses_live", worktreeId: wt.id, createdAt: "" } + const delivered = mock(async () => ({ data: undefined })) + const client = { + session: { + get: mock(async () => ({ data: { id: "ses_live", directory: worktree, title: "Live" } as Session })), + status: mock(async () => ({ data: {} })), + promptAsync: delivered, + }, + permission: { list: mock(async () => ({ data: [] })) }, + question: { list: mock(async () => ({ data: [] })) }, + } as unknown as KiloClient + + await prompt({ client, root, state, sessionID: "ses_live", text: "Continue", messageID: "amr_live", managed }) + expect(delivered).toHaveBeenCalledWith(expect.objectContaining({ directory: worktree }), { throwOnError: true }) + }) + it("waits for a busy managed session to become idle before prompting", async () => { const managed = state.addWorktree({ branch: "fix/wait", path: worktree, parentBranch: "main" }) state.addSession("ses_wait", managed.id) @@ -352,6 +444,28 @@ describe("Agent Manager orchestration domain", () => { ).rejects.toMatchObject({ code: "unknown_session", } satisfies Partial) + await expect( + prompt({ + client, + root, + state, + sessionID: "ses_unknown", + text: "Continue", + messageID: "amr_mismatch", + managed: { id: "ses_target", worktreeId: managed.id, createdAt: "" }, + }), + ).rejects.toMatchObject({ code: "unknown_session" } satisfies Partial) + await expect( + prompt({ + client, + root, + state, + sessionID: "ses_foreign", + text: "Continue", + messageID: "amr_foreign", + managed: { id: "ses_foreign", worktreeId: "wt_foreign", createdAt: "" }, + }), + ).rejects.toMatchObject({ code: "stale_session" } satisfies Partial) await expect( prompt({ client, root, state, sessionID: "ses_target", text: "Continue", messageID: "amr_cross" }), ).rejects.toMatchObject({ diff --git a/packages/kilo-vscode/tests/unit/worktree-state-manager.test.ts b/packages/kilo-vscode/tests/unit/worktree-state-manager.test.ts index 854b6b420d..540c63b653 100644 --- a/packages/kilo-vscode/tests/unit/worktree-state-manager.test.ts +++ b/packages/kilo-vscode/tests/unit/worktree-state-manager.test.ts @@ -199,6 +199,27 @@ describe("WorktreeStateManager", () => { manager.removeSession("s1") expect(manager.getSession("s1")).toBeUndefined() }) + + it("persists stopped worktree sessions across reloads", async () => { + const wt = manager.addWorktree({ branch: "fix", path: "/tmp/fix", parentBranch: "main" }) + manager.closeSession("ses-stopped", wt.id) + await manager.flush() + + const restored = new WorktreeStateManager(root, () => undefined) + await restored.load() + + expect(restored.isSessionClosed("ses-stopped")).toBe(true) + restored.addSession("ses-stopped", wt.id) + expect(restored.isSessionClosed("ses-stopped")).toBe(false) + await restored.flush() + }) + + it("removes stopped-session records when their worktree is deleted", () => { + const wt = manager.addWorktree({ branch: "fix", path: "/tmp/fix", parentBranch: "main" }) + manager.closeSession("ses-stopped", wt.id) + manager.removeWorktree(wt.id) + expect(manager.isSessionClosed("ses-stopped")).toBe(false) + }) }) describe("directoryFor", () => { @@ -375,6 +396,24 @@ describe("WorktreeStateManager", () => { expect(worktree?.remote).toBe("origin") expect(manager.getSession("sess-recovered")?.worktreeId).toBe(worktree?.id) }) + + it("does not recover a session that was explicitly stopped", () => { + const wt = manager.addWorktree({ branch: "fix-recovered", path: "/tmp/recovered", parentBranch: "main" }) + manager.closeSession("sess-stopped", wt.id) + const result = restoreWorktrees(manager, [ + { + branch: "fix-recovered", + path: "/tmp/recovered", + parentBranch: "main", + createdAt: Date.UTC(2026, 0, 1), + sessionId: "sess-stopped", + }, + ]) + + expect(result).toEqual({ worktrees: 0, sessions: 0 }) + expect(manager.getSession("sess-stopped")).toBeUndefined() + expect(manager.isSessionClosed("sess-stopped")).toBe(true) + }) }) describe("tab order", () => {