feat(run): integrate session origins and agent detection into continuation state

- 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

@@ -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<string, "direct" | "appended">,
+): 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")
 

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

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)
+  const isTrackedSession = boulder.session_ids.includes(sessionID)
-  if (!requiredAgentName || !isAgentRegistered(requiredAgentName)) {
+  const sessionOrigin = boulder.session_origins?.[sessionID]
+  if (!isTrackedSession) {
     return false
   }
 
-  const sessionAgent = getSessionAgent(sessionID)
+  const isTrackedDescendant = await isTrackedDescendantSession(client, sessionID, boulder.session_ids)
-  const agentKey = getAgentConfigKey(sessionAgent ?? "")
-  const requiredAgentKey = getAgentConfigKey(requiredAgentName)
-  const isEligibleSubagent =
-    agentKey === requiredAgentKey
-    || (requiredAgentKey === getAgentConfigKey("atlas") && agentKey === getAgentConfigKey("sisyphus"))
 
-  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<boolean> {
+  const ancestorSessionIDs = trackedSessionIDs.filter((trackedSessionID) => trackedSessionID !== sessionID)
+  if (ancestorSessionIDs.length === 0) {
     return false
   }
 
   return isSessionInBoulderLineage({
     client,
     sessionID,
-    boulderSessionIDs: boulder.session_ids,
+    boulderSessionIDs: ancestorSessionIDs,
   })
 }