diff --git a/src/features/boulder-state/storage.test.ts b/src/features/boulder-state/storage.test.ts index f1a2671c6d..9d685e7982 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 99aed0106a..c42fc88128 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 b231e165fb..f56dcdaa26 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.test.ts b/src/hooks/atlas/index.test.ts index 109ed3de96..a1b165e7c2 100644 --- a/src/hooks/atlas/index.test.ts +++ b/src/hooks/atlas/index.test.ts @@ -858,8 +858,8 @@ describe("atlas hook", () => { expect(callArgs.body.parts[0].text).toContain("2 remaining") }) - test("should not inject when last agent is not Atlas", async () => { - // given - boulder state with incomplete plan, but last agent is NOT Atlas + test("should not inject when last agent does not match boulder agent", async () => { + // given - boulder state with incomplete plan, but last agent does NOT match const planPath = join(TEST_DIR, "test-plan.md") writeFileSync(planPath, "# Plan\n- [ ] Task 1\n- [ ] Task 2") @@ -868,10 +868,11 @@ describe("atlas hook", () => { started_at: "2026-01-02T10:00:00Z", session_ids: [MAIN_SESSION_ID], plan_name: "test-plan", + agent: "atlas", } writeBoulderState(TEST_DIR, state) - // given - last agent is NOT Atlas + // given - last agent is NOT the boulder agent cleanupMessageStorage(MAIN_SESSION_ID) setupMessageStorage(MAIN_SESSION_ID, "sisyphus") @@ -886,10 +887,44 @@ describe("atlas hook", () => { }, }) - // then - should NOT call prompt because agent is not Atlas + // then - should NOT call prompt because agent does not match expect(mockInput._promptMock).not.toHaveBeenCalled() }) + test("should inject when last agent matches boulder agent even if non-Atlas", async () => { + // given - boulder state expects sisyphus and last agent is sisyphus + const planPath = join(TEST_DIR, "test-plan.md") + writeFileSync(planPath, "# Plan\n- [ ] Task 1\n- [ ] Task 2") + + const state: BoulderState = { + active_plan: planPath, + started_at: "2026-01-02T10:00:00Z", + session_ids: [MAIN_SESSION_ID], + plan_name: "test-plan", + agent: "sisyphus", + } + writeBoulderState(TEST_DIR, state) + + cleanupMessageStorage(MAIN_SESSION_ID) + setupMessageStorage(MAIN_SESSION_ID, "sisyphus") + + const mockInput = createMockPluginInput() + const hook = createAtlasHook(mockInput) + + // when + await hook.handler({ + event: { + type: "session.idle", + properties: { sessionID: MAIN_SESSION_ID }, + }, + }) + + // then - should call prompt for sisyphus + expect(mockInput._promptMock).toHaveBeenCalled() + const callArgs = mockInput._promptMock.mock.calls[0][0] + expect(callArgs.body.agent).toBe("sisyphus") + }) + test("should debounce rapid continuation injections (prevent infinite loop)", async () => { // given - boulder state with incomplete plan const planPath = join(TEST_DIR, "test-plan.md") diff --git a/src/hooks/atlas/index.ts b/src/hooks/atlas/index.ts index 3d70859e6f..3bf4e78ac9 100644 --- a/src/hooks/atlas/index.ts +++ b/src/hooks/atlas/index.ts @@ -26,6 +26,13 @@ function isSisyphusPath(filePath: string): boolean { const WRITE_EDIT_TOOLS = ["Write", "Edit", "write", "edit"] +function getLastAgentFromSession(sessionID: string): string | null { + const messageDir = getMessageDir(sessionID) + if (!messageDir) return null + const nearest = findNearestMessageWithFields(messageDir) + return nearest?.agent?.toLowerCase() ?? null +} + const DIRECT_WORK_REMINDER = ` --- @@ -431,7 +438,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 +484,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 }], }, @@ -549,8 +556,14 @@ export function createAtlasHook( return } - if (!isCallerOrchestrator(sessionID)) { - log(`[${HOOK_NAME}] Skipped: last agent is not Atlas`, { sessionID }) + const requiredAgent = (boulderState.agent ?? "atlas").toLowerCase() + const lastAgent = getLastAgentFromSession(sessionID) + if (!lastAgent || lastAgent !== requiredAgent) { + log(`[${HOOK_NAME}] Skipped: last agent does not match boulder agent`, { + sessionID, + lastAgent: lastAgent ?? "unknown", + requiredAgent, + }) return } @@ -568,7 +581,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/prometheus-md-only/index.test.ts b/src/hooks/prometheus-md-only/index.test.ts index d564424027..07dff9bf58 100644 --- a/src/hooks/prometheus-md-only/index.test.ts +++ b/src/hooks/prometheus-md-only/index.test.ts @@ -352,6 +352,121 @@ describe("prometheus-md-only", () => { }) }) + describe("boulder state priority over message files (fixes #927)", () => { + const BOULDER_DIR = join(tmpdir(), `boulder-test-${randomUUID()}`) + const BOULDER_FILE = join(BOULDER_DIR, ".sisyphus", "boulder.json") + + beforeEach(() => { + mkdirSync(join(BOULDER_DIR, ".sisyphus"), { recursive: true }) + }) + + afterEach(() => { + rmSync(BOULDER_DIR, { recursive: true, force: true }) + }) + + //#given session was started with prometheus (first message), but /start-work set boulder agent to atlas + //#when user types "continue" after interruption (memory cleared, falls back to message files) + //#then should use boulder state agent (atlas), not message file agent (prometheus) + test("should prioritize boulder agent over message file agent", async () => { + // given - prometheus in message files (from /plan) + setupMessageStorage(TEST_SESSION_ID, "prometheus") + + // given - atlas in boulder state (from /start-work) + writeFileSync(BOULDER_FILE, JSON.stringify({ + active_plan: "/test/plan.md", + started_at: new Date().toISOString(), + session_ids: [TEST_SESSION_ID], + plan_name: "test-plan", + agent: "atlas" + })) + + const hook = createPrometheusMdOnlyHook({ + client: {}, + directory: BOULDER_DIR, + } as never) + + const input = { + tool: "Write", + sessionID: TEST_SESSION_ID, + callID: "call-1", + } + const output = { + args: { filePath: "/path/to/code.ts" }, + } + + // when / then - should NOT block because boulder says atlas, not prometheus + await expect( + hook["tool.execute.before"](input, output) + ).resolves.toBeUndefined() + }) + + test("should use prometheus from boulder state when set", async () => { + // given - atlas in message files (from some other agent) + setupMessageStorage(TEST_SESSION_ID, "atlas") + + // given - prometheus in boulder state (edge case, but should honor it) + writeFileSync(BOULDER_FILE, JSON.stringify({ + active_plan: "/test/plan.md", + started_at: new Date().toISOString(), + session_ids: [TEST_SESSION_ID], + plan_name: "test-plan", + agent: "prometheus" + })) + + const hook = createPrometheusMdOnlyHook({ + client: {}, + directory: BOULDER_DIR, + } as never) + + const input = { + tool: "Write", + sessionID: TEST_SESSION_ID, + callID: "call-1", + } + const output = { + args: { filePath: "/path/to/code.ts" }, + } + + // when / then - should block because boulder says prometheus + await expect( + hook["tool.execute.before"](input, output) + ).rejects.toThrow("can only write/edit .md files") + }) + + test("should fall back to message files when session not in boulder", async () => { + // given - prometheus in message files + setupMessageStorage(TEST_SESSION_ID, "prometheus") + + // given - boulder state exists but for different session + writeFileSync(BOULDER_FILE, JSON.stringify({ + active_plan: "/test/plan.md", + started_at: new Date().toISOString(), + session_ids: ["other-session-id"], + plan_name: "test-plan", + agent: "atlas" + })) + + const hook = createPrometheusMdOnlyHook({ + client: {}, + directory: BOULDER_DIR, + } as never) + + const input = { + tool: "Write", + sessionID: TEST_SESSION_ID, + callID: "call-1", + } + const output = { + args: { filePath: "/path/to/code.ts" }, + } + + // when / then - should block because falls back to message files (prometheus) + await expect( + hook["tool.execute.before"](input, output) + ).rejects.toThrow("can only write/edit .md files") + }) + }) + describe("without message storage", () => { test("should handle missing session gracefully (no agent found)", async () => { // given diff --git a/src/hooks/prometheus-md-only/index.ts b/src/hooks/prometheus-md-only/index.ts index 1c311dc024..7c9131f98f 100644 --- a/src/hooks/prometheus-md-only/index.ts +++ b/src/hooks/prometheus-md-only/index.ts @@ -4,6 +4,7 @@ import { join, resolve, relative, isAbsolute } from "node:path" import { HOOK_NAME, PROMETHEUS_AGENT, ALLOWED_EXTENSIONS, ALLOWED_PATH_PREFIX, BLOCKED_TOOLS, PLANNING_CONSULT_WARNING, PROMETHEUS_WORKFLOW_REMINDER } from "./constants" import { findNearestMessageWithFields, findFirstMessageWithAgent, MESSAGE_STORAGE } from "../../features/hook-message-injector" import { getSessionAgent } from "../../features/claude-code-session-state" +import { readBoulderState } from "../../features/boulder-state" import { log } from "../../shared/logger" import { SYSTEM_DIRECTIVE_PREFIX } from "../../shared/system-directive" import { getAgentDisplayName } from "../../shared/agent-display-names" @@ -70,8 +71,31 @@ function getAgentFromMessageFiles(sessionID: string): string | undefined { return findFirstMessageWithAgent(messageDir) ?? findNearestMessageWithFields(messageDir)?.agent } -function getAgentFromSession(sessionID: string): string | undefined { - return getSessionAgent(sessionID) ?? getAgentFromMessageFiles(sessionID) +/** + * Get the effective agent for the session. + * Priority order: + * 1. In-memory session agent (most recent, set by /start-work) + * 2. Boulder state agent (persisted across restarts, fixes #927) + * 3. Message files (fallback for sessions without boulder state) + * + * This fixes issue #927 where after interruption: + * - In-memory map is cleared (process restart) + * - Message files return "prometheus" (oldest message from /plan) + * - But boulder.json has agent: "atlas" (set by /start-work) + */ +function getAgentFromSession(sessionID: string, directory: string): string | undefined { + // Check in-memory first (current session) + const memoryAgent = getSessionAgent(sessionID) + if (memoryAgent) return memoryAgent + + // Check boulder state (persisted across restarts) - fixes #927 + const boulderState = readBoulderState(directory) + if (boulderState?.session_ids.includes(sessionID) && boulderState.agent) { + return boulderState.agent + } + + // Fallback to message files + return getAgentFromMessageFiles(sessionID) } export function createPrometheusMdOnlyHook(ctx: PluginInput) { @@ -80,7 +104,7 @@ export function createPrometheusMdOnlyHook(ctx: PluginInput) { input: { tool: string; sessionID: string; callID: string }, output: { args: Record; message?: string } ): Promise => { - const agentName = getAgentFromSession(input.sessionID) + const agentName = getAgentFromSession(input.sessionID, ctx.directory) if (agentName !== PROMETHEUS_AGENT) { return diff --git a/src/hooks/start-work/index.ts b/src/hooks/start-work/index.ts index ce432e46ea..600814bd56 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 += `