feat(hooks/atlas): wire per-task timers via startTaskTimer/endTaskTimer

src/hooks/atlas/tool-execute-after-task-timers.test.ts

@@ -0,0 +1,222 @@
+/// <reference types="bun-types" />
+
+import { afterAll, afterEach, beforeEach, describe, expect, it, mock, spyOn } from "bun:test"
+import { existsSync, mkdirSync, rmSync, writeFileSync } from "node:fs"
+import { tmpdir } from "node:os"
+import { join } from "node:path"
+import type { PluginInput } from "@opencode-ai/plugin"
+import type { Project } from "@opencode-ai/sdk"
+import { readBoulderState, writeBoulderState } from "../../features/boulder-state"
+import { createToolExecuteBeforeHandler } from "./tool-execute-before"
+
+const isCallerOrchestratorMock = mock(async () => true)
+const collectGitDiffStatsMock = mock(() => ({
+  filesChanged: 0,
+  insertions: 0,
+  deletions: 0,
+}))
+
+mock.module("../../shared/session-utils", () => ({
+  isCallerOrchestrator: isCallerOrchestratorMock,
+}))
+
+mock.module("../../shared/git-worktree", () => ({
+  collectGitDiffStats: collectGitDiffStatsMock,
+  formatFileChanges: mock(() => "No file changes"),
+}))
+
+afterAll(() => { mock.restore() })
+
+const { createToolExecuteAfterHandler } = await import("./tool-execute-after")
+
+type SessionGetInput = { path: { id: string } }
+type SessionGetResult = {
+  data: { parentID: string | undefined }
+  error?: undefined
+  request: Request
+  response: Response
+}
+
+describe("createToolExecuteAfterHandler task timers", () => {
+  let testDirectory = ""
+
+  beforeEach(() => {
+    testDirectory = join(tmpdir(), `atlas-task-timers-${crypto.randomUUID()}`)
+    if (!existsSync(testDirectory)) {
+      mkdirSync(testDirectory, { recursive: true })
+    }
+    isCallerOrchestratorMock.mockClear()
+    collectGitDiffStatsMock.mockClear()
+  })
+
+  afterEach(() => {
+    if (testDirectory && existsSync(testDirectory)) {
+      rmSync(testDirectory, { recursive: true, force: true })
+    }
+  })
+
+  function createProject(): Project {
+    return {
+      id: "project-1",
+      worktree: testDirectory,
+      time: { created: Date.now() },
+    }
+  }
+
+  function createSessionGetResult(parentID: string | undefined): SessionGetResult {
+    return {
+      data: { parentID },
+      error: undefined,
+      request: new Request("https://example.com/session"),
+      response: new Response(null, { status: 200 }),
+    } as SessionGetResult
+  }
+
+  function createHandlers(parentSessionIDs?: Record<string, string | undefined>) {
+    const project = createProject()
+    const client = {
+      session: {
+        get: async (input: SessionGetInput) => createSessionGetResult(parentSessionIDs?.[input.path.id]),
+      },
+    } as unknown as PluginInput["client"]
+
+    if (parentSessionIDs) {
+      spyOn(client.session, "get").mockImplementation((input) => Promise.resolve(
+        createSessionGetResult(parentSessionIDs[input?.path?.id ?? ""]),
+      ) as never)
+    }
+
+    const pendingFilePaths = new Map<string, string>()
+    const pendingTaskRefs = new Map()
+    const ctx = {
+      client,
+      project,
+      directory: testDirectory,
+      worktree: testDirectory,
+      serverUrl: new URL("https://example.com"),
+      $: Bun.$,
+    } satisfies PluginInput
+
+    return {
+      beforeHandler: createToolExecuteBeforeHandler({ ctx, pendingFilePaths, pendingTaskRefs }),
+      afterHandler: createToolExecuteAfterHandler({
+        ctx,
+        pendingFilePaths,
+        pendingTaskRefs,
+        autoCommit: true,
+        getState: () => ({ promptFailureCount: 0 }),
+      }),
+    }
+  }
+
+  it("starts task timer for todo:1 when delegated task session is tracked", async () => {
+    // given
+    const parentSessionID = "ses_parent"
+    const childSessionID = "ses_child"
+    const planPath = join(testDirectory, "task-timer-plan.md")
+    writeFileSync(planPath, "# Plan\n\n## TODOs\n- [ ] 1. Implement auth flow\n", "utf-8")
+    writeBoulderState(testDirectory, {
+      schema_version: 2,
+      active_work_id: "work-1",
+      active_plan: planPath,
+      started_at: "2026-01-02T10:00:00Z",
+      session_ids: [parentSessionID],
+      plan_name: "task-timer-plan",
+      works: {
+        "work-1": {
+          work_id: "work-1",
+          active_plan: planPath,
+          plan_name: "task-timer-plan",
+          started_at: "2026-01-02T10:00:00Z",
+          session_ids: [parentSessionID],
+          status: "active",
+        },
+      },
+    })
+    const { beforeHandler, afterHandler } = createHandlers({
+      [childSessionID]: parentSessionID,
+    })
+
+    await beforeHandler(
+      { tool: "task", sessionID: parentSessionID, callID: "call-task-timer-1" },
+      { args: { prompt: "Implement auth flow" } },
+    )
+
+    // when
+    await afterHandler(
+      { tool: "task", sessionID: parentSessionID, callID: "call-task-timer-1" },
+      {
+        title: "Sisyphus Task",
+        output: "Task completed\n<task_metadata>\nsession_id: ses_child\n</task_metadata>",
+        metadata: {
+          sessionId: childSessionID,
+          agent: "sisyphus-junior",
+          category: "deep",
+        },
+      },
+    )
+
+    // then
+    const taskSession = readBoulderState(testDirectory)?.works?.["work-1"]?.task_sessions?.["todo:1"]
+    expect(taskSession).toBeDefined()
+    expect(taskSession?.started_at).toBeString()
+    expect(taskSession?.status).toBe("running")
+    expect(taskSession?.session_id).toBe(childSessionID)
+  })
+
+  it("ends task timer when todo:1 checkbox transitions to checked", async () => {
+    // given
+    const parentSessionID = "ses_parent_2"
+    const childSessionID = "ses_child_2"
+    const planPath = join(testDirectory, "task-timer-complete-plan.md")
+    writeFileSync(planPath, "# Plan\n\n## TODOs\n- [ ] 1. Implement auth flow\n", "utf-8")
+    writeBoulderState(testDirectory, {
+      schema_version: 2,
+      active_work_id: "work-1",
+      active_plan: planPath,
+      started_at: "2026-01-02T10:00:00Z",
+      session_ids: [parentSessionID],
+      plan_name: "task-timer-complete-plan",
+      works: {
+        "work-1": {
+          work_id: "work-1",
+          active_plan: planPath,
+          plan_name: "task-timer-complete-plan",
+          started_at: "2026-01-02T10:00:00Z",
+          session_ids: [parentSessionID],
+          status: "active",
+        },
+      },
+    })
+    const { beforeHandler, afterHandler } = createHandlers({
+      [childSessionID]: parentSessionID,
+    })
+
+    await beforeHandler(
+      { tool: "task", sessionID: parentSessionID, callID: "call-task-timer-2" },
+      { args: { prompt: "Implement auth flow" } },
+    )
+    writeFileSync(planPath, "# Plan\n\n## TODOs\n- [x] 1. Implement auth flow\n", "utf-8")
+
+    // when
+    await afterHandler(
+      { tool: "task", sessionID: parentSessionID, callID: "call-task-timer-2" },
+      {
+        title: "Sisyphus Task",
+        output: "Task completed\n<task_metadata>\nsession_id: ses_child_2\n</task_metadata>",
+        metadata: {
+          sessionId: childSessionID,
+          agent: "sisyphus-junior",
+          category: "deep",
+        },
+      },
+    )
+
+    // then
+    const taskSession = readBoulderState(testDirectory)?.works?.["work-1"]?.task_sessions?.["todo:1"]
+    expect(taskSession).toBeDefined()
+    expect(taskSession?.ended_at).toBeString()
+    expect(taskSession?.status).toBe("completed")
+    expect((taskSession?.elapsed_ms ?? 0) > 0).toBe(true)
+  })
+})

