From d12c74120d582b0b098e467113d5812860490870 Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Sun, 5 Apr 2026 17:15:08 +0900 Subject: [PATCH] feat(run): integrate session origins and agent detection into continuation state MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Update continuation-state to use session_origins from boulder state - Add isTrackedDescendantSession helper for lineage detection - Integrate getLastAgentFromSession for reliable agent detection - Update completion tests for new lineage-aware continuation logic - Add JSON backend tests for continuation state 🤖 Generated with assistance of OhMyOpenCode --- src/cli/run/completion-continuation.test.ts | 307 ++++++++++++++++-- .../continuation-state.json-backend.test.ts | 171 ++++++++++ src/cli/run/continuation-state.ts | 55 +++- 3 files changed, 496 insertions(+), 37 deletions(-) create mode 100644 src/cli/run/continuation-state.json-backend.test.ts diff --git a/src/cli/run/completion-continuation.test.ts b/src/cli/run/completion-continuation.test.ts index 36cd8a6f3..20758c344 100644 --- a/src/cli/run/completion-continuation.test.ts +++ b/src/cli/run/completion-continuation.test.ts @@ -3,18 +3,11 @@ import { mkdtempSync, mkdirSync, rmSync, writeFileSync } from "node:fs" import { join } from "node:path" import { tmpdir } from "node:os" import type { RunContext } from "./types" -import { - _resetForTesting, - registerAgentName, - setSessionAgent, - subagentSessions, -} from "../../features/claude-code-session-state" import { writeState as writeRalphLoopState } from "../../hooks/ralph-loop/storage" const testDirs: string[] = [] afterEach(() => { - _resetForTesting() while (testDirs.length > 0) { const dir = testDirs.pop() if (dir) { @@ -42,6 +35,7 @@ function createMockContext(directory: string): RunContext { parentID: undefined, }, })), + messages: mock(async () => ({ data: [] })), }, } as unknown as RunContext["client"], sessionID: "test-session", @@ -50,7 +44,12 @@ function createMockContext(directory: string): RunContext { } } -function writeBoulderStateFile(directory: string, activePlanPath: string, sessionIDs: string[]): void { +function writeBoulderStateFile( + directory: string, + activePlanPath: string, + sessionIDs: string[], + sessionOrigins?: Record, +): void { const sisyphusDir = join(directory, ".sisyphus") mkdirSync(sisyphusDir, { recursive: true }) writeFileSync( @@ -59,6 +58,7 @@ function writeBoulderStateFile(directory: string, activePlanPath: string, sessio active_plan: activePlanPath, started_at: new Date().toISOString(), session_ids: sessionIDs, + session_origins: sessionOrigins, plan_name: "test-plan", agent: "atlas", }), @@ -103,26 +103,31 @@ describe("checkCompletionConditions continuation coverage", () => { expect(result).toBe(true) }) - it("returns false when current session is a descendant of an active boulder session with unchecked plan items", async () => { + it("returns false when current session is an appended descendant of an active boulder session with unchecked plan items", async () => { // given spyOn(console, "log").mockImplementation(() => {}) - registerAgentName("atlas") const directory = createTempDir() const planPath = join(directory, ".sisyphus", "plans", "active-descendant-plan.md") mkdirSync(join(directory, ".sisyphus", "plans"), { recursive: true }) writeFileSync(planPath, "- [ ] unfinished task\n", "utf-8") - writeBoulderStateFile(directory, planPath, ["root-session"]) + writeBoulderStateFile(directory, planPath, ["root-session", "child-session"], { + "root-session": "direct", + "child-session": "appended", + }) const ctx = createMockContext(directory) ctx.sessionID = "child-session" - subagentSessions.add("child-session") - setSessionAgent("child-session", "atlas") ctx.client.session.get = mock(async ({ path }: { path: { id: string } }) => ({ data: { id: path.id, parentID: path.id === "child-session" ? "root-session" : undefined, }, })) as unknown as RunContext["client"]["session"]["get"] + ctx.client.session.messages = mock(async ({ path }: { path: { id: string } }) => ({ + data: path.id === "child-session" + ? [{ info: { agent: "atlas", providerID: "openai", modelID: "gpt-5.4" } }] + : [], + })) as unknown as RunContext["client"]["session"]["messages"] const { checkCompletionConditions } = await import("./completion") @@ -133,10 +138,9 @@ describe("checkCompletionConditions continuation coverage", () => { expect(result).toBe(false) }) - it("returns true when current session is only in lineage but is not a registered subagent", async () => { + it("returns true when current session is only in lineage and is not explicitly tracked in boulder", async () => { // given spyOn(console, "log").mockImplementation(() => {}) - registerAgentName("atlas") const directory = createTempDir() const planPath = join(directory, ".sisyphus", "plans", "lineage-non-subagent-plan.md") mkdirSync(join(directory, ".sisyphus", "plans"), { recursive: true }) @@ -151,6 +155,7 @@ describe("checkCompletionConditions continuation coverage", () => { parentID: path.id === "lineage-only-session" ? "root-session" : undefined, }, })) as unknown as RunContext["client"]["session"]["get"] + ctx.client.session.messages = mock(async () => ({ data: [] })) as unknown as RunContext["client"]["session"]["messages"] const { checkCompletionConditions } = await import("./completion") @@ -161,26 +166,288 @@ describe("checkCompletionConditions continuation coverage", () => { expect(result).toBe(true) }) - it("returns true when descendant subagent has agent mismatch and atlas would not continue it", async () => { + it("returns true when appended descendant has agent mismatch and atlas would not continue it", async () => { // given spyOn(console, "log").mockImplementation(() => {}) - registerAgentName("atlas") const directory = createTempDir() const planPath = join(directory, ".sisyphus", "plans", "lineage-agent-mismatch-plan.md") mkdirSync(join(directory, ".sisyphus", "plans"), { recursive: true }) writeFileSync(planPath, "- [ ] unfinished task\n", "utf-8") - writeBoulderStateFile(directory, planPath, ["root-session"]) + writeBoulderStateFile(directory, planPath, ["root-session", "mismatch-subagent-session"], { + "root-session": "direct", + "mismatch-subagent-session": "appended", + }) const ctx = createMockContext(directory) ctx.sessionID = "mismatch-subagent-session" - subagentSessions.add("mismatch-subagent-session") - setSessionAgent("mismatch-subagent-session", "sisyphus-junior") ctx.client.session.get = mock(async ({ path }: { path: { id: string } }) => ({ data: { id: path.id, parentID: path.id === "mismatch-subagent-session" ? "root-session" : undefined, }, })) as unknown as RunContext["client"]["session"]["get"] + ctx.client.session.messages = mock(async ({ path }: { path: { id: string } }) => ({ + data: path.id === "mismatch-subagent-session" + ? [{ info: { agent: "sisyphus-junior", providerID: "openai", modelID: "gpt-5.4" } }] + : [], + })) as unknown as RunContext["client"]["session"]["messages"] + + const { checkCompletionConditions } = await import("./completion") + + // when + const result = await checkCompletionConditions(ctx) + + // then + expect(result).toBe(true) + }) + + it("returns true when mismatched descendant was already appended into boulder session_ids", async () => { + // given + spyOn(console, "log").mockImplementation(() => {}) + const directory = createTempDir() + const planPath = join(directory, ".sisyphus", "plans", "appended-mismatch-plan.md") + mkdirSync(join(directory, ".sisyphus", "plans"), { recursive: true }) + writeFileSync(planPath, "- [ ] unfinished task\n", "utf-8") + writeBoulderStateFile(directory, planPath, ["root-session", "appended-mismatch-session"], { + "root-session": "direct", + "appended-mismatch-session": "appended", + }) + + const ctx = createMockContext(directory) + ctx.sessionID = "appended-mismatch-session" + ctx.client.session.get = mock(async ({ path }: { path: { id: string } }) => ({ + data: { + id: path.id, + parentID: path.id === "appended-mismatch-session" ? "root-session" : undefined, + }, + })) as unknown as RunContext["client"]["session"]["get"] + ctx.client.session.messages = mock(async ({ path }: { path: { id: string } }) => ({ + data: path.id === "appended-mismatch-session" + ? [{ info: { agent: "sisyphus-junior", providerID: "openai", modelID: "gpt-5.4" } }] + : [], + })) as unknown as RunContext["client"]["session"]["messages"] + + const { checkCompletionConditions } = await import("./completion") + + // when + const result = await checkCompletionConditions(ctx) + + // then + expect(result).toBe(true) + }) + + it("returns true when appended descendant cannot prove lineage because parent lookup fails", async () => { + // given + spyOn(console, "log").mockImplementation(() => {}) + const directory = createTempDir() + const planPath = join(directory, ".sisyphus", "plans", "appended-unresolved-lineage-plan.md") + mkdirSync(join(directory, ".sisyphus", "plans"), { recursive: true }) + writeFileSync(planPath, "- [ ] unfinished task\n", "utf-8") + writeBoulderStateFile(directory, planPath, ["root-session", "ses_appended_descendant"], { + "root-session": "direct", + "ses_appended_descendant": "appended", + }) + + const ctx = createMockContext(directory) + ctx.sessionID = "ses_appended_descendant" + ctx.client.session.get = mock(async () => { + throw new Error("session lookup failed") + }) as unknown as RunContext["client"]["session"]["get"] + ctx.client.session.messages = mock(async ({ path }: { path: { id: string } }) => ({ + data: path.id === "ses_appended_descendant" + ? [{ info: { agent: "atlas", providerID: "openai", modelID: "gpt-5.4" } }] + : [], + })) as unknown as RunContext["client"]["session"]["messages"] + + const { checkCompletionConditions } = await import("./completion") + + // when + const result = await checkCompletionConditions(ctx) + + // then + expect(result).toBe(true) + }) + + it("returns false when current session is directly tracked in boulder session_ids even if it has a parent session", async () => { + // given + spyOn(console, "log").mockImplementation(() => {}) + const directory = createTempDir() + const planPath = join(directory, ".sisyphus", "plans", "direct-tracked-child-plan.md") + mkdirSync(join(directory, ".sisyphus", "plans"), { recursive: true }) + writeFileSync(planPath, "- [ ] unfinished task\n", "utf-8") + writeBoulderStateFile(directory, planPath, ["ses_direct_child"]) + + const ctx = createMockContext(directory) + ctx.sessionID = "ses_direct_child" + ctx.client.session.get = mock(async ({ path }: { path: { id: string } }) => ({ + data: { + id: path.id, + parentID: path.id === "ses_direct_child" ? "ses_parent" : undefined, + }, + })) as unknown as RunContext["client"]["session"]["get"] + + const { checkCompletionConditions } = await import("./completion") + + // when + const result = await checkCompletionConditions(ctx) + + // then + expect(result).toBe(false) + }) + + it("returns false when current session is directly tracked among multiple boulder session_ids and has no parent session", async () => { + // given + spyOn(console, "log").mockImplementation(() => {}) + const directory = createTempDir() + const planPath = join(directory, ".sisyphus", "plans", "multi-tracked-direct-plan.md") + mkdirSync(join(directory, ".sisyphus", "plans"), { recursive: true }) + writeFileSync(planPath, "- [ ] unfinished task\n", "utf-8") + writeBoulderStateFile(directory, planPath, ["ses_other_tracked", "ses_direct_tracked"], { + "ses_other_tracked": "direct", + "ses_direct_tracked": "direct", + }) + + const ctx = createMockContext(directory) + ctx.sessionID = "ses_direct_tracked" + ctx.client.session.get = mock(async ({ path }: { path: { id: string } }) => ({ + data: { + id: path.id, + parentID: undefined, + }, + })) as unknown as RunContext["client"]["session"]["get"] + + const { checkCompletionConditions } = await import("./completion") + + // when + const result = await checkCompletionConditions(ctx) + + // then + expect(result).toBe(false) + }) + + it("returns true when multi-session tracked child is missing provenance and lineage cannot be proven", async () => { + // given + spyOn(console, "log").mockImplementation(() => {}) + const directory = createTempDir() + const planPath = join(directory, ".sisyphus", "plans", "unknown-origin-multi-session-plan.md") + mkdirSync(join(directory, ".sisyphus", "plans"), { recursive: true }) + writeFileSync(planPath, "- [ ] unfinished task\n", "utf-8") + writeBoulderStateFile(directory, planPath, ["ses_root_tracked", "ses_unknown_child"]) + + const ctx = createMockContext(directory) + ctx.sessionID = "ses_unknown_child" + ctx.client.session.get = mock(async () => { + throw new Error("lineage unavailable") + }) as unknown as RunContext["client"]["session"]["get"] + + const { checkCompletionConditions } = await import("./completion") + + // when + const result = await checkCompletionConditions(ctx) + + // then + expect(result).toBe(true) + }) + + it("returns false when directly tracked child session has a tracked ancestor and mismatched agent metadata", async () => { + // given + spyOn(console, "log").mockImplementation(() => {}) + const directory = createTempDir() + const planPath = join(directory, ".sisyphus", "plans", "multi-tracked-direct-child-plan.md") + mkdirSync(join(directory, ".sisyphus", "plans"), { recursive: true }) + writeFileSync(planPath, "- [ ] unfinished task\n", "utf-8") + writeBoulderStateFile(directory, planPath, ["ses_root_tracked", "ses_direct_child"], { + "ses_root_tracked": "direct", + "ses_direct_child": "direct", + }) + + const ctx = createMockContext(directory) + ctx.sessionID = "ses_direct_child" + ctx.client.session.get = mock(async ({ path }: { path: { id: string } }) => ({ + data: { + id: path.id, + parentID: path.id === "ses_direct_child" ? "ses_root_tracked" : undefined, + }, + })) as unknown as RunContext["client"]["session"]["get"] + ctx.client.session.messages = mock(async ({ path }: { path: { id: string } }) => ({ + data: path.id === "ses_direct_child" + ? [{ info: { agent: "sisyphus-junior", providerID: "openai", modelID: "gpt-5.4" } }] + : [], + })) as unknown as RunContext["client"]["session"]["messages"] + + const { checkCompletionConditions } = await import("./completion") + + // when + const result = await checkCompletionConditions(ctx) + + // then + expect(result).toBe(false) + }) + + it("returns false when latest appended descendant message is compaction but previous real agent still matches atlas", async () => { + // given + spyOn(console, "log").mockImplementation(() => {}) + const directory = createTempDir() + const planPath = join(directory, ".sisyphus", "plans", "compaction-descendant-plan.md") + mkdirSync(join(directory, ".sisyphus", "plans"), { recursive: true }) + writeFileSync(planPath, "- [ ] unfinished task\n", "utf-8") + writeBoulderStateFile(directory, planPath, ["root-session", "ses_child_after_compaction"], { + "root-session": "direct", + "ses_child_after_compaction": "appended", + }) + + const ctx = createMockContext(directory) + ctx.sessionID = "ses_child_after_compaction" + ctx.client.session.get = mock(async ({ path }: { path: { id: string } }) => ({ + data: { + id: path.id, + parentID: path.id === "ses_child_after_compaction" ? "root-session" : undefined, + }, + })) as unknown as RunContext["client"]["session"]["get"] + ctx.client.session.messages = mock(async ({ path }: { path: { id: string } }) => ({ + data: path.id === "ses_child_after_compaction" + ? [ + { info: { agent: "atlas", providerID: "openai", modelID: "gpt-5.4" } }, + { info: { agent: "compaction", providerID: "openai", modelID: "gpt-5.4" } }, + ] + : [], + })) as unknown as RunContext["client"]["session"]["messages"] + + const { checkCompletionConditions } = await import("./completion") + + // when + const result = await checkCompletionConditions(ctx) + + // then + expect(result).toBe(false) + }) + + it("returns true for untracked descendant continuation on SQLite-shaped misordered messages because lineage alone is no longer sufficient", async () => { + // given + spyOn(console, "log").mockImplementation(() => {}) + const directory = createTempDir() + const planPath = join(directory, ".sisyphus", "plans", "sqlite-ordered-descendant-plan.md") + mkdirSync(join(directory, ".sisyphus", "plans"), { recursive: true }) + writeFileSync(planPath, "- [ ] unfinished task\n", "utf-8") + writeBoulderStateFile(directory, planPath, ["root-session"]) + + const ctx = createMockContext(directory) + ctx.sessionID = "ses_sqlite_descendant" + ctx.client.session.get = mock(async ({ path }: { path: { id: string } }) => ({ + data: { + id: path.id, + parentID: path.id === "ses_sqlite_descendant" ? "root-session" : undefined, + }, + })) as unknown as RunContext["client"]["session"]["get"] + ctx.client.session.messages = mock(async ({ path }: { path: { id: string } }) => ({ + data: path.id === "ses_sqlite_descendant" + ? [ + { id: "msg_0001", info: { agent: "atlas", providerID: "openai", modelID: "gpt-5.4", time: { created: 100 } } }, + { id: "msg_0003", info: { agent: "compaction", providerID: "openai", modelID: "gpt-5.4", time: { created: 200 } } }, + { id: "msg_0002", info: { agent: "sisyphus-junior", providerID: "openai", modelID: "gpt-5.4", time: { created: 100 } } }, + ] + : [], + })) as unknown as RunContext["client"]["session"]["messages"] const { checkCompletionConditions } = await import("./completion") diff --git a/src/cli/run/continuation-state.json-backend.test.ts b/src/cli/run/continuation-state.json-backend.test.ts new file mode 100644 index 000000000..c2652538b --- /dev/null +++ b/src/cli/run/continuation-state.json-backend.test.ts @@ -0,0 +1,171 @@ +declare const require: (name: string) => any +const { afterEach, describe, expect, mock, test, afterAll } = require("bun:test") +import { mkdtempSync, mkdirSync, rmSync, writeFileSync } from "node:fs" +import { join } from "node:path" +import { tmpdir } from "node:os" + +const testDirs: string[] = [] + +const TEST_STORAGE_ROOT = join(tmpdir(), `omo-run-json-storage-${Date.now()}`) +const TEST_MESSAGE_STORAGE = join(TEST_STORAGE_ROOT, "message") + +mock.module("../../shared/opencode-storage-detection", () => ({ + isSqliteBackend: () => false, +})) + +mock.module("../../shared/opencode-message-dir", () => ({ + getMessageDir: (sessionID: string) => { + const directPath = join(TEST_MESSAGE_STORAGE, sessionID) + return require("node:fs").existsSync(directPath) ? directPath : null + }, +})) + +afterAll(() => { mock.restore() }) + +afterEach(() => { + while (testDirs.length > 0) { + const dir = testDirs.pop() + if (dir) { + rmSync(dir, { recursive: true, force: true }) + } + } +}) + +function createTempDir(): string { + const directory = mkdtempSync(join(tmpdir(), "omo-run-json-backend-")) + testDirs.push(directory) + return directory +} + +function writeJsonMessage(sessionID: string, fileName: string, agent: string): void { + const messageDir = join(TEST_MESSAGE_STORAGE, sessionID) + mkdirSync(messageDir, { recursive: true }) + writeFileSync( + join(messageDir, fileName), + JSON.stringify({ + agent, + model: { providerID: "openai", modelID: "gpt-5.4" }, + time: { created: fileName.includes("002") ? 200 : 100 }, + }), + "utf-8", + ) +} + +describe("getContinuationState JSON backend descendant coverage", () => { + test("returns active boulder for explicitly tracked appended descendant on JSON message storage backend", async () => { + // given + const directory = createTempDir() + const plansDir = join(directory, ".sisyphus", "plans") + mkdirSync(plansDir, { recursive: true }) + const planPath = join(plansDir, "json-descendant-plan.md") + writeFileSync(planPath, "- [ ] unfinished task\n", "utf-8") + mkdirSync(join(directory, ".sisyphus"), { recursive: true }) + writeFileSync(join(directory, ".sisyphus", "boulder.json"), JSON.stringify({ + active_plan: planPath, + started_at: new Date().toISOString(), + session_ids: ["ses_root_session", "ses_child_session"], + session_origins: { + "ses_root_session": "direct", + "ses_child_session": "appended", + }, + plan_name: "json-descendant-plan", + agent: "atlas", + }), "utf-8") + writeJsonMessage("ses_child_session", "msg_001.json", "atlas") + writeJsonMessage("ses_child_session", "msg_002.json", "compaction") + + const { getContinuationState } = await import("./continuation-state") + + // when + const state = await getContinuationState(directory, "ses_child_session", { + session: { + get: async ({ path }: { path: { id: string } }) => ({ + data: { + id: path.id, + parentID: path.id === "ses_child_session" ? "ses_root_session" : undefined, + }, + }), + }, + } as never) + + // then + expect(state.hasActiveBoulder).toBe(true) + }) + + test("prefers earliest JSON agent by time.created instead of filename order for first-message fallback helpers", async () => { + // given + const directory = createTempDir() + const sessionID = "ses_json_first_agent" + writeJsonMessage(sessionID, "msg_ffff0000_000001.json", "later-agent") + writeFileSync( + join(TEST_MESSAGE_STORAGE, sessionID, "msg_00000000_000999.json"), + JSON.stringify({ + agent: "earliest-agent", + model: { providerID: "openai", modelID: "gpt-5.4" }, + time: { created: 10 }, + }), + "utf-8", + ) + + const { findFirstMessageWithAgent } = await import("../../features/hook-message-injector") + + // when + const result = findFirstMessageWithAgent(join(TEST_MESSAGE_STORAGE, sessionID)) + + // then + expect(result).toBe("earliest-agent") + rmSync(directory, { recursive: true, force: true }) + }) + + test("prefers newest JSON agent by time.created even when filenames look reversed and timestamps tie-break by filename only", async () => { + // given + const directory = createTempDir() + const plansDir = join(directory, ".sisyphus", "plans") + mkdirSync(plansDir, { recursive: true }) + const planPath = join(plansDir, "json-random-id-plan.md") + writeFileSync(planPath, "- [ ] unfinished task\n", "utf-8") + mkdirSync(join(directory, ".sisyphus"), { recursive: true }) + writeFileSync(join(directory, ".sisyphus", "boulder.json"), JSON.stringify({ + active_plan: planPath, + started_at: new Date().toISOString(), + session_ids: ["ses_root_random"], + plan_name: "json-random-id-plan", + agent: "atlas", + }), "utf-8") + const sessionID = "ses_child_random" + const messageDir = join(TEST_MESSAGE_STORAGE, sessionID) + mkdirSync(messageDir, { recursive: true }) + writeFileSync(join(messageDir, "msg_a91f00ab_000001.json"), JSON.stringify({ + agent: "atlas", + model: { providerID: "openai", modelID: "gpt-5.4" }, + time: { created: 100 }, + }), "utf-8") + writeFileSync(join(messageDir, "msg_f0e1d2c3_000002.json"), JSON.stringify({ + agent: "compaction", + model: { providerID: "openai", modelID: "gpt-5.4" }, + time: { created: 200 }, + }), "utf-8") + writeFileSync(join(messageDir, "msg_d4c3b2a1_000003.json"), JSON.stringify({ + agent: "sisyphus-junior", + model: { providerID: "openai", modelID: "gpt-5.4" }, + time: { created: 100 }, + }), "utf-8") + + const { getContinuationState } = await import("./continuation-state") + + // when + const state = await getContinuationState(directory, sessionID, { + session: { + get: async ({ path }: { path: { id: string } }) => ({ + data: { + id: path.id, + parentID: path.id === sessionID ? "ses_root_random" : undefined, + }, + }), + }, + } as never) + + // then + expect(state.hasActiveBoulder).toBe(false) + }) +}) diff --git a/src/cli/run/continuation-state.ts b/src/cli/run/continuation-state.ts index a9645b6ee..d827af5e3 100644 --- a/src/cli/run/continuation-state.ts +++ b/src/cli/run/continuation-state.ts @@ -1,15 +1,11 @@ import { getPlanProgress, readBoulderState } from "../../features/boulder-state" -import { - getSessionAgent, - isAgentRegistered, - subagentSessions, -} from "../../features/claude-code-session-state" import { getActiveContinuationMarkerReason, isContinuationMarkerActive, readContinuationMarker, } from "../../features/run-continuation-state" import { isSessionInBoulderLineage } from "../../hooks/atlas/boulder-session-lineage" +import { getLastAgentFromSession } from "../../hooks/atlas/session-last-agent" import { getAgentConfigKey } from "../../shared/agent-display-names" import { readState as readRalphLoopState } from "../../hooks/ralph-loop/storage" import type { RunContext } from "./types" @@ -50,30 +46,55 @@ async function hasActiveBoulderContinuation( const progress = getPlanProgress(boulder.active_plan) if (progress.isComplete) return false - if (boulder.session_ids.includes(sessionID)) return true if (!client) return false - if (!subagentSessions.has(sessionID)) return false - const requiredAgentName = boulder.agent ?? (isAgentRegistered("atlas") ? "atlas" : undefined) - if (!requiredAgentName || !isAgentRegistered(requiredAgentName)) { + const isTrackedSession = boulder.session_ids.includes(sessionID) + const sessionOrigin = boulder.session_origins?.[sessionID] + if (!isTrackedSession) { return false } - const sessionAgent = getSessionAgent(sessionID) - const agentKey = getAgentConfigKey(sessionAgent ?? "") - const requiredAgentKey = getAgentConfigKey(requiredAgentName) - const isEligibleSubagent = - agentKey === requiredAgentKey - || (requiredAgentKey === getAgentConfigKey("atlas") && agentKey === getAgentConfigKey("sisyphus")) + const isTrackedDescendant = await isTrackedDescendantSession(client, sessionID, boulder.session_ids) - if (!isEligibleSubagent) { + if (isTrackedSession && sessionOrigin === "direct") { + return true + } + + if (isTrackedSession && sessionOrigin !== "direct" && !isTrackedDescendant) { + return false + } + + const sessionAgent = await getLastAgentFromSession(sessionID, client) + if (!sessionAgent) { + return false + } + + const requiredAgentKey = getAgentConfigKey(boulder.agent ?? "atlas") + const sessionAgentKey = getAgentConfigKey(sessionAgent) + if ( + sessionAgentKey !== requiredAgentKey + && !(requiredAgentKey === getAgentConfigKey("atlas") && sessionAgentKey === getAgentConfigKey("sisyphus")) + ) { + return false + } + + return isTrackedSession || isTrackedDescendant +} + +async function isTrackedDescendantSession( + client: RunContext["client"], + sessionID: string, + trackedSessionIDs: string[], +): Promise { + const ancestorSessionIDs = trackedSessionIDs.filter((trackedSessionID) => trackedSessionID !== sessionID) + if (ancestorSessionIDs.length === 0) { return false } return isSessionInBoulderLineage({ client, sessionID, - boulderSessionIDs: boulder.session_ids, + boulderSessionIDs: ancestorSessionIDs, }) }