Merge pull request #1477 from kaizen403/fix/boulder-agent-tracking

fix: track agent in boulder state to fix session continuation (fixes #927)

Author: code-yeongyu

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()
+    })
   })
 })

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 } : {}),
   }
 }

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 {

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")

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<void> {
+  async function injectContinuation(sessionID: string, planName: string, remaining: number, total: number, agent?: string): Promise<void> {
     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
       }
 

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

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<string, unknown>; message?: string }
     ): Promise<void> => {
-      const agentName = getAgentFromSession(input.sessionID)
+      const agentName = getAgentFromSession(input.sessionID, ctx.directory)
 
       if (agentName !== PROMETHEUS_AGENT) {
         return

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 += `