luna/oh-my-opencode
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
YeonGyu-Kim
2026-05-11 13:39:47 +09:00
parent f2a5ef0966
commit 127112e1e2
2 changed files with 306 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/hooks/atlas/tool-execute-after-task-timers.test.ts
+222
View File
@@ -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
+84 -15
View File
@@ -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, {
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 = `
<system-reminder>
@@ -181,8 +250,8 @@ ${
? ""
: `<system-reminder>\n${followupReminder}\n</system-reminder>`
}`
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,