diff --git a/scripts/generate-site-data.ts b/scripts/generate-site-data.ts index 4e19515cc..90bb11a53 100644 --- a/scripts/generate-site-data.ts +++ b/scripts/generate-site-data.ts @@ -311,7 +311,7 @@ const TOOL_GROUPS: Array<{ title: string; note: string; tools: string[] }> = [ { title: 'Workflow and activity navigation', note: 'Load workflow structure and advance through activities.', tools: ['get_workflow', 'next_activity', 'get_activity'] }, { title: 'Checkpoint flow', note: 'Yield to the orchestrator, present decisions to the user, and resume.', tools: ['yield_checkpoint', 'resume_checkpoint', 'present_checkpoint', 'respond_checkpoint'] }, { title: 'Techniques and resources', note: 'Fetch technique definitions and lazy-loaded reference material.', tools: ['get_technique', 'get_resource'] }, - { title: 'Trace', note: 'Execution history for debugging and audit.', tools: ['get_trace'] }, + { title: 'Trace', note: 'Execution history and per-dispatch cost accounting.', tools: ['get_trace', 'record_usage'] }, ]; /** Plain-language one-line summaries for the site (source descriptions stay authoritative for MCP). */ diff --git a/site/api/tools.html b/site/api/tools.html index 3cecd31d8..69675f79a 100644 --- a/site/api/tools.html +++ b/site/api/tools.html @@ -79,7 +79,7 @@

MCP tool reference

  • Workflow and activity navigation: get_workflow, next_activity, get_activity
  • Checkpoint flow: yield_checkpoint, resume_checkpoint, present_checkpoint, respond_checkpoint
  • Techniques and resources: get_technique, get_resource
    • -
  • Trace: get_trace
    • +
  • Trace: get_trace, record_usage

    • Bootstrap

    @@ -247,7 +247,6 @@

    next_activity

    step_manifestobject[]noSteps completed in the previous activity, for example [{ "step_id": "detect-review-mode", "output": "is_review_mode=false" }]. Omit if no steps ran. activity_manifestobject[]noHistory of completed activities with outcomes and transition conditions. variables_changedobjectnoVariable assignments the completing activity produced — relay the worker's activity_complete variables_changed map verbatim. The server writes them into the session variable bag and records one variable_set history event per name, so the bag a later get_workflow_status / inspect_session returns reflects worker outputs and survives a lost agent context. Declared types are validated warn-only: a mismatch is stored as written and surfaced in _meta.validation. Omit when the activity changed nothing. - usageobjectnoHarness-reported token usage for the activity this call EXITS — relay the figure the harness surfaced for the worker that just completed (subagent token counts plus any cache/model fields), as reported. Recorded as an activity_usage history event keyed to the exited activity, so inspect_session reports per-activity cost. Workers cannot self-measure: omit the parameter entirely when the harness surfaces nothing rather than passing zeros. @@ -412,7 +411,7 @@

    get_resource

    Trace

    -

    Execution history for debugging and audit.

    +

    Execution history and per-dispatch cost accounting.

    get_trace

    Retrieve the session execution trace.

    @@ -431,6 +430,24 @@

    get_trace

    +
    +

    record_usage

    +

    Orchestrator tool: record harness-reported token usage for ONE completed dispatch.

    +
    + Full description +

    Orchestrator tool: record harness-reported token usage for ONE completed dispatch. Call as each dispatch finishes — the first worker, a continue, a fresh worker after a timeout, a resume after a checkpoint yield, an out-of-band dispatch, and the terminal activity.

    +
    +
    + + + + + + + +
    ParameterTypeRequiredDescription
    session_indexstringyesSix-character token from start_session. Use the same value for every call in this session.
    activitystringyesActivity the dispatch ran, whether or not the session is still on it.
    usageobjectyesHarness-reported token usage for this ONE dispatch, as reported. Omit the call entirely when the harness surfaced nothing rather than passing zeros — the worker cannot self-measure, so absence must stay distinguishable from a measured zero.
    +
    +
    diff --git a/src/schema/state.schema.ts b/src/schema/state.schema.ts index 57a60d7f8..74f3df185 100644 --- a/src/schema/state.schema.ts +++ b/src/schema/state.schema.ts @@ -27,11 +27,12 @@ export const HistoryEventTypeSchema = z.enum([ // session variable bag at session creation. ONE event per session; `data` // carries { variables: }. 'variables_seeded', - // Per-activity cost accounting (#324 B1): harness-reported token usage for - // the activity a next_activity call exits, relayed by the orchestrator. - // `activity` is the exited activity; `data` carries { usage: }. - // A worker cannot self-measure, so absence means the harness surfaced - // nothing — never a zero. + // Per-dispatch cost accounting (#324 B1, #346 DI-33): harness-reported token + // usage for ONE completed dispatch, recorded by record_usage as it finishes. + // `activity` is the activity that ran; `data` carries { usage: }. + // One event per dispatch, so a resumed or re-dispatched activity contributes a + // row per pass. A worker cannot self-measure, so absence means the harness + // surfaced nothing — never a zero. 'activity_usage', ]); export type HistoryEventType = z.infer; diff --git a/src/tools/workflow-tools.ts b/src/tools/workflow-tools.ts index deb383d84..2b6ef28d4 100644 --- a/src/tools/workflow-tools.ts +++ b/src/tools/workflow-tools.ts @@ -50,11 +50,11 @@ const activityManifestSchema = z.array(z.object({ transition_condition: z.string().optional(), })).optional().describe('Orchestrator activity-completion manifest: [{activity_id, outcome, transition_condition?}].'); -const usageSchema = z.record(z.unknown()).optional().describe( - 'Harness-reported token usage for the activity this call EXITS — relay the figure the harness surfaced for the worker that just completed ' + - '(subagent token counts plus any cache/model fields), as reported. Recorded as an `activity_usage` history event keyed to the exited activity, ' + - 'so inspect_session reports per-activity cost. Workers cannot self-measure: omit the parameter entirely when the harness surfaces nothing rather ' + - 'than passing zeros.', +const usageSchema = z.record(z.unknown()).describe( + 'Harness-reported token usage for ONE completed dispatch — subagent token counts plus any cache/model \n' + + 'fields, as reported. Recorded as an `activity_usage` history event keyed to the activity that ran, so \n' + + 'inspect_session reports cost per dispatch. Workers cannot self-measure: omit the record_usage call \n' + + 'entirely when the harness surfaces nothing rather than passing zeros.', ); const variablesChangedSchema = z.record(z.unknown()).optional().describe( @@ -245,11 +245,11 @@ export function projectChildren(s: SessionFile): Array> } /** - * Usage projection (#324 B1): one entry per `activity_usage` event, in the - * order the orchestrator reported them. An activity appears once per exit, so - * a resumed activity contributes a row per pass rather than a merged total — - * summing is the caller's call, since harnesses differ on whether a resumed - * worker re-reports cumulative figures. + * Usage projection (#324 B1, #346 DI-33): one entry per `activity_usage` event, + * in the order the orchestrator recorded them — one per dispatch. An activity + * dispatched several times contributes a row per dispatch, and summing is the + * caller's call, since harnesses differ on whether a resumed worker re-reports + * cumulative figures. */ export function projectUsage(s: SessionFile): Array> { return (s.history ?? []) @@ -440,9 +440,8 @@ export function registerWorkflowTools(server: McpServer, config: ServerConfig): step_manifest: stepManifestSchema, activity_manifest: activityManifestSchema, variables_changed: variablesChangedSchema, - usage: usageSchema, }, - withAuditLog('next_activity', withSessionStoreErrors(async ({ session_index, activity_id, transition_condition, step_manifest, activity_manifest, variables_changed, usage }) => { + withAuditLog('next_activity', withSessionStoreErrors(async ({ session_index, activity_id, transition_condition, step_manifest, activity_manifest, variables_changed }) => { const loadOpts = await sessionLoadOpts(); const loaded = await loadSessionForTool(planningRootDir, session_index, loadOpts); const { state } = loaded; @@ -478,12 +477,6 @@ export function registerWorkflowTools(server: McpServer, config: ServerConfig): ? validateTransitionCondition(view, result.value, activity_id, transition_condition) : null; - // usage measures the activity being exited, so the entry transition has - // nothing to attribute it to. Say so rather than dropping it silently. - const usageWarning = (usage && !state.currentActivity) - ? `usage supplied on the entry transition to '${activity_id}', which exits no activity — not recorded. Pass usage on the call that leaves the measured activity.` - : null; - const activityManifestWarnings: string[] = []; if (activity_manifest) { if (activity_manifest.length === 0) { @@ -509,14 +502,6 @@ export function registerWorkflowTools(server: McpServer, config: ServerConfig): if (!draft.completedActivities.includes(exitingActivity)) { draft.completedActivities.push(exitingActivity); } - // #324 B1: attribute the harness's usage figure to the activity being - // exited — it measures the worker that just finished, not the one - // about to start. Dropped on the first transition, which exits nothing. - if (usage) { - draft.history.push({ - timestamp: now, type: 'activity_usage', activity: exitingActivity, data: { usage }, - }); - } } // Persist the completing activity's worker outputs into the bag. These // are attributed to the activity being exited, not the one entered — @@ -553,7 +538,6 @@ export function registerWorkflowTools(server: McpServer, config: ServerConfig): validateActivityTransition(view, result.value, activity_id), validateWorkflowVersion(view, result.value), condWarning, - usageWarning, ...manifestWarnings, ...activityManifestWarnings, ...variableWarnings, @@ -1122,6 +1106,45 @@ export function registerWorkflowTools(server: McpServer, config: ServerConfig): }; }), traceOpts)); + server.tool('record_usage', 'Orchestrator tool: record harness-reported token usage for ONE completed dispatch. Call as each dispatch finishes — the first worker, a continue, a fresh worker after a timeout, a resume after a checkpoint yield, an out-of-band dispatch, and the terminal activity.', + { + ...sessionIndexParam, + activity: z.string().describe('Activity the dispatch ran, whether or not the session is still on it.'), + usage: usageSchema.describe( + 'Harness-reported token usage for this ONE dispatch, as reported. Omit the call entirely when the harness ' + + 'surfaced nothing rather than passing zeros — the worker cannot self-measure, so absence must stay ' + + 'distinguishable from a measured zero.', + ), + }, + withAuditLog('record_usage', withSessionStoreErrors(async ({ session_index, activity, usage }) => { + const loadOpts = await sessionLoadOpts(); + const loaded = await loadSessionForTool(planningRootDir, session_index, loadOpts); + const { state } = loaded; + + const recordedAt = new Date().toISOString(); + const next = advanceSession(state, (draft) => { + // One entry per dispatch, which is what projectUsage reports. Attribution is the + // caller's `activity`: the dispatch ran for it, whether or not the session is still + // there by the time the figure arrives. + draft.history.push({ + timestamp: recordedAt, type: 'activity_usage', activity, data: { usage }, + }); + }); + await saveSessionForTool(loaded, next); + + const recorded = (next.history ?? []).filter(e => e.type === 'activity_usage').length; + return { + content: [{ type: 'text' as const, text: JSON.stringify({ + status: 'recorded', + activity, + session_index, + usage_events: recorded, + message: `Usage recorded for one dispatch of '${activity}'. The session now carries ${recorded} usage event(s); the cost artifact reconciles that count against the run's actual dispatch count.`, + }, null, 2) }], + _meta: { session_index, validation: buildValidation() }, + }; + }), traceOpts)); + server.tool('resume_checkpoint', 'Worker tool: continue after the orchestrator resolves a checkpoint. Verifies no activeCheckpoint and returns variable updates to apply.', { ...sessionIndexParam, diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts index 0a26a15e2..4537f1884 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -869,6 +869,54 @@ describe('mcp-server integration', () => { // ============== Trace Integration ============== + describe('per-dispatch usage accounting (DI-33)', () => { + it('record_usage adds one usage event per call, for dispatches no transition exits', async () => { + const before = parseToolResponse(await client.callTool({ + name: 'inspect_session', + arguments: { session_index: sessionToken, view: 'usage' }, + })); + const baseline = Array.isArray(before) ? before.length : (before.usage?.length ?? 0); + + // Two dispatches of the SAME activity — a first pass and a resume after a + // checkpoint yield. next_activity could account for at most one of them, + // because only one transition exits the activity. + for (const total of [111, 222]) { + const res = await client.callTool({ + name: 'record_usage', + arguments: { + session_index: sessionToken, + activity: 'start-work-package', + usage: { input_tokens: total, output_tokens: 7, total_tokens: total + 7 }, + }, + }); + expect(res.isError).toBeFalsy(); + expect(parseToolResponse(res).status).toBe('recorded'); + } + + const after = parseToolResponse(await client.callTool({ + name: 'inspect_session', + arguments: { session_index: sessionToken, view: 'usage' }, + })); + const rows = Array.isArray(after) ? after : (after.usage ?? []); + expect(rows.length).toBe(baseline + 2); + + // Both passes survive as separate rows rather than one overwriting the other: + // a merged total is what hid the missing third of the run. + const mine = rows.filter((r: { activity?: string }) => r.activity === 'start-work-package'); + const totals = mine.map((r: { usage?: { input_tokens?: number } }) => r.usage?.input_tokens); + expect(totals).toContain(111); + expect(totals).toContain(222); + }); + + it('record_usage rejects an unknown session', async () => { + const res = await client.callTool({ + name: 'record_usage', + arguments: { session_index: 'NOPE00', activity: 'x', usage: { total_tokens: 1 } }, + }); + expect(res.isError).toBeTruthy(); + }); + }); + describe('trace lifecycle', () => { it('session creation initializes trace (IT-6)', async () => { const result = await client.callTool({ diff --git a/tests/variable-seeding.test.ts b/tests/variable-seeding.test.ts index da3893ca5..262dc4067 100644 --- a/tests/variable-seeding.test.ts +++ b/tests/variable-seeding.test.ts @@ -316,15 +316,15 @@ describe('B7 seeding + setVariable type validation (fixture corpus)', () => { // #324 B1: per-activity token accounting. The worker cannot self-measure, so // the orchestrator relays what the harness reported for the activity it exits. - describe('next_activity usage accounting (#324 B1)', () => { + describe('per-dispatch usage accounting (#324 B1, #346 DI-33)', () => { const usage = { input_tokens: 1200, output_tokens: 340, cache_read_input_tokens: 8000 }; - it('records usage against the exited activity and surfaces it on the usage view', async () => { + it('records usage against the named activity and surfaces it on the usage view', async () => { const slug = '2026-07-28-usage-recorded'; const started = await call('start_session', { workflow_id: 'seed-fixture', agent_id: 'orchestrator', planning_folder: planningFolder(slug) }); const sessionIndex = (started._meta as Record).session_index as string; await call('next_activity', { session_index: sessionIndex, activity_id: 'checkpoint-activity' }); - await call('next_activity', { session_index: sessionIndex, activity_id: 'followup-activity', usage }); + await call('record_usage', { session_index: sessionIndex, activity: 'checkpoint-activity', usage }); const events = readSession(slug).history.filter((h: { type: string }) => h.type === 'activity_usage'); expect(events).toHaveLength(1); @@ -337,19 +337,35 @@ describe('B7 seeding + setVariable type validation (fixture corpus)', () => { expect(rows[0]).toMatchObject({ activity: 'checkpoint-activity', usage }); }); - it('warns instead of dropping silently when usage rides the entry transition', async () => { - const slug = '2026-07-28-usage-entry'; + it('records a dispatch the graph never transitions away from', async () => { + // The case the transition-keyed ledger could not reach: an activity is entered, + // dispatched, and the run ends there. Nothing exits it, so nothing could have + // carried its cost. + const slug = '2026-07-28-usage-terminal'; const started = await call('start_session', { workflow_id: 'seed-fixture', agent_id: 'orchestrator', planning_folder: planningFolder(slug) }); const sessionIndex = (started._meta as Record).session_index as string; - const result = await call('next_activity', { session_index: sessionIndex, activity_id: 'checkpoint-activity', usage }); + await call('next_activity', { session_index: sessionIndex, activity_id: 'checkpoint-activity' }); + await call('record_usage', { session_index: sessionIndex, activity: 'checkpoint-activity', usage }); - const validation = (result._meta as { validation: { status: string; warnings: string[] } }).validation; - expect(validation.status).toBe('warning'); - expect(validation.warnings.join('\n')).toMatch(/usage supplied on the entry transition/); - expect(readSession(slug).history.filter((h: { type: string }) => h.type === 'activity_usage')).toHaveLength(0); + const rows = readSession(slug).history.filter((h: { type: string }) => h.type === 'activity_usage'); + expect(rows).toHaveLength(1); + expect(rows[0].activity).toBe('checkpoint-activity'); + }); + + it('keeps each dispatch of one activity as its own row', async () => { + const slug = '2026-07-28-usage-repeat'; + const started = await call('start_session', { workflow_id: 'seed-fixture', agent_id: 'orchestrator', planning_folder: planningFolder(slug) }); + const sessionIndex = (started._meta as Record).session_index as string; + await call('next_activity', { session_index: sessionIndex, activity_id: 'checkpoint-activity' }); + await call('record_usage', { session_index: sessionIndex, activity: 'checkpoint-activity', usage: { total_tokens: 10 } }); + await call('record_usage', { session_index: sessionIndex, activity: 'checkpoint-activity', usage: { total_tokens: 20 } }); + + const rows = readSession(slug).history.filter((h: { type: string }) => h.type === 'activity_usage'); + expect(rows).toHaveLength(2); + expect(rows.map((r: { data?: { usage?: { total_tokens?: number } } }) => r.data?.usage?.total_tokens)).toEqual([10, 20]); }); - it('omitting usage records nothing and leaves the usage view empty', async () => { + it('recording nothing leaves the usage view empty', async () => { const slug = '2026-07-28-usage-absent'; const started = await call('start_session', { workflow_id: 'seed-fixture', agent_id: 'orchestrator', planning_folder: planningFolder(slug) }); const sessionIndex = (started._meta as Record).session_index as string;