diff --git a/src/features/boulder-state/storage.test.ts b/src/features/boulder-state/storage.test.ts index f1a2671c6..9d685e798 100644 --- a/src/features/boulder-state/storage.test.ts +++ b/src/features/boulder-state/storage.test.ts @@ -246,5 +246,33 @@ describe("boulder-state", () => { expect(state.plan_name).toBe("auth-refactor") expect(state.started_at).toBeDefined() }) + + test("should include agent field when provided", () => { + //#given - plan path, session id, and agent type + const planPath = "/path/to/feature.md" + const sessionId = "ses-xyz789" + const agent = "atlas" + + //#when - createBoulderState is called with agent + const state = createBoulderState(planPath, sessionId, agent) + + //#then - state should include the agent field + expect(state.agent).toBe("atlas") + expect(state.active_plan).toBe(planPath) + expect(state.session_ids).toEqual([sessionId]) + expect(state.plan_name).toBe("feature") + }) + + test("should allow agent to be undefined", () => { + //#given - plan path and session id without agent + const planPath = "/path/to/legacy.md" + const sessionId = "ses-legacy" + + //#when - createBoulderState is called without agent + const state = createBoulderState(planPath, sessionId) + + //#then - state should not have agent field (backward compatible) + expect(state.agent).toBeUndefined() + }) }) }) diff --git a/src/features/boulder-state/storage.ts b/src/features/boulder-state/storage.ts index 99aed0106..c42fc8812 100644 --- a/src/features/boulder-state/storage.ts +++ b/src/features/boulder-state/storage.ts @@ -139,12 +139,14 @@ export function getPlanName(planPath: string): string { */ export function createBoulderState( planPath: string, - sessionId: string + sessionId: string, + agent?: string ): BoulderState { return { active_plan: planPath, started_at: new Date().toISOString(), session_ids: [sessionId], plan_name: getPlanName(planPath), + ...(agent !== undefined ? { agent } : {}), } } diff --git a/src/features/boulder-state/types.ts b/src/features/boulder-state/types.ts index b231e165f..f56dcdaa2 100644 --- a/src/features/boulder-state/types.ts +++ b/src/features/boulder-state/types.ts @@ -14,6 +14,8 @@ export interface BoulderState { session_ids: string[] /** Plan name derived from filename */ plan_name: string + /** Agent type to use when resuming (e.g., 'atlas') */ + agent?: string } export interface PlanProgress { diff --git a/src/hooks/atlas/index.ts b/src/hooks/atlas/index.ts index 3d70859e6..2583606e5 100644 --- a/src/hooks/atlas/index.ts +++ b/src/hooks/atlas/index.ts @@ -431,7 +431,7 @@ export function createAtlasHook( return state } - async function injectContinuation(sessionID: string, planName: string, remaining: number, total: number): Promise { + async function injectContinuation(sessionID: string, planName: string, remaining: number, total: number, agent?: string): Promise { const hasRunningBgTasks = backgroundManager ? backgroundManager.getTasksByParentSession(sessionID).some(t => t.status === "running") : false @@ -477,7 +477,7 @@ export function createAtlasHook( await ctx.client.session.prompt({ path: { id: sessionID }, body: { - agent: "atlas", + agent: agent ?? "atlas", ...(model !== undefined ? { model } : {}), parts: [{ type: "text", text: prompt }], }, @@ -568,7 +568,7 @@ export function createAtlasHook( state.lastContinuationInjectedAt = now const remaining = progress.total - progress.completed - injectContinuation(sessionID, boulderState.plan_name, remaining, progress.total) + injectContinuation(sessionID, boulderState.plan_name, remaining, progress.total, boulderState.agent) return } diff --git a/src/hooks/start-work/index.ts b/src/hooks/start-work/index.ts index ce432e46e..600814bd5 100644 --- a/src/hooks/start-work/index.ts +++ b/src/hooks/start-work/index.ts @@ -102,7 +102,7 @@ All ${progress.total} tasks are done. Create a new plan with: /plan "your task"` if (existingState) { clearBoulderState(ctx.directory) } - const newState = createBoulderState(matchedPlan, sessionId) + const newState = createBoulderState(matchedPlan, sessionId, "atlas") writeBoulderState(ctx.directory, newState) contextInfo = ` @@ -187,7 +187,7 @@ All ${plans.length} plan(s) are complete. Create a new plan with: /plan "your ta } else if (incompletePlans.length === 1) { const planPath = incompletePlans[0] const progress = getPlanProgress(planPath) - const newState = createBoulderState(planPath, sessionId) + const newState = createBoulderState(planPath, sessionId, "atlas") writeBoulderState(ctx.directory, newState) contextInfo += `