diff --git a/.engineering b/.engineering index 90926a4e..8e9203ff 160000 --- a/.engineering +++ b/.engineering @@ -1 +1 @@ -Subproject commit 90926a4ee6070686e22f978d0acefece9b5f773c +Subproject commit 8e9203fff3443fa946ec9d399cd4a696c6786821 diff --git a/schemas/session-file.schema.json b/schemas/session-file.schema.json index feb4769a..07bb447a 100644 --- a/schemas/session-file.schema.json +++ b/schemas/session-file.schema.json @@ -133,6 +133,107 @@ }, "default": {} }, + "usage": { + "type": "object", + "properties": { + "perActivity": { + "type": "object", + "additionalProperties": { + "type": "object", + "properties": { + "input_tokens": { + "type": "integer", + "minimum": 0 + }, + "output_tokens": { + "type": "integer", + "minimum": 0 + }, + "total_tokens": { + "type": "integer", + "minimum": 0 + }, + "cache_read_tokens": { + "type": "integer", + "minimum": 0 + }, + "cache_write_5m_tokens": { + "type": "integer", + "minimum": 0 + }, + "cache_write_1h_tokens": { + "type": "integer", + "minimum": 0 + }, + "model": { + "type": "string" + }, + "cost_usd": { + "anyOf": [ + { + "type": "number", + "minimum": 0 + }, + { + "type": "null" + } + ] + }, + "priceTableVersion": { + "type": "string" + } + }, + "required": [ + "input_tokens", + "output_tokens", + "total_tokens", + "cost_usd" + ], + "additionalProperties": false + }, + "default": {} + }, + "workflowTotal": { + "type": "object", + "properties": { + "input_tokens": { + "type": "integer", + "minimum": 0 + }, + "output_tokens": { + "type": "integer", + "minimum": 0 + }, + "total_tokens": { + "type": "integer", + "minimum": 0 + }, + "cost_usd": { + "anyOf": [ + { + "type": "number", + "minimum": 0 + }, + { + "type": "null" + } + ] + } + }, + "required": [ + "input_tokens", + "output_tokens", + "total_tokens", + "cost_usd" + ], + "additionalProperties": false + } + }, + "required": [ + "workflowTotal" + ], + "additionalProperties": false + }, "history": { "type": "array", "items": { @@ -170,7 +271,8 @@ "technique_fetched", "resource_fetched", "technique_bundled", - "variables_seeded" + "variables_seeded", + "usage_recorded" ] }, "activity": { diff --git a/schemas/state.schema.json b/schemas/state.schema.json index 374eebc3..a5c703be 100644 --- a/schemas/state.schema.json +++ b/schemas/state.schema.json @@ -204,7 +204,8 @@ "technique_fetched", "resource_fetched", "technique_bundled", - "variables_seeded" + "variables_seeded", + "usage_recorded" ] }, "activity": { diff --git a/src/config.ts b/src/config.ts index 587236b6..d27db05a 100644 --- a/src/config.ts +++ b/src/config.ts @@ -31,8 +31,18 @@ export interface ServerConfig { traceStore?: TraceStore; /** Minimum seconds between checkpoint issuance and response. Default 3. Set to 0 for testing. */ minCheckpointResponseSeconds?: number; + /** + * Per-model USD-per-MTok rates ({ input, output }) for cost estimation. + * Cache rates are derived from base input via universal multipliers. + */ + priceTable?: PriceTable; + /** Version string stamped on each cost estimate; env-overridable via PRICE_TABLE_VERSION. */ + priceTableVersion?: string; } +/** USD-per-MTok base input/output rates keyed by model id. */ +export type PriceTable = Record; + /** Config shape after startup — traceStore is guaranteed present. */ export interface ResolvedServerConfig extends ServerConfig { traceStore: TraceStore; @@ -49,6 +59,19 @@ const PROJECT_ROOT = resolve(import.meta.dirname, '..'); export const DEFAULT_BUNDLE_HEADROOM_FRACTION = 0.8; export const DEFAULT_BUNDLE_CHARS_PER_TOKEN = 4; +/** + * Build-time price table seeded from Anthropic pricing (2026-07-15). + * Re-confirm figures at build time when pricing changes. + */ +export const DEFAULT_PRICE_TABLE: PriceTable = { + 'claude-opus-4-8': { input: 5, output: 25 }, + 'claude-sonnet-5': { input: 2, output: 10 }, + 'claude-haiku-4-5': { input: 1, output: 5 }, + 'claude-fable-5': { input: 10, output: 50 }, +}; + +export const DEFAULT_PRICE_TABLE_VERSION = '2026-07-15'; + function envOrDefault(key: string, fallback: string): string { const value = process.env[key]?.trim(); return value || fallback; @@ -131,5 +154,7 @@ export function loadConfig(argv: readonly string[] = process.argv.slice(2)): Ser serverVersion: envOrDefault('SERVER_VERSION', '2.1.0'), bundleHeadroomFraction: envNumberOrDefault('BUNDLE_HEADROOM_FRACTION', DEFAULT_BUNDLE_HEADROOM_FRACTION), bundleCharsPerToken: envNumberOrDefault('BUNDLE_CHARS_PER_TOKEN', DEFAULT_BUNDLE_CHARS_PER_TOKEN), + priceTable: DEFAULT_PRICE_TABLE, + priceTableVersion: envOrDefault('PRICE_TABLE_VERSION', DEFAULT_PRICE_TABLE_VERSION), }; } diff --git a/src/schema/session.schema.ts b/src/schema/session.schema.ts index 1b8c51f3..e5bddd14 100644 --- a/src/schema/session.schema.ts +++ b/src/schema/session.schema.ts @@ -48,6 +48,36 @@ export const ActiveCheckpointSchema = z.object({ }); export type ActiveCheckpoint = z.infer; +/** Per-activity token usage entry with cost frozen at capture time. */ +export const ActivityUsageEntrySchema = z.object({ + input_tokens: z.number().int().nonnegative(), + output_tokens: z.number().int().nonnegative(), + total_tokens: z.number().int().nonnegative(), + cache_read_tokens: z.number().int().nonnegative().optional(), + cache_write_5m_tokens: z.number().int().nonnegative().optional(), + cache_write_1h_tokens: z.number().int().nonnegative().optional(), + model: z.string().optional(), + cost_usd: z.number().nonnegative().nullable(), + priceTableVersion: z.string().optional(), +}); +export type ActivityUsageEntry = z.infer; + +/** Rolled-up per-workflow token totals (child-inclusive at completion). */ +export const WorkflowUsageTotalSchema = z.object({ + input_tokens: z.number().int().nonnegative(), + output_tokens: z.number().int().nonnegative(), + total_tokens: z.number().int().nonnegative(), + cost_usd: z.number().nonnegative().nullable(), +}); +export type WorkflowUsageTotal = z.infer; + +/** Durable usage roll-up persisted on SessionFile. */ +export const SessionUsageSchema = z.object({ + perActivity: z.record(ActivityUsageEntrySchema).default({}), + workflowTotal: WorkflowUsageTotalSchema, +}); +export type SessionUsage = z.infer; + /** * Base (non-recursive) shape of `SessionFile`. Used as the building block for * the recursive `SessionFileSchema` below; the only structural extension is @@ -101,6 +131,12 @@ const SessionFileBaseSchema = z.object({ */ checkpointResponses: z.record(CheckpointResponseSchema).default({}), + /** + * Rolled-up token usage and cost estimate. Absent until the first + * `usage_recorded` event; child-inclusive at workflow completion. + */ + usage: SessionUsageSchema.optional(), + /** Append-only event log for the session. */ history: z.array(HistoryEntrySchema).default([]), @@ -177,6 +213,7 @@ export interface SessionFile { completedActivities: string[]; skippedActivities: string[]; checkpointResponses: Record; + usage?: SessionUsage; history: HistoryEntry[]; status: 'running' | 'completed' | 'aborted'; triggeredWorkflows: EmbeddedSessionRef[]; diff --git a/src/schema/state.schema.ts b/src/schema/state.schema.ts index 8a11e6d6..62864f78 100644 --- a/src/schema/state.schema.ts +++ b/src/schema/state.schema.ts @@ -27,6 +27,10 @@ export const HistoryEventTypeSchema = z.enum([ // session variable bag at session creation. ONE event per session; `data` // carries { variables: }. 'variables_seeded', + // Token-use tracking (#232): per-activity native usage relayed by the + // orchestrator at the next_activity transition seam. `data` carries the + // token figures, model, cost_usd (nullable), and priceTableVersion. + 'usage_recorded', ]); export type HistoryEventType = z.infer; diff --git a/src/tools/workflow-tools.ts b/src/tools/workflow-tools.ts index 95bbfaf4..ef9c8961 100644 --- a/src/tools/workflow-tools.ts +++ b/src/tools/workflow-tools.ts @@ -1,7 +1,7 @@ import { z } from 'zod'; import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import type { ServerConfig } from '../config.js'; -import { DEFAULT_BUNDLE_HEADROOM_FRACTION, DEFAULT_BUNDLE_CHARS_PER_TOKEN } from '../config.js'; +import { DEFAULT_BUNDLE_HEADROOM_FRACTION, DEFAULT_BUNDLE_CHARS_PER_TOKEN, DEFAULT_PRICE_TABLE, DEFAULT_PRICE_TABLE_VERSION } from '../config.js'; import { listWorkflows, listWorkflowsWithDiagnostics, loadWorkflow, loadWorkflowWithDiagnostics, getActivity, getCheckpoint, readActivityRaw, buildFragmentsLookup, TERMINAL_SENTINEL } from '../loaders/workflow-loader.js'; import { injectCheckpointFragmentBodies, resolveCheckpointFragment, scanCheckpointRefLines } from '../loaders/fragment-resolver.js'; import { resolveTechniques, formatTechniqueBundle, composeActivityTechnique, projectTechnique, projectTechniqueToYaml } from '../loaders/technique-loader.js'; @@ -28,6 +28,7 @@ import { import type { SessionFile } from '../schema/session.schema.js'; import { buildValidation, validateWorkflowVersion, validateActivityTransition, validateStepManifest, validateTechniqueFetches, validateTransitionCondition, validateActivityManifest } from '../utils/validation.js'; import type { StepManifestEntry, ActivityManifestEntry } from '../utils/validation.js'; +import { finalizeUsageTree, recordActivityUsage, type UsageInput } from '../utils/usage.js'; import { createTraceToken, decodeTraceToken } from '../trace.js'; import type { TraceEvent, TraceTokenPayload } from '../trace.js'; @@ -42,6 +43,16 @@ const activityManifestSchema = z.array(z.object({ transition_condition: z.string().optional(), })).optional().describe('Activity completion manifest from the orchestrator. Reports the sequence of activities completed so far with outcomes and transition conditions.'); +const usageParamSchema = z.object({ + input_tokens: z.number().int().nonnegative(), + output_tokens: z.number().int().nonnegative(), + total_tokens: z.number().int().nonnegative().optional(), + cache_read_tokens: z.number().int().nonnegative().optional(), + cache_write_5m_tokens: z.number().int().nonnegative().optional(), + cache_write_1h_tokens: z.number().int().nonnegative().optional(), + model: z.string().min(1).optional(), +}).optional().describe('Native model usage for the activity being EXITED, relayed by the orchestrator/harness (the worker cannot self-measure). Omit entirely when no usage is available — the run still completes. Absent sub-figures are omitted, not fabricated.'); + /** * Wrap a tool handler so any thrown `SessionStoreError` is re-thrown with a * user-facing message. Keeps the per-handler logic terse — handlers can @@ -123,7 +134,7 @@ export async function composeActivityArtifacts( /** The views `inspect_session` can project. `summary` is the default composite. */ export const INSPECT_SESSION_VIEWS = [ - 'summary', 'identity', 'variables', 'checkpoints', 'activities', 'history', 'children', + 'summary', 'identity', 'variables', 'checkpoints', 'activities', 'history', 'children', 'usage', ] as const; export type InspectSessionView = (typeof INSPECT_SESSION_VIEWS)[number]; @@ -215,9 +226,14 @@ export function projectChildren(s: SessionFile): Array> }); } +/** Usage projection: rolled-up per-activity and per-workflow token figures. */ +export function projectUsage(s: SessionFile): Record | undefined { + return s.usage ? { ...s.usage } : undefined; +} + /** Summary (default) view: the composite of all projections for the addressed session. */ export function projectSummary(s: SessionFile): Record { - return { + const summary: Record = { identity: projectIdentity(s), activities: projectActivities(s), variables: s.variables ?? {}, @@ -225,6 +241,9 @@ export function projectSummary(s: SessionFile): Record { history: projectHistory(s), children: projectChildren(s), }; + const usage = projectUsage(s); + if (usage) summary['usage'] = usage; + return summary; } /** @@ -247,6 +266,7 @@ export function projectSessionView( case 'activities': return projectActivities(s); case 'history': return projectHistory(s); case 'children': return projectChildren(s); + case 'usage': return projectUsage(s) ?? {}; case 'summary': default: return projectSummary(s); @@ -377,8 +397,9 @@ export function registerWorkflowTools(server: McpServer, config: ServerConfig): transition_condition: z.string().optional().describe('The transition condition that led to this activity (from the transitions field of the previous activity). Enables server-side validation of condition-activity consistency.'), step_manifest: stepManifestSchema, activity_manifest: activityManifestSchema, + usage: usageParamSchema, }, - withAuditLog('next_activity', withSessionStoreErrors(async ({ session_index, activity_id, transition_condition, step_manifest, activity_manifest }) => { + withAuditLog('next_activity', withSessionStoreErrors(async ({ session_index, activity_id, transition_condition, step_manifest, activity_manifest, usage }) => { const loaded = await loadSessionForTool(config.workspaceDir, session_index); const { state } = loaded; @@ -431,11 +452,23 @@ export function registerWorkflowTools(server: McpServer, config: ServerConfig): ...activityManifestWarnings, ); + const priceTable = config.priceTable ?? DEFAULT_PRICE_TABLE; + const priceTableVersion = config.priceTableVersion ?? DEFAULT_PRICE_TABLE_VERSION; + const next = advanceSession(state, (draft) => { const now = new Date().toISOString(); // Exit-prior: any non-empty previous activity is recorded as // completed once we transition off it. if (draft.currentActivity) { + if (usage) { + recordActivityUsage( + draft, + draft.currentActivity, + usage as UsageInput, + priceTable, + priceTableVersion, + ); + } draft.history.push({ timestamp: now, type: 'activity_exited', activity: draft.currentActivity }); if (!draft.completedActivities.includes(draft.currentActivity)) { draft.completedActivities.push(draft.currentActivity); @@ -450,6 +483,7 @@ export function registerWorkflowTools(server: McpServer, config: ServerConfig): // across the work-package, prism, and meta workflows; the TERMINAL_SENTINEL // is the contentless terminal reached via an explicit terminal transition. if (activity_id === 'complete' || isTerminal) { + finalizeUsageTree(draft); draft.history.push({ timestamp: now, type: 'workflow_completed' }); draft.status = 'completed'; } @@ -1302,6 +1336,9 @@ export function registerWorkflowTools(server: McpServer, config: ServerConfig): }; } + const usage = projectUsage(state); + if (usage) response['usage'] = usage; + // get_workflow_status reads but does not advance — keep the on-disk state stable. response['session_index'] = session_index; diff --git a/src/utils/session/store.ts b/src/utils/session/store.ts index b8bebdcb..24c5e06b 100644 --- a/src/utils/session/store.ts +++ b/src/utils/session/store.ts @@ -101,6 +101,7 @@ const TOP_LEVEL_KEY_PRIORITY = [ 'skippedActivities', 'variables', 'checkpointResponses', + 'usage', 'history', 'triggeredWorkflows', 'parentSession', diff --git a/src/utils/usage.ts b/src/utils/usage.ts new file mode 100644 index 00000000..61fe1cc5 --- /dev/null +++ b/src/utils/usage.ts @@ -0,0 +1,180 @@ +import type { PriceTable } from '../config.js'; +import type { + ActivityUsageEntry, + SessionFile, + SessionUsage, + WorkflowUsageTotal, +} from '../schema/session.schema.js'; + +/** Native usage relayed by the orchestrator at the next_activity transition seam. */ +export interface UsageInput { + input_tokens: number; + output_tokens: number; + total_tokens?: number; + cache_read_tokens?: number; + cache_write_5m_tokens?: number; + cache_write_1h_tokens?: number; + model?: string; +} + +const MTok = 1_000_000; + +/** Empty workflow total for incremental roll-up initialization. */ +export function emptyWorkflowTotal(): WorkflowUsageTotal { + return { input_tokens: 0, output_tokens: 0, total_tokens: 0, cost_usd: 0 }; +} + +/** Derive total_tokens from input + output. */ +export function deriveTotalTokens(usage: UsageInput): number { + return usage.input_tokens + usage.output_tokens; +} + +/** + * Estimate USD cost from the price table with derived cache multipliers + * (5m write ×1.25, 1h write ×2, cache read ×0.1). Unknown model → null. + */ +export function estimateCost( + usage: UsageInput, + priceTable: PriceTable, +): number | null { + if (!usage.model) return null; + const rates = priceTable[usage.model]; + if (!rates) return null; + + let cost = + usage.input_tokens * rates.input + + usage.output_tokens * rates.output; + + if (usage.cache_read_tokens) { + cost += usage.cache_read_tokens * rates.input * 0.1; + } + if (usage.cache_write_5m_tokens) { + cost += usage.cache_write_5m_tokens * rates.input * 1.25; + } + if (usage.cache_write_1h_tokens) { + cost += usage.cache_write_1h_tokens * rates.input * 2; + } + + return cost / MTok; +} + +function sumOptional(a: number | undefined, b: number | undefined): number | undefined { + if (a === undefined && b === undefined) return undefined; + return (a ?? 0) + (b ?? 0); +} + +/** Merge two activity usage entries (re-entered activities accumulate). */ +export function mergeActivityUsage( + existing: ActivityUsageEntry, + incoming: ActivityUsageEntry, +): ActivityUsageEntry { + return { + input_tokens: existing.input_tokens + incoming.input_tokens, + output_tokens: existing.output_tokens + incoming.output_tokens, + total_tokens: existing.total_tokens + incoming.total_tokens, + cache_read_tokens: sumOptional(existing.cache_read_tokens, incoming.cache_read_tokens), + cache_write_5m_tokens: sumOptional(existing.cache_write_5m_tokens, incoming.cache_write_5m_tokens), + cache_write_1h_tokens: sumOptional(existing.cache_write_1h_tokens, incoming.cache_write_1h_tokens), + model: incoming.model ?? existing.model, + cost_usd: + existing.cost_usd !== null && incoming.cost_usd !== null + ? existing.cost_usd + incoming.cost_usd + : null, + priceTableVersion: incoming.priceTableVersion ?? existing.priceTableVersion, + }; +} + +/** Add an activity entry into the running workflow total. */ +export function addToWorkflowTotal( + total: WorkflowUsageTotal, + entry: Pick, +): WorkflowUsageTotal { + return { + input_tokens: total.input_tokens + entry.input_tokens, + output_tokens: total.output_tokens + entry.output_tokens, + total_tokens: total.total_tokens + entry.total_tokens, + cost_usd: + total.cost_usd !== null && entry.cost_usd !== null + ? total.cost_usd + entry.cost_usd + : null, + }; +} + +/** + * Child-inclusive per-workflow total by pure in-tree traversal of embedded + * child SessionFile state under triggeredWorkflows. + */ +export function sumUsageTree(state: SessionFile): WorkflowUsageTotal { + const base = state.usage?.workflowTotal ?? emptyWorkflowTotal(); + let result = { ...base }; + for (const child of state.triggeredWorkflows) { + if (child.state) { + result = addToWorkflowTotal(result, sumUsageTree(child.state)); + } + } + return result; +} + +/** Build a per-activity usage entry from relayed input and config. */ +export function buildActivityUsageEntry( + usage: UsageInput, + priceTable: PriceTable, + priceTableVersion: string, +): ActivityUsageEntry { + return { + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + total_tokens: deriveTotalTokens(usage), + cost_usd: estimateCost(usage, priceTable), + priceTableVersion, + ...(usage.model ? { model: usage.model } : {}), + ...(usage.cache_read_tokens !== undefined ? { cache_read_tokens: usage.cache_read_tokens } : {}), + ...(usage.cache_write_5m_tokens !== undefined ? { cache_write_5m_tokens: usage.cache_write_5m_tokens } : {}), + ...(usage.cache_write_1h_tokens !== undefined ? { cache_write_1h_tokens: usage.cache_write_1h_tokens } : {}), + }; +} + +/** Record usage for the exited activity into session state. */ +export function recordActivityUsage( + draft: SessionFile, + activityId: string, + usage: UsageInput, + priceTable: PriceTable, + priceTableVersion: string, +): void { + const entry = buildActivityUsageEntry(usage, priceTable, priceTableVersion); + const now = new Date().toISOString(); + + if (!draft.usage) { + draft.usage = { perActivity: {}, workflowTotal: emptyWorkflowTotal() }; + } + + const existing = draft.usage.perActivity[activityId]; + draft.usage.perActivity[activityId] = existing + ? mergeActivityUsage(existing, entry) + : entry; + + draft.usage.workflowTotal = addToWorkflowTotal(draft.usage.workflowTotal, entry); + + const { input_tokens, output_tokens, total_tokens, cost_usd, priceTableVersion: recordedVersion, ...optional } = entry; + draft.history.push({ + timestamp: now, + type: 'usage_recorded', + activity: activityId, + data: { + activityId, + input_tokens, + output_tokens, + total_tokens, + cost_usd, + priceTableVersion: recordedVersion, + ...optional, + }, + }); +} + +/** Fold child workflow usage into the parent total at completion. */ +export function finalizeUsageTree(draft: SessionFile): void { + if (!draft.usage) return; + draft.usage.workflowTotal = sumUsageTree(draft); +} diff --git a/tests/config.test.ts b/tests/config.test.ts index 5d831dc7..825aebfe 100644 --- a/tests/config.test.ts +++ b/tests/config.test.ts @@ -92,5 +92,23 @@ describe('loadConfig — workspace argument', () => { expect(config.bundleHeadroomFraction).toBe(0.8); expect(config.bundleCharsPerToken).toBe(4); }); + + it('loads default price table and version (PR233-TC-09)', () => { + const config = loadConfig(['--workspace=/tmp/ws']); + expect(config.priceTable?.['claude-sonnet-5']).toEqual({ input: 2, output: 10 }); + expect(config.priceTableVersion).toBe('2026-07-15'); + }); + + it('applies PRICE_TABLE_VERSION env override (PR233-TC-09)', () => { + const prev = process.env['PRICE_TABLE_VERSION']; + process.env['PRICE_TABLE_VERSION'] = 'custom-version'; + try { + const config = loadConfig(['--workspace=/tmp/ws']); + expect(config.priceTableVersion).toBe('custom-version'); + } finally { + if (prev === undefined) delete process.env['PRICE_TABLE_VERSION']; + else process.env['PRICE_TABLE_VERSION'] = prev; + } + }); }); }); diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts index bb73cbef..c16f49b6 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -329,6 +329,114 @@ describe('mcp-server integration', () => { }); expect(result.isError).toBe(true); }); + + it('accepts usage param and records per-activity entry (PR233-TC-10)', async () => { + await transitionToActivity(client, sessionToken, 'start-work-package'); + const result = await client.callTool({ + name: 'next_activity', + arguments: { + session_index: sessionToken, + activity_id: 'codebase-comprehension', + usage: { + input_tokens: 1000, + output_tokens: 500, + model: 'claude-sonnet-5', + }, + }, + }); + expect(result.isError).toBeFalsy(); + + const inspect = await client.callTool({ + name: 'inspect_session', + arguments: { session_index: sessionToken, view: 'usage' }, + }); + const usageView = parseToolResponse(inspect) as { perActivity?: Record }; + expect(usageView.perActivity?.['start-work-package']).toMatchObject({ + input_tokens: 1000, + output_tokens: 500, + total_tokens: 1500, + model: 'claude-sonnet-5', + priceTableVersion: '2026-07-15', + }); + expect(usageView.perActivity?.['start-work-package'].cost_usd).not.toBeNull(); + const historyInspect = await client.callTool({ + name: 'inspect_session', + arguments: { session_index: sessionToken, view: 'history' }, + }); + const historyView = parseToolResponse(historyInspect) as { byType?: Record }; + expect(historyView.byType?.usage_recorded).toBeGreaterThan(0); + }); + + it('completes without usage when param omitted (PR233-TC-11)', async () => { + await transitionToActivity(client, sessionToken, 'start-work-package'); + const result = await client.callTool({ + name: 'next_activity', + arguments: { session_index: sessionToken, activity_id: 'codebase-comprehension' }, + }); + expect(result.isError).toBeFalsy(); + const inspect = await client.callTool({ + name: 'inspect_session', + arguments: { session_index: sessionToken, view: 'usage' }, + }); + const usageView = parseToolResponse(inspect); + expect(usageView).toEqual({}); + const historyInspect = await client.callTool({ + name: 'inspect_session', + arguments: { session_index: sessionToken, view: 'history' }, + }); + const historyView = parseToolResponse(historyInspect) as { byType?: Record }; + expect(historyView.byType?.usage_recorded ?? 0).toBe(0); + }); + + it('records cache figures when supplied (PR233-TC-12)', async () => { + await transitionToActivity(client, sessionToken, 'start-work-package'); + await client.callTool({ + name: 'next_activity', + arguments: { + session_index: sessionToken, + activity_id: 'codebase-comprehension', + usage: { + input_tokens: 100, + output_tokens: 50, + cache_read_tokens: 200, + cache_write_5m_tokens: 300, + model: 'claude-sonnet-5', + }, + }, + }); + const inspect = await client.callTool({ + name: 'inspect_session', + arguments: { session_index: sessionToken, view: 'usage' }, + }); + const usageView = parseToolResponse(inspect) as { perActivity?: Record }; + expect(usageView.perActivity?.['start-work-package']).toMatchObject({ + cache_read_tokens: 200, + cache_write_5m_tokens: 300, + }); + }); + + it('records tokens with cost_usd null for unknown model (PR233-TC-15)', async () => { + await transitionToActivity(client, sessionToken, 'start-work-package'); + const result = await client.callTool({ + name: 'next_activity', + arguments: { + session_index: sessionToken, + activity_id: 'codebase-comprehension', + usage: { + input_tokens: 100, + output_tokens: 50, + model: 'unknown-model-xyz', + }, + }, + }); + expect(result.isError).toBeFalsy(); + const inspect = await client.callTool({ + name: 'inspect_session', + arguments: { session_index: sessionToken, view: 'usage' }, + }); + const usageView = parseToolResponse(inspect) as { perActivity?: Record }; + expect(usageView.perActivity?.['start-work-package'].cost_usd).toBeNull(); + }); }); describe('tool: get_activity', () => { diff --git a/tests/session-schema.test.ts b/tests/session-schema.test.ts index c0d9da44..8b673ee9 100644 --- a/tests/session-schema.test.ts +++ b/tests/session-schema.test.ts @@ -8,6 +8,7 @@ import { PARENT_CHAIN_DEPTH_WARN_THRESHOLD, type SessionFile, } from '../src/schema/session.schema.js'; +import { HistoryEntrySchema } from '../src/schema/state.schema.js'; const VALID_INDEX = 'ABCDEF'; @@ -325,4 +326,52 @@ describe('SessionFile schema', () => { expect(d).toBe(1024); }); }); + + describe('usage field (PR233)', () => { + it('accepts SessionFile with usage present (PR233-TC-01)', () => { + const data = minimalSession({ + usage: { + perActivity: { + implement: { + input_tokens: 1000, + output_tokens: 500, + total_tokens: 1500, + model: 'claude-sonnet-5', + cost_usd: 0.01, + priceTableVersion: '2026-07-15', + }, + }, + workflowTotal: { + input_tokens: 1000, + output_tokens: 500, + total_tokens: 1500, + cost_usd: 0.01, + }, + }, + }); + expect(validateSessionFile(data).usage?.perActivity.implement.total_tokens).toBe(1500); + }); + + it('accepts SessionFile with usage absent (PR233-TC-02)', () => { + const data = minimalSession(); + expect(validateSessionFile(data).usage).toBeUndefined(); + }); + + it('validates usage_recorded history entry (PR233-TC-03)', () => { + const result = HistoryEntrySchema.safeParse({ + timestamp: '2026-07-15T12:00:00.000Z', + type: 'usage_recorded', + activity: 'implement', + data: { + activityId: 'implement', + input_tokens: 100, + output_tokens: 50, + total_tokens: 150, + cost_usd: null, + priceTableVersion: '2026-07-15', + }, + }); + expect(result.success).toBe(true); + }); + }); }); diff --git a/tests/session-store.test.ts b/tests/session-store.test.ts index 29df5c6e..1f14b8bc 100644 --- a/tests/session-store.test.ts +++ b/tests/session-store.test.ts @@ -69,6 +69,19 @@ describe('session-store primitives', () => { expect(() => canonicaliseJson(Number.POSITIVE_INFINITY)).toThrow(); expect(() => canonicaliseJson(Number.NaN)).toThrow(); }); + + it('places usage before history per TOP_LEVEL_KEY_PRIORITY (PR233-TC-04)', () => { + const bytes = canonicaliseJson({ + history: [], + usage: { perActivity: {}, workflowTotal: { input_tokens: 0, output_tokens: 0, total_tokens: 0, cost_usd: 0 } }, + variables: {}, + checkpointResponses: {}, + }); + const usageIdx = bytes.indexOf('"usage"'); + const historyIdx = bytes.indexOf('"history"'); + expect(usageIdx).toBeGreaterThan(-1); + expect(historyIdx).toBeGreaterThan(usageIdx); + }); }); describe('writeSessionFile atomicity', () => { diff --git a/tests/usage.test.ts b/tests/usage.test.ts new file mode 100644 index 00000000..1cb1c2da --- /dev/null +++ b/tests/usage.test.ts @@ -0,0 +1,176 @@ +import { describe, it, expect } from 'vitest'; +import { + DEFAULT_PRICE_TABLE, + DEFAULT_PRICE_TABLE_VERSION, +} from '../src/config.js'; +import type { SessionFile } from '../src/schema/session.schema.js'; +import { + addToWorkflowTotal, + buildActivityUsageEntry, + deriveTotalTokens, + emptyWorkflowTotal, + estimateCost, + mergeActivityUsage, + recordActivityUsage, + sumUsageTree, +} from '../src/utils/usage.js'; + +describe('estimateCost', () => { + it('computes base input/output cost per the price table (PR233-TC-05)', () => { + const cost = estimateCost( + { input_tokens: 1_000_000, output_tokens: 1_000_000, model: 'claude-sonnet-5' }, + DEFAULT_PRICE_TABLE, + ); + expect(cost).toBe(12); + }); + + it('derives cache rates from base input (PR233-TC-06)', () => { + const cost = estimateCost( + { + input_tokens: 0, + output_tokens: 0, + cache_read_tokens: 1_000_000, + cache_write_5m_tokens: 1_000_000, + cache_write_1h_tokens: 1_000_000, + model: 'claude-sonnet-5', + }, + DEFAULT_PRICE_TABLE, + ); + // read: 2 * 0.1 = 0.2, 5m: 2 * 1.25 = 2.5, 1h: 2 * 2 = 4 → 6.7 + expect(cost).toBeCloseTo(6.7); + }); + + it('returns null for a model absent from the price table (PR233-TC-07)', () => { + const cost = estimateCost( + { input_tokens: 100, output_tokens: 100, model: 'unknown-model' }, + DEFAULT_PRICE_TABLE, + ); + expect(cost).toBeNull(); + }); +}); + +describe('deriveTotalTokens', () => { + it('derives input + output regardless of relayed total_tokens', () => { + expect( + deriveTotalTokens({ input_tokens: 100, output_tokens: 50, total_tokens: 999 }), + ).toBe(150); + }); +}); + +describe('sumUsageTree', () => { + function session(usage?: SessionFile['usage'], children: SessionFile[] = []): SessionFile { + return { + schemaVersion: 1, + sessionIndex: 'ABCDEF', + workflowId: 'work-package', + workflowVersion: '3.11.0', + agentId: 'worker', + seq: 0, + ts: 0, + startedAt: '2026-01-01T00:00:00.000Z', + currentActivity: '', + currentTechnique: '', + condition: '', + variables: {}, + completedActivities: [], + skippedActivities: [], + checkpointResponses: {}, + history: [], + status: 'running', + triggeredWorkflows: children.map((state, i) => ({ + workflowId: 'child', + sessionIndex: 'BBBBBB', + triggeredAt: '2026-01-01T00:00:00.000Z', + triggeredFrom: { activityId: 'dispatch' }, + status: 'completed' as const, + state, + })), + ...(usage ? { usage } : {}), + }; + } + + it('sums parent + nested child workflowTotal (PR233-TC-08)', () => { + const parentTotal = { input_tokens: 100, output_tokens: 50, total_tokens: 150, cost_usd: 1 }; + const childTotal = { input_tokens: 200, output_tokens: 100, total_tokens: 300, cost_usd: 2 }; + const child = session({ perActivity: {}, workflowTotal: childTotal }); + const parent = session({ perActivity: {}, workflowTotal: parentTotal }, [child]); + expect(sumUsageTree(parent)).toEqual(addToWorkflowTotal(parentTotal, childTotal)); + }); + + it('returns empty total when usage is absent', () => { + expect(sumUsageTree(session())).toEqual(emptyWorkflowTotal()); + }); +}); + +describe('recordActivityUsage', () => { + it('builds stamped activity entries with price table version', () => { + const entry = buildActivityUsageEntry( + { input_tokens: 1000, output_tokens: 500, model: 'claude-sonnet-5' }, + DEFAULT_PRICE_TABLE, + DEFAULT_PRICE_TABLE_VERSION, + ); + expect(entry.priceTableVersion).toBe(DEFAULT_PRICE_TABLE_VERSION); + expect(entry.cost_usd).not.toBeNull(); + }); + + it('merges re-entered activity usage', () => { + const draft: SessionFile = { + schemaVersion: 1, + sessionIndex: 'ABCDEF', + workflowId: 'work-package', + workflowVersion: '3.11.0', + agentId: 'worker', + seq: 0, + ts: 0, + startedAt: '2026-01-01T00:00:00.000Z', + currentActivity: 'implement', + currentTechnique: '', + condition: '', + variables: {}, + completedActivities: [], + skippedActivities: [], + checkpointResponses: {}, + history: [], + status: 'running', + triggeredWorkflows: [], + }; + recordActivityUsage( + draft, + 'implement', + { input_tokens: 100, output_tokens: 50, model: 'claude-sonnet-5' }, + DEFAULT_PRICE_TABLE, + DEFAULT_PRICE_TABLE_VERSION, + ); + recordActivityUsage( + draft, + 'implement', + { input_tokens: 200, output_tokens: 100, model: 'claude-sonnet-5' }, + DEFAULT_PRICE_TABLE, + DEFAULT_PRICE_TABLE_VERSION, + ); + expect(draft.usage?.perActivity.implement.input_tokens).toBe(300); + expect(draft.history.filter((e) => e.type === 'usage_recorded')).toHaveLength(2); + }); +}); + +describe('mergeActivityUsage', () => { + it('propagates null cost when either side is unpriced', () => { + const merged = mergeActivityUsage( + { + input_tokens: 1, + output_tokens: 1, + total_tokens: 2, + cost_usd: 0.01, + priceTableVersion: 'v1', + }, + { + input_tokens: 1, + output_tokens: 1, + total_tokens: 2, + cost_usd: null, + model: 'unknown', + }, + ); + expect(merged.cost_usd).toBeNull(); + }); +});