From 169ccb6b052b349a8f58b24058a1f65fd36444e3 Mon Sep 17 00:00:00 2001
From: Rishi Vhavle <rishivhavle21@gmail.com>
Date: Wed, 4 Feb 2026 13:06:34 +0530
Subject: [PATCH] fix: use boulder agent instead of hardcoded Atlas check for
 continuation

Address code review: continuation was blocked unless last agent was Atlas,
making the new agent parameter ineffective. Now the idle handler checks if
the last session agent matches boulderState.agent (defaults to 'atlas'),
allowing non-Atlas agents to resume when properly configured.

- Add getLastAgentFromSession helper for agent lookup
- Replace isCallerOrchestrator gate with boulder-agent-aware check
- Add test for non-Atlas agent continuation scenario
---
 src/hooks/atlas/index.test.ts | 43 +++++++++++++++++++++++++++++++----
 src/hooks/atlas/index.ts      | 17 ++++++++++++--
 2 files changed, 54 insertions(+), 6 deletions(-)

diff --git a/src/hooks/atlas/index.test.ts b/src/hooks/atlas/index.test.ts
index 109ed3de9..a1b165e7c 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 2583606e5..3bf4e78ac 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 = `
 
 ---
@@ -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
         }