feat(atlas): persist preferred task session reuse

2026-03-17 17:25:46 +09:00
parent 1070b9170f
commit 5c6194372e
14 changed files with 665 additions and 18 deletions
@@ -7,6 +7,7 @@ import type { AtlasHookOptions, SessionState } from "./types"
 export function createAtlasHook(ctx: PluginInput, options?: AtlasHookOptions) {
  const sessions = new Map<string, SessionState>()
  const pendingFilePaths = new Map<string, string>()
+  const pendingTaskRefs = new Map<string, { key: string; label: string; title: string } | null>()
  const autoCommit = options?.autoCommit ?? true

  function getState(sessionID: string): SessionState {
@@ -20,7 +21,7 @@ export function createAtlasHook(ctx: PluginInput, options?: AtlasHookOptions) {

  return {
    handler: createAtlasEventHandler({ ctx, options, sessions, getState }),
-    "tool.execute.before": createToolExecuteBeforeHandler({ ctx, pendingFilePaths }),
-    "tool.execute.after": createToolExecuteAfterHandler({ ctx, pendingFilePaths, autoCommit, getState }),
+    "tool.execute.before": createToolExecuteBeforeHandler({ ctx, pendingFilePaths, pendingTaskRefs }),
+    "tool.execute.after": createToolExecuteAfterHandler({ ctx, pendingFilePaths, pendingTaskRefs, autoCommit, getState }),
  }
 }
@@ -15,6 +15,8 @@ export async function injectBoulderContinuation(input: {
  total: number
  agent?: string
  worktreePath?: string
+  preferredTaskSessionId?: string
+  preferredTaskTitle?: string
  backgroundManager?: BackgroundManager
  sessionState: SessionState
 }): Promise<void> {
@@ -26,6 +28,8 @@ export async function injectBoulderContinuation(input: {
    total,
    agent,
    worktreePath,
+    preferredTaskSessionId,
+    preferredTaskTitle,
    backgroundManager,
    sessionState,
  } = input
@@ -40,9 +44,13 @@ export async function injectBoulderContinuation(input: {
  }

  const worktreeContext = worktreePath ? `\n\n[Worktree: ${worktreePath}]` : ""
+  const preferredSessionContext = preferredTaskSessionId
+    ? `\n\n[Preferred reuse session for current top-level plan task${preferredTaskTitle ? `: ${preferredTaskTitle}` : ""}: ${preferredTaskSessionId}]`
+    : ""
  const prompt =
    BOULDER_CONTINUATION_PROMPT.replace(/{PLAN_NAME}/g, planName) +
    `\n\n[Status: ${total - remaining}/${total} completed, ${remaining} remaining]` +
+    preferredSessionContext +
    worktreeContext

  try {
@@ -1,5 +1,10 @@
 import type { PluginInput } from "@opencode-ai/plugin"
-import { getPlanProgress, readBoulderState } from "../../features/boulder-state"
+import {
+  getPlanProgress,
+  getTaskSessionState,
+  readBoulderState,
+  readCurrentTopLevelTask,
+} from "../../features/boulder-state"
 import { log } from "../../shared/logger"
 import { injectBoulderContinuation } from "./boulder-continuation-injector"
 import { HOOK_NAME } from "./hook-name"
@@ -31,6 +36,14 @@ async function injectContinuation(input: {
  input.sessionState.lastContinuationInjectedAt = Date.now()

  try {
+    const currentBoulder = readBoulderState(input.ctx.directory)
+    const currentTask = currentBoulder
+      ? readCurrentTopLevelTask(currentBoulder.active_plan)
+      : null
+    const preferredTaskSession = currentTask
+      ? getTaskSessionState(input.ctx.directory, currentTask.key)
+      : null
+
    await injectBoulderContinuation({
      ctx: input.ctx,
      sessionID: input.sessionID,
@@ -39,6 +52,8 @@ async function injectContinuation(input: {
      total: input.progress.total,
      agent: input.agent,
      worktreePath: input.worktreePath,
+      preferredTaskSessionId: preferredTaskSession?.session_id,
+      preferredTaskTitle: preferredTaskSession?.task_title,
      backgroundManager: input.options?.backgroundManager,
      sessionState: input.sessionState,
    })
@@ -404,12 +404,174 @@ describe("atlas hook", () => {

      // then - should include verification instructions
      expect(output.output).toContain("LYING")
-      expect(output.output).toContain("PHASE 1")
-      expect(output.output).toContain("PHASE 2")
+     expect(output.output).toContain("PHASE 1")
+     expect(output.output).toContain("PHASE 2")
      
      cleanupMessageStorage(sessionID)
    })

+     test("should persist preferred subagent session for the current top-level task", async () => {
+       // given - boulder state with a current top-level task, Atlas caller
+       const sessionID = "session-task-session-track-test"
+       setupMessageStorage(sessionID, "atlas")
+
+      const planPath = join(TEST_DIR, "task-session-plan.md")
+      writeFileSync(planPath, `# Plan
+
+## TODOs
+- [ ] 1. Implement auth flow
+  - [ ] nested acceptance checkbox
+`)
+
+      const state: BoulderState = {
+        active_plan: planPath,
+        started_at: "2026-01-02T10:00:00Z",
+        session_ids: ["session-1"],
+        plan_name: "task-session-plan",
+      }
+      writeBoulderState(TEST_DIR, state)
+
+      const hook = createAtlasHook(createMockPluginInput())
+      const output = {
+        title: "Sisyphus Task",
+        output: `Task completed successfully
+
+<task_metadata>
+session_id: ses_auth_flow_123
+</task_metadata>`,
+        metadata: {
+          agent: "sisyphus-junior",
+          category: "deep",
+        },
+      }
+
+      // when
+      await hook["tool.execute.after"](
+        { tool: "task", sessionID },
+        output
+      )
+
+      // then
+     const updatedState = readBoulderState(TEST_DIR)
+      expect(updatedState?.task_sessions?.["todo:1"]?.session_id).toBe("ses_auth_flow_123")
+      expect(updatedState?.task_sessions?.["todo:1"]?.task_title).toBe("Implement auth flow")
+      expect(updatedState?.task_sessions?.["todo:1"]?.agent).toBe("sisyphus-junior")
+      expect(updatedState?.task_sessions?.["todo:1"]?.category).toBe("deep")
+
+      cleanupMessageStorage(sessionID)
+    })
+
+     test("should preserve the delegated task key even after the plan advances to the next task", async () => {
+       // given - Atlas caller starts task 1, then the plan advances before task output is processed
+       const sessionID = "session-stable-task-key-test"
+       setupMessageStorage(sessionID, "atlas")
+
+      const planPath = join(TEST_DIR, "stable-task-key-plan.md")
+      writeFileSync(planPath, `# Plan
+
+## TODOs
+- [ ] 1. Implement auth flow
+- [ ] 2. Add API validation
+`)
+
+      writeBoulderState(TEST_DIR, {
+        active_plan: planPath,
+        started_at: "2026-01-02T10:00:00Z",
+        session_ids: ["session-1"],
+        plan_name: "stable-task-key-plan",
+      })
+
+      const hook = createAtlasHook(createMockPluginInput())
+
+      // when - Atlas delegates task 1
+      await hook["tool.execute.before"](
+        { tool: "task", sessionID, callID: "call-task-1" },
+        { args: { prompt: "Implement auth flow" } }
+      )
+
+      // and the plan is advanced before the task output is processed
+      writeFileSync(planPath, `# Plan
+
+## TODOs
+- [x] 1. Implement auth flow
+- [ ] 2. Add API validation
+`)
+
+      await hook["tool.execute.after"](
+        { tool: "task", sessionID, callID: "call-task-1" },
+        {
+          title: "Sisyphus Task",
+          output: `Task completed successfully
+
+<task_metadata>
+session_id: ses_auth_flow_123
+</task_metadata>`,
+          metadata: {
+            agent: "sisyphus-junior",
+            category: "deep",
+          },
+        }
+      )
+
+      // then - the completed task session is still recorded against task 1, not task 2
+     const updatedState = readBoulderState(TEST_DIR)
+      expect(updatedState?.task_sessions?.["todo:1"]?.session_id).toBe("ses_auth_flow_123")
+      expect(updatedState?.task_sessions?.["todo:2"]).toBeUndefined()
+
+      cleanupMessageStorage(sessionID)
+    })
+
+     test("should not overwrite the current task mapping when task() explicitly resumes an older session", async () => {
+       // given - current plan is on task 2, but Atlas explicitly resumes an older session for a previous task
+       const sessionID = "session-cross-task-resume-test"
+       setupMessageStorage(sessionID, "atlas")
+
+      const planPath = join(TEST_DIR, "cross-task-resume-plan.md")
+      writeFileSync(planPath, `# Plan
+
+## TODOs
+- [x] 1. Implement auth flow
+- [ ] 2. Add API validation
+`)
+
+      writeBoulderState(TEST_DIR, {
+        active_plan: planPath,
+        started_at: "2026-01-02T10:00:00Z",
+        session_ids: ["session-1"],
+        plan_name: "cross-task-resume-plan",
+      })
+
+      const hook = createAtlasHook(createMockPluginInput())
+
+      // when - Atlas resumes an explicit prior session
+      await hook["tool.execute.before"](
+        { tool: "task", sessionID, callID: "call-resume-old-task" },
+        { args: { prompt: "Follow up on previous task", session_id: "ses_old_task_111" } }
+      )
+
+      await hook["tool.execute.after"](
+        { tool: "task", sessionID, callID: "call-resume-old-task" },
+        {
+          title: "Sisyphus Task",
+          output: `Task continued successfully
+
+<task_metadata>
+session_id: ses_old_task_111
+</task_metadata>`,
+          metadata: {
+            agent: "sisyphus-junior",
+            category: "deep",
+          },
+        }
+      )
+
+      // then - Atlas does not poison task 2's preferred session mapping
+      const updatedState = readBoulderState(TEST_DIR)
+      expect(updatedState?.task_sessions?.["todo:2"]).toBeUndefined()
+
+      cleanupMessageStorage(sessionID)
+    })
+
    describe("completion gate output ordering", () => {
      const COMPLETION_GATE_SESSION = "completion-gate-order-test"

@@ -1147,6 +1309,48 @@ describe("atlas hook", () => {
      expect(callArgs.body.parts[0].text).toContain("2 remaining")
    })

+    test("should include preferred reuse session in continuation prompt for current top-level task", async () => {
+      // given - boulder state with tracked preferred session
+      const planPath = join(TEST_DIR, "preferred-session-plan.md")
+      writeFileSync(planPath, `# Plan
+
+## TODOs
+- [ ] 1. Implement auth flow
+`)
+
+      writeBoulderState(TEST_DIR, {
+        active_plan: planPath,
+        started_at: "2026-01-02T10:00:00Z",
+        session_ids: [MAIN_SESSION_ID],
+        plan_name: "preferred-session-plan",
+        task_sessions: {
+          "todo:1": {
+            task_key: "todo:1",
+            task_label: "1",
+            task_title: "Implement auth flow",
+            session_id: "ses_auth_flow_123",
+            updated_at: "2026-01-02T10:00:00Z",
+          },
+        },
+      })
+
+      const mockInput = createMockPluginInput()
+      const hook = createAtlasHook(mockInput)
+
+      // when
+      await hook.handler({
+        event: {
+          type: "session.idle",
+          properties: { sessionID: MAIN_SESSION_ID },
+        },
+      })
+
+      // then
+      const callArgs = mockInput._promptMock.mock.calls[0][0]
+      expect(callArgs.body.parts[0].text).toContain("Preferred reuse session for current top-level plan task")
+      expect(callArgs.body.parts[0].text).toContain("ses_auth_flow_123")
+    })
+
    test("should inject when last agent is sisyphus and boulder targets atlas explicitly", async () => {
       // given - boulder explicitly set to atlas, but last agent is sisyphus (initial state after /start-work)
       const planPath = join(TEST_DIR, "test-plan.md")
@@ -0,0 +1,53 @@
+import { describe, expect, test } from "bun:test"
+
+import { extractSessionIdFromOutput } from "./subagent-session-id"
+
+describe("extractSessionIdFromOutput", () => {
+  test("extracts Session ID blocks from background output", () => {
+    // given
+    const output = `Background task launched.\n\nSession ID: ses_bg_12345`
+
+    // when
+    const result = extractSessionIdFromOutput(output)
+
+    // then
+    expect(result).toBe("ses_bg_12345")
+  })
+
+  test("extracts session_id from task metadata blocks", () => {
+    // given
+    const output = `Task completed.\n\n<task_metadata>\nsession_id: ses_sync_12345\n</task_metadata>`
+
+    // when
+    const result = extractSessionIdFromOutput(output)
+
+    // then
+    expect(result).toBe("ses_sync_12345")
+  })
+
+  test("returns undefined when no session id is present", () => {
+    // given
+    const output = "Task completed without metadata"
+
+    // when
+    const result = extractSessionIdFromOutput(output)
+
+    // then
+    expect(result).toBeUndefined()
+  })
+
+  test("prefers the session id inside the trailing task_metadata block", () => {
+    // given
+    const output = `The previous attempt mentioned session_id: ses_wrong_body_123 but that was only context.
+
+<task_metadata>
+session_id: ses_real_metadata_456
+</task_metadata>`
+
+    // when
+    const result = extractSessionIdFromOutput(output)
+
+    // then
+    expect(result).toBe("ses_real_metadata_456")
+  })
+})
@@ -1,4 +1,10 @@
-export function extractSessionIdFromOutput(output: string): string {
-  const match = output.match(/Session ID:\s*(ses_[a-zA-Z0-9]+)/)
-  return match?.[1] ?? "<session_id>"
+export function extractSessionIdFromOutput(output: string): string | undefined {
+  const taskMetadataMatches = [...output.matchAll(/<task_metadata>[\s\S]*?session_id:\s*(ses_[a-zA-Z0-9_]+)[\s\S]*?<\/task_metadata>/gi)]
+  const lastTaskMetadataMatch = taskMetadataMatches.at(-1)
+  if (lastTaskMetadataMatch) {
+    return lastTaskMetadataMatch[1]
+  }
+
+  const explicitSessionMatches = [...output.matchAll(/Session ID:\s*(ses_[a-zA-Z0-9_]+)/g)]
+  return explicitSessionMatches.at(-1)?.[1]
 }
@@ -1,5 +1,12 @@
 import type { PluginInput } from "@opencode-ai/plugin"
-import { appendSessionId, getPlanProgress, readBoulderState } from "../../features/boulder-state"
+import {
+  appendSessionId,
+  getPlanProgress,
+  getTaskSessionState,
+  readBoulderState,
+  readCurrentTopLevelTask,
+  upsertTaskSessionState,
+} from "../../features/boulder-state"
 import { log } from "../../shared/logger"
 import { isCallerOrchestrator } from "../../shared/session-utils"
 import { collectGitDiffStats, formatFileChanges } from "../../shared/git-worktree"
@@ -18,13 +25,18 @@ import { isWriteOrEditToolName } from "./write-edit-tool-policy"
 import type { SessionState } from "./types"
 import type { ToolExecuteAfterInput, ToolExecuteAfterOutput } from "./types"

+function resolvePreferredSessionId(currentSessionId?: string, trackedSessionId?: string): string {
+  return currentSessionId ?? trackedSessionId ?? "<session_id>"
+}
+
 export function createToolExecuteAfterHandler(input: {
  ctx: PluginInput
  pendingFilePaths: Map<string, string>
+  pendingTaskRefs: Map<string, { key: string; label: string; title: string } | null>
  autoCommit: boolean
  getState: (sessionID: string) => SessionState
 }): (toolInput: ToolExecuteAfterInput, toolOutput: ToolExecuteAfterOutput) => Promise<void> {
-  const { ctx, pendingFilePaths, autoCommit, getState } = input
+  const { ctx, pendingFilePaths, pendingTaskRefs, autoCommit, getState } = input
  return async (toolInput, toolOutput): Promise<void> => {
    // Guard against undefined output (e.g., from /review command - see issue #1035)
    if (!toolOutput) {
@@ -68,10 +80,21 @@ export function createToolExecuteAfterHandler(input: {
      const gitStats = collectGitDiffStats(ctx.directory)
      const fileChanges = formatFileChanges(gitStats)
      const subagentSessionId = extractSessionIdFromOutput(toolOutput.output)
+      const pendingTaskRef = toolInput.callID ? pendingTaskRefs.get(toolInput.callID) : undefined
+      if (toolInput.callID) {
+        pendingTaskRefs.delete(toolInput.callID)
+      }

      const boulderState = readBoulderState(ctx.directory)
      if (boulderState) {
        const progress = getPlanProgress(boulderState.active_plan)
+        const shouldSkipTaskSessionUpdate = pendingTaskRef === null
+        const currentTask = shouldSkipTaskSessionUpdate
+          ? null
+          : pendingTaskRef ?? readCurrentTopLevelTask(boulderState.active_plan)
+        const trackedTaskSession = currentTask
+          ? getTaskSessionState(ctx.directory, currentTask.key)
+          : null
        const sessionState = toolInput.sessionID ? getState(toolInput.sessionID) : undefined

        if (toolInput.sessionID && !boulderState.session_ids?.includes(toolInput.sessionID)) {
@@ -82,6 +105,22 @@ export function createToolExecuteAfterHandler(input: {
          })
        }

+        if (currentTask && subagentSessionId) {
+          upsertTaskSessionState(ctx.directory, {
+            taskKey: currentTask.key,
+            taskLabel: currentTask.label,
+            taskTitle: currentTask.title,
+            sessionId: subagentSessionId,
+            agent: toolOutput.metadata?.agent as string | undefined,
+            category: toolOutput.metadata?.category as string | undefined,
+          })
+        }
+
+        const preferredSessionId = resolvePreferredSessionId(
+          subagentSessionId,
+          trackedTaskSession?.session_id,
+        )
+
        // Preserve original subagent response - critical for debugging failed tasks
        const originalResponse = toolOutput.output
        const shouldPauseForApproval = sessionState
@@ -102,11 +141,11 @@ export function createToolExecuteAfterHandler(input: {
        }

        const leadReminder = shouldPauseForApproval
-          ? buildFinalWaveApprovalReminder(boulderState.plan_name, progress, subagentSessionId)
-          : buildCompletionGate(boulderState.plan_name, subagentSessionId)
+          ? buildFinalWaveApprovalReminder(boulderState.plan_name, progress, preferredSessionId)
+          : buildCompletionGate(boulderState.plan_name, preferredSessionId)
        const followupReminder = shouldPauseForApproval
          ? null
-          : buildOrchestratorReminder(boulderState.plan_name, progress, subagentSessionId, autoCommit, false)
+          : buildOrchestratorReminder(boulderState.plan_name, progress, preferredSessionId, autoCommit, false)

        toolOutput.output = `
 <system-reminder>
@@ -132,10 +171,13 @@ ${
          plan: boulderState.plan_name,
          progress: `${progress.completed}/${progress.total}`,
          fileCount: gitStats.length,
+          preferredSessionId,
          waitingForFinalWaveApproval: shouldPauseForApproval,
        })
      } else {
-        toolOutput.output += `\n<system-reminder>\n${buildStandaloneVerificationReminder(subagentSessionId)}\n</system-reminder>`
+        toolOutput.output += `\n<system-reminder>\n${buildStandaloneVerificationReminder(
+          resolvePreferredSessionId(subagentSessionId),
+        )}\n</system-reminder>`

        log(`[${HOOK_NAME}] Verification reminder appended for orchestrator`, {
          sessionID: toolInput.sessionID,
@@ -2,6 +2,7 @@ import { log } from "../../shared/logger"
 import { SYSTEM_DIRECTIVE_PREFIX } from "../../shared/system-directive"
 import { isCallerOrchestrator } from "../../shared/session-utils"
 import type { PluginInput } from "@opencode-ai/plugin"
+import { readBoulderState, readCurrentTopLevelTask } from "../../features/boulder-state"
 import { HOOK_NAME } from "./hook-name"
 import { ORCHESTRATOR_DELEGATION_REQUIRED, SINGLE_TASK_DIRECTIVE } from "./system-reminder-templates"
 import { isSisyphusPath } from "./sisyphus-path"
@@ -10,11 +11,12 @@ import { isWriteOrEditToolName } from "./write-edit-tool-policy"
 export function createToolExecuteBeforeHandler(input: {
  ctx: PluginInput
  pendingFilePaths: Map<string, string>
+  pendingTaskRefs: Map<string, { key: string; label: string; title: string } | null>
 }): (
  toolInput: { tool: string; sessionID?: string; callID?: string },
  toolOutput: { args: Record<string, unknown>; message?: string }
 ) => Promise<void> {
-  const { ctx, pendingFilePaths } = input
+  const { ctx, pendingFilePaths, pendingTaskRefs } = input

  return async (toolInput, toolOutput): Promise<void> => {
    if (!(await isCallerOrchestrator(toolInput.sessionID, ctx.client))) {
@@ -43,6 +45,25 @@ export function createToolExecuteBeforeHandler(input: {

    // Check task - inject single-task directive
    if (toolInput.tool === "task") {
+      if (toolInput.callID) {
+        const requestedSessionId = toolOutput.args.session_id as string | undefined
+        if (requestedSessionId) {
+          pendingTaskRefs.set(toolInput.callID, null)
+        } else {
+          const boulderState = readBoulderState(ctx.directory)
+          const currentTask = boulderState
+            ? readCurrentTopLevelTask(boulderState.active_plan)
+            : null
+          if (currentTask) {
+            pendingTaskRefs.set(toolInput.callID, {
+              key: currentTask.key,
+              label: currentTask.label,
+              title: currentTask.title,
+            })
+          }
+        }
+      }
+
      const prompt = toolOutput.args.prompt as string | undefined
      if (prompt && !prompt.includes(SYSTEM_DIRECTIVE_PREFIX)) {
        toolOutput.args.prompt = `<system-reminder>${SINGLE_TASK_DIRECTIVE}</system-reminder>\n` + prompt
@@ -1,5 +1,14 @@
 import { VERIFICATION_REMINDER } from "./system-reminder-templates"

+function buildReuseHint(sessionId: string): string {
+  return `
+**PREFERRED REUSE SESSION FOR THE CURRENT TOP-LEVEL PLAN TASK**
+
+- Reuse \`${sessionId}\` first if verification fails or the result needs follow-up.
+- Start a fresh subagent session only when reuse is unavailable or would cross task boundaries.
+`
+}
+
 export function buildCompletionGate(planName: string, sessionId: string): string {
  return `
 **COMPLETION GATE — DO NOT PROCEED UNTIL THIS IS DONE**
@@ -25,7 +34,8 @@ task(session_id="${sessionId}", prompt="fix: checkbox not recorded correctly")

 **Your completion is NOT tracked until the checkbox is marked in the plan file.**

-**VERIFICATION_REMINDER**`
+**VERIFICATION_REMINDER**
+${buildReuseHint(sessionId)}`
 }

 function buildVerificationReminder(sessionId: string): string {
@@ -38,7 +48,9 @@ ${VERIFICATION_REMINDER}
 **If ANY verification fails, use this immediately:**
 \`\`\`
 task(session_id="${sessionId}", prompt="fix: [describe the specific failure]")
-\`\`\``
+\`\`\`
+
+${buildReuseHint(sessionId)}`
 }

 export function buildOrchestratorReminder(