Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 28 additions & 0 deletions src/features/boulder-state/storage.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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()
})
})
})
4 changes: 3 additions & 1 deletion src/features/boulder-state/storage.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 } : {}),
}
}
2 changes: 2 additions & 0 deletions src/features/boulder-state/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
43 changes: 39 additions & 4 deletions src/hooks/atlas/index.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand All @@ -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")

Expand All @@ -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")
Expand Down
23 changes: 18 additions & 5 deletions src/hooks/atlas/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 = `

---
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 }],
},
Expand Down Expand Up @@ -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
}

Expand All @@ -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
}

Expand Down
115 changes: 115 additions & 0 deletions src/hooks/prometheus-md-only/index.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
30 changes: 27 additions & 3 deletions src/hooks/prometheus-md-only/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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) {
Expand All @@ -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
Expand Down
4 changes: 2 additions & 2 deletions src/hooks/start-work/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 = `
Expand Down Expand Up @@ -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 += `
Expand Down