src/hooks/atlas/tool-execute-after.ts

@@ -1,11 +1,16 @@
 import type { PluginInput } from "@opencode-ai/plugin"
 import {
+  endTaskTimer,
+  getWorkForSession,
   getPlanProgress,
   getTaskSessionState,
   readBoulderState,
   resolveBoulderPlanPath,
+  resolveBoulderPlanPathForWork,
+  startTaskTimer,
   upsertTaskSessionState,
 } from "../../features/boulder-state"
+import { existsSync, readFileSync } from "node:fs"
 import { log } from "../../shared/logger"
 import { isCallerOrchestrator } from "../../shared/session-utils"
 import { syncBackgroundLaunchSessionTracking } from "./background-launch-session-tracking"
@@ -26,6 +31,34 @@ import { isWriteOrEditToolName } from "./write-edit-tool-policy"
 import type { PendingTaskRef, SessionState } from "./types"
 import type { ToolExecuteAfterInput, ToolExecuteAfterOutput } from "./types"
 
+function isTrackedTaskChecked(planPath: string, taskKey: string): boolean {
+  if (!existsSync(planPath)) {
+    return false
+  }
+
+  const [section, label] = taskKey.split(":")
+  if (!section || !label) {
+    return false
+  }
+
+  const escapedLabel = label.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")
+  const matcher = section === "todo"
+    ? new RegExp(`^\\s*[-*]\\s*\\[[xX]\\]\\s*${escapedLabel}\\.\\s+`, "m")
+    : section === "final-wave"
+      ? new RegExp(`^\\s*[-*]\\s*\\[[xX]\\]\\s*${escapedLabel.toUpperCase()}\\.\\s+`, "m")
+      : null
+  if (!matcher) {
+    return false
+  }
+
+  try {
+    const content = readFileSync(planPath, "utf-8")
+    return matcher.test(content)
+  } catch {
+    return false
+  }
+}
+
 export function createToolExecuteAfterHandler(input: {
   ctx: PluginInput
   pendingFilePaths: Map<string, string>
@@ -100,7 +133,29 @@ export function createToolExecuteAfterHandler(input: {
       const extractedSessionId = metadataSessionId ?? extractSessionIdFromOutput(toolOutput.output)
 
       if (boulderState) {
-        const planPath = resolveBoulderPlanPath(ctx.directory, boulderState)
+        const sessionWork = toolInput.sessionID
+          ? getWorkForSession(ctx.directory, toolInput.sessionID)
+          : null
+        const planPath = sessionWork
+          ? resolveBoulderPlanPathForWork(ctx.directory, sessionWork)
+          : resolveBoulderPlanPath(ctx.directory, boulderState)
+        const workScopedBoulderState = sessionWork
+          ? {
+              ...boulderState,
+              active_plan: sessionWork.active_plan,
+              plan_name: sessionWork.plan_name,
+              status: sessionWork.status,
+              started_at: sessionWork.started_at,
+              ended_at: sessionWork.ended_at,
+              elapsed_ms: sessionWork.elapsed_ms,
+              updated_at: sessionWork.updated_at,
+              session_ids: [...sessionWork.session_ids],
+              session_origins: sessionWork.session_origins ? { ...sessionWork.session_origins } : {},
+              agent: sessionWork.agent,
+              worktree_path: sessionWork.worktree_path,
+              task_sessions: sessionWork.task_sessions ? { ...sessionWork.task_sessions } : {},
+            }
+          : boulderState
         const progress = getPlanProgress(planPath)
         const {
           currentTask,
@@ -112,7 +167,7 @@ export function createToolExecuteAfterHandler(input: {
           : null
         const sessionState = toolInput.sessionID ? getState(toolInput.sessionID) : undefined
 
-        const lineageSessionIDs = boulderState.session_ids
+        const lineageSessionIDs = sessionWork?.session_ids ?? boulderState.session_ids
         const subagentSessionId = await validateSubagentSessionId({
           client: ctx.client,
           sessionID: extractedSessionId,
@@ -120,14 +175,28 @@ export function createToolExecuteAfterHandler(input: {
         })
 
         if (currentTask && subagentSessionId && !shouldSkipTaskSessionUpdate) {
-          upsertTaskSessionState(ctx.directory, {
+          if (sessionWork) {
-            taskKey: currentTask.key,
+            startTaskTimer(ctx.directory, sessionWork.work_id, {
-            taskLabel: currentTask.label,
+              taskKey: currentTask.key,
-            taskTitle: currentTask.title,
+              taskLabel: currentTask.label,
-            sessionId: subagentSessionId,
+              taskTitle: currentTask.title,
-            agent: typeof toolOutput.metadata?.agent === "string" ? toolOutput.metadata.agent : undefined,
+              sessionId: subagentSessionId,
-            category: typeof toolOutput.metadata?.category === "string" ? toolOutput.metadata.category : undefined,
+              agent: typeof toolOutput.metadata?.agent === "string" ? toolOutput.metadata.agent : undefined,
-          })
+              category: typeof toolOutput.metadata?.category === "string" ? toolOutput.metadata.category : undefined,
+            })
+            if (isTrackedTaskChecked(planPath, currentTask.key)) {
+              endTaskTimer(ctx.directory, sessionWork.work_id, currentTask.key)
+            }
+          } else {
+            upsertTaskSessionState(ctx.directory, {
+              taskKey: currentTask.key,
+              taskLabel: currentTask.label,
+              taskTitle: currentTask.title,
+              sessionId: subagentSessionId,
+              agent: typeof toolOutput.metadata?.agent === "string" ? toolOutput.metadata.agent : undefined,
+              category: typeof toolOutput.metadata?.category === "string" ? toolOutput.metadata.category : undefined,
+            })
+          }
         }
 
         const preferredSessionId = resolvePreferredSessionId(
@@ -155,11 +224,11 @@ export function createToolExecuteAfterHandler(input: {
         }
 
         const leadReminder = shouldPauseForApproval
-          ? buildFinalWaveApprovalReminder(boulderState.plan_name, progress, preferredSessionId)
+          ? buildFinalWaveApprovalReminder(workScopedBoulderState.plan_name, progress, preferredSessionId)
-          : buildCompletionGate(boulderState.plan_name, preferredSessionId)
+          : buildCompletionGate(workScopedBoulderState.plan_name, preferredSessionId)
         const followupReminder = shouldPauseForApproval
           ? null
-          : buildOrchestratorReminder(boulderState.plan_name, progress, preferredSessionId, autoCommit, false)
+          : buildOrchestratorReminder(workScopedBoulderState.plan_name, progress, preferredSessionId, autoCommit, false)
 
         toolOutput.output = `
 <system-reminder>
@@ -181,8 +250,8 @@ ${
     ? ""
     : `<system-reminder>\n${followupReminder}\n</system-reminder>`
 }`
-        log(`[${HOOK_NAME}] Output transformed for orchestrator mode (boulder)`, {
+          log(`[${HOOK_NAME}] Output transformed for orchestrator mode (boulder)`, {
-          plan: boulderState.plan_name,
+          plan: workScopedBoulderState.plan_name,
           progress: `${progress.completed}/${progress.total}`,
           fileCount: gitStats.length,
           preferredSessionId,