feat(server): account for token usage per dispatch (#346 DI-33)

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). */

site/api/tools.html

@@ -79,7 +79,7 @@ <h1>MCP tool reference</h1>
         <li>Workflow and activity navigation: <a href="#get_workflow"><code>get_workflow</code></a>, <a href="#next_activity"><code>next_activity</code></a>, <a href="#get_activity"><code>get_activity</code></a></li>
         <li>Checkpoint flow: <a href="#yield_checkpoint"><code>yield_checkpoint</code></a>, <a href="#resume_checkpoint"><code>resume_checkpoint</code></a>, <a href="#present_checkpoint"><code>present_checkpoint</code></a>, <a href="#respond_checkpoint"><code>respond_checkpoint</code></a></li>
         <li>Techniques and resources: <a href="#get_technique"><code>get_technique</code></a>, <a href="#get_resource"><code>get_resource</code></a></li>
-        <li>Trace: <a href="#get_trace"><code>get_trace</code></a></li>
+        <li>Trace: <a href="#get_trace"><code>get_trace</code></a>, <a href="#record_usage"><code>record_usage</code></a></li>
       </ul>
       </nav>
       <h2>Bootstrap</h2>
@@ -247,7 +247,6 @@ <h3><code>next_activity</code></h3>
           <tr><td><code>step_manifest</code></td><td><code>object[]</code></td><td>no</td><td>Steps completed in the previous activity, for example <code>[{ &quot;step_id&quot;: &quot;detect-review-mode&quot;, &quot;output&quot;: &quot;is_review_mode=false&quot; }]</code>. Omit if no steps ran.</td></tr>
           <tr><td><code>activity_manifest</code></td><td><code>object[]</code></td><td>no</td><td>History of completed activities with outcomes and transition conditions.</td></tr>
           <tr><td><code>variables_changed</code></td><td><code>object</code></td><td>no</td><td>Variable assignments the completing activity produced — relay the worker's <code>activity_complete</code> <code>variables_changed</code> map verbatim. The server writes them into the session variable bag and records one <code>variable_set</code> 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.</td></tr>
-          <tr><td><code>usage</code></td><td><code>object</code></td><td>no</td><td>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 <code>activity_usage</code> 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.</td></tr>
           </tbody>
         </table>
         </div>
@@ -412,7 +411,7 @@ <h3><code>get_resource</code></h3>
         </div>
       </section>
       <h2>Trace</h2>
-      <p>Execution history for debugging and audit.</p>
+      <p>Execution history and per-dispatch cost accounting.</p>
       <section class="tool" id="get_trace">
         <h3><code>get_trace</code></h3>
         <p class="tool-summary">Retrieve the session execution trace.</p>
@@ -431,6 +430,24 @@ <h3><code>get_trace</code></h3>
         </table>
         </div>
       </section>
+      <section class="tool" id="record_usage">
+        <h3><code>record_usage</code></h3>
+        <p class="tool-summary">Orchestrator tool: record harness-reported token usage for ONE completed dispatch.</p>
+        <details class="tool-details">
+          <summary>Full description</summary>
+          <p>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.</p>
+        </details>
+        <div class="table-wrap">
+        <table class="param-table">
+          <thead><tr><th scope="col">Parameter</th><th scope="col">Type</th><th scope="col">Required</th><th scope="col">Description</th></tr></thead>
+          <tbody>
+          <tr><td><code>session_index</code></td><td><code>string</code></td><td>yes</td><td>Six-character token from <code>start_session</code>. Use the same value for every call in this session.</td></tr>
+          <tr><td><code>activity</code></td><td><code>string</code></td><td>yes</td><td>Activity the dispatch ran, whether or not the session is still on it.</td></tr>
+          <tr><td><code>usage</code></td><td><code>object</code></td><td>yes</td><td>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.</td></tr>
+          </tbody>
+        </table>
+        </div>
+      </section>
       <!-- END GENERATED -->
 </main>
 

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: <the seeded map> }.
   '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: <as reported> }.
-  // 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: <as reported> }.
+  // 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<typeof HistoryEventTypeSchema>;

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<Record<string, unknown>>
 }
 
 /**
- * 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<Record<string, unknown>> {
   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,

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({