From 127112e1e2057cd6f5c86684c4a72d869e56bd87 Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Mon, 11 May 2026 13:39:47 +0900 Subject: [PATCH] feat(hooks/atlas): wire per-task timers via startTaskTimer/endTaskTimer --- .../tool-execute-after-task-timers.test.ts | 222 ++++++++++++++++++ src/hooks/atlas/tool-execute-after.ts | 99 ++++++-- 2 files changed, 306 insertions(+), 15 deletions(-) create mode 100644 src/hooks/atlas/tool-execute-after-task-timers.test.ts diff --git a/src/hooks/atlas/tool-execute-after-task-timers.test.ts b/src/hooks/atlas/tool-execute-after-task-timers.test.ts new file mode 100644 index 000000000..64182c93e --- /dev/null +++ b/src/hooks/atlas/tool-execute-after-task-timers.test.ts @@ -0,0 +1,222 @@ +/// + +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) { + 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() + 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\nsession_id: ses_child\n", + 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\nsession_id: ses_child_2\n", + 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) + }) +}) diff --git a/src/hooks/atlas/tool-execute-after.ts b/src/hooks/atlas/tool-execute-after.ts index 3869c291f..4cef75ac1 100644 --- a/src/hooks/atlas/tool-execute-after.ts +++ b/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 @@ -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, { - 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, - }) + if (sessionWork) { + startTaskTimer(ctx.directory, sessionWork.work_id, { + 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, + }) + 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) - : buildCompletionGate(boulderState.plan_name, preferredSessionId) + ? buildFinalWaveApprovalReminder(workScopedBoulderState.plan_name, progress, 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 = ` @@ -181,8 +250,8 @@ ${ ? "" : `\n${followupReminder}\n` }` - log(`[${HOOK_NAME}] Output transformed for orchestrator mode (boulder)`, { - plan: boulderState.plan_name, + log(`[${HOOK_NAME}] Output transformed for orchestrator mode (boulder)`, { + plan: workScopedBoulderState.plan_name, progress: `${progress.completed}/${progress.total}`, fileCount: gitStats.length, preferredSessionId,