From 3bf656019134a8dfdd9b4ff58e992e736d1b8fa6 Mon Sep 17 00:00:00 2001 From: Mike Clay Date: Thu, 30 Jul 2026 06:36:47 +0100 Subject: [PATCH] feat(server): account for token usage per dispatch (#346 DI-33) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `next_activity` recorded an `activity_usage` event keyed to the activity a transition EXITS, so any dispatch not followed by a transition never reached the ledger: a resume after a checkpoint yield, a re-dispatch after a timeout, an out-of-band dispatch, an abandoned session, and the terminal activity, which nothing exits. The review run that surfaced this reported 2,751,188 tokens against an actual 4,130,657 across 33 dispatches. Cost belongs to a dispatch, not to a graph edge. `record_usage { session_index, activity, usage }` writes one event per dispatch, and `next_activity` no longer takes `usage` at all — one write path, so there is no double-count to warn about and no attribution question to answer. The `usageWarning` that fired when usage rode an entry transition goes with it: that case only existed because an edge was being asked to carry a dispatch's cost. The event type is unchanged, so projectUsage and the inspect_session usage view need no edit — that projection already documented one entry per exit and left summing to the caller, which is the shape per-dispatch accounting wants. Absence stays distinguishable from zero: omit the call when the harness surfaces nothing, since a worker cannot self-measure. Tests: the #324 B1 suite encoded transition attribution, including the entry-edge warning, so it moves onto record_usage rather than being kept alongside. Added cover for the case the old model could not reach — a dispatch the graph never transitions away from — and for two dispatches of one activity staying separate rows, since a merged total is what hid the shortfall. The entry-edge warning test is dropped: the case no longer exists. The site's tool-group drift guard caught record_usage being ungrouped; it sits under Trace, whose note now covers cost accounting. Refs #346 Co-Authored-By: Claude Opus 5 (1M context) --- scripts/generate-site-data.ts | 2 +- site/api/tools.html | 23 ++++++++-- src/schema/state.schema.ts | 11 ++--- src/tools/workflow-tools.ts | 77 ++++++++++++++++++++++------------ tests/mcp-server.test.ts | 48 +++++++++++++++++++++ tests/variable-seeding.test.ts | 38 ++++++++++++----- 6 files changed, 152 insertions(+), 47 deletions(-) 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;