diff --git a/README.md b/README.md index 19c3a8ee6..e43b9758e 100644 --- a/README.md +++ b/README.md @@ -30,6 +30,7 @@ dashboard — any provider, any deploy target, no lock-in. - **Project instructions** — `createAgent({ projectInstructions: true })` appends the nearest `AGENTS.md` / `CLAUDE.md` to the instructions (opt-in) >>>>>>> origin/main - **Parallel tool calls** — when the model asks for several tools in one turn they run concurrently (cap it with `toolConcurrency`, or `1` for sequential), and results still reach the transcript in the model's call order +- **Usage & cost** — every result carries `usage` (tokens, USD cost, per-model breakdown, delegated children included, continuing across resumes) and `formatUsage(result.usage)` for logs - **Skills** — `defineSkill()` / `loadSkills('./skills')`: only each skill's name and description sit in the system prompt; the model loads the full markdown on demand through an auto-registered `load_skill` tool - **Agents as directories** — `loadAgentDir('./my-agent')` turns `instructions.md`, `tools/*.ts`, `skills/` and `subagents//` into the same object `createAgent()` returns, so you can start with files and drop to code later — see [Agent directories](docs/agent-directories.md) - **Multi-agent delegation** — wrap a child agent as a tool with `createDelegateTool()`, with a `maxDepth` guard against delegation loops diff --git a/docs/api-overview.md b/docs/api-overview.md index 2e67fe19b..3eaf0c930 100644 --- a/docs/api-overview.md +++ b/docs/api-overview.md @@ -333,6 +333,45 @@ const usd = estimateCost({ inputTokens: used, outputTokens: 500 }, 'my-llama'); The built-in context windows and prices are a dated snapshot (see the retrieval date and sources at the top of `src/models/modelData.ts`). Providers change prices and models, so override entries with `registerModel` when you need billing-grade numbers. +### Usage and cost of a run + +Every `ExecutionResult` (and `agent.send()` result) carries `usage`, the running total of the whole run: + +```ts +import { AgentExecutor, formatUsage } from '@loushy/build-ai-agent'; + +const result = await AgentExecutor.execute({ agent, input: 'Compare 3 cities', provider }); + +result.usage.inputTokens; // all model calls of the run, delegated children included +result.usage.costUsd; // number, or undefined if any model used has no known price +result.usage.estimated; // true if some call reported no usage and was estimated +result.usage.byModel['gpt-4o-mini']; // { inputTokens, outputTokens, costUsd?, calls } +result.stepUsage?.[0]; // { step, model, usage, estimated, costUsd? } per model call +console.log(formatUsage(result.usage)); // 1,234 in / 567 out tokens · $0.0042 (2 model calls) +``` + +- **Reported vs estimated.** Providers report `promptTokens`/`completionTokens`; the built-in providers also pass on `cachedInputTokens`/`reasoningTokens` when the 'ai' SDK's provider metadata carries them (`usage.cachedInputTokens` and `usage.reasoningTokens` only appear then). A backend that reports nothing gives `undefined` usage, never zeros. For that step the executor falls back to `estimateTokens` and sets `usage.estimated` (a leading `~` in `formatUsage`). A custom provider should leave `GenerateResult.usage` unset rather than fill in zeros. +- **Cost.** `costUsd` is the sum of `estimateCost` per model. It is `undefined`, never a misleading partial sum, as soon as any model used has unknown pricing; `byModel` shows which ones are priced. +- **Delegation.** A delegated child's usage is added to the parent's totals and `byModel`, and is also shown on its own as `usage.delegated` (`{ inputTokens, outputTokens, totalTokens, costUsd, modelCalls, estimated, runs }`). +- **Resume.** A run resumed from a checkpoint, or after an approval, continues from the saved totals instead of restarting at zero. Checkpoints written by older versions start from their saved token counts with an unknown cost. +- **Events and traces.** The `finish` event (and every lifecycle event that carried `usage`) now carries the running totals; `text-complete` also has `stepUsage`, and `onLLMResponse` receives the call's usage as a third argument. The `chat` span's `gen_ai.usage.*` attributes use the same numbers, with `loushy.usage.estimated` set to `true` when they are estimates. +- **Streaming.** `agent.stream()` events carry the same accounting: `step.done` and `run.done` have `usage` with `inputTokens`, `outputTokens`, `estimated` and `costUsd` (run-level also `modelCalls`), alongside the older `promptTokens`/`completionTokens`. +- `promptTokens` and `completionTokens` on `usage` remain as deprecated aliases of `inputTokens` and `outputTokens`. + +Prices come from the model registry above, so to get a cost for a custom or fine-tuned model, register it under the id you pass as the model: + +```ts +import { registerModel } from '@loushy/build-ai-agent'; + +registerModel({ + id: 'ft:gpt-4o-mini:acme', + provider: 'openai', + contextWindow: 128000, + inputCostPerMTok: 0.3, + outputCostPerMTok: 1.2, +}); +``` + ## Flows, evals, observability and security - `FlowBuilder` / `FlowExecutor` - multi-step workflow graphs. diff --git a/llms-full.txt b/llms-full.txt index ca4649afb..7445135b1 100644 --- a/llms-full.txt +++ b/llms-full.txt @@ -28,6 +28,7 @@ dashboard — any provider, any deploy target, no lock-in. - **Project instructions** — `createAgent({ projectInstructions: true })` appends the nearest `AGENTS.md` / `CLAUDE.md` to the instructions (opt-in) >>>>>>> origin/main - **Parallel tool calls** — when the model asks for several tools in one turn they run concurrently (cap it with `toolConcurrency`, or `1` for sequential), and results still reach the transcript in the model's call order +- **Usage & cost** — every result carries `usage` (tokens, USD cost, per-model breakdown, delegated children included, continuing across resumes) and `formatUsage(result.usage)` for logs - **Skills** — `defineSkill()` / `loadSkills('./skills')`: only each skill's name and description sit in the system prompt; the model loads the full markdown on demand through an auto-registered `load_skill` tool - **Agents as directories** — `loadAgentDir('./my-agent')` turns `instructions.md`, `tools/*.ts`, `skills/` and `subagents//` into the same object `createAgent()` returns, so you can start with files and drop to code later — see [Agent directories](https://github.com/LinuxDevil/agent-sdk/blob/main/docs/agent-directories.md) - **Multi-agent delegation** — wrap a child agent as a tool with `createDelegateTool()`, with a `maxDepth` guard against delegation loops @@ -1481,6 +1482,45 @@ const usd = estimateCost({ inputTokens: used, outputTokens: 500 }, 'my-llama'); The built-in context windows and prices are a dated snapshot (see the retrieval date and sources at the top of `src/models/modelData.ts`). Providers change prices and models, so override entries with `registerModel` when you need billing-grade numbers. +### Usage and cost of a run + +Every `ExecutionResult` (and `agent.send()` result) carries `usage`, the running total of the whole run: + +```ts +import { AgentExecutor, formatUsage } from '@loushy/build-ai-agent'; + +const result = await AgentExecutor.execute({ agent, input: 'Compare 3 cities', provider }); + +result.usage.inputTokens; // all model calls of the run, delegated children included +result.usage.costUsd; // number, or undefined if any model used has no known price +result.usage.estimated; // true if some call reported no usage and was estimated +result.usage.byModel['gpt-4o-mini']; // { inputTokens, outputTokens, costUsd?, calls } +result.stepUsage?.[0]; // { step, model, usage, estimated, costUsd? } per model call +console.log(formatUsage(result.usage)); // 1,234 in / 567 out tokens · $0.0042 (2 model calls) +``` + +- **Reported vs estimated.** Providers report `promptTokens`/`completionTokens`; the built-in providers also pass on `cachedInputTokens`/`reasoningTokens` when the 'ai' SDK's provider metadata carries them (`usage.cachedInputTokens` and `usage.reasoningTokens` only appear then). A backend that reports nothing gives `undefined` usage, never zeros. For that step the executor falls back to `estimateTokens` and sets `usage.estimated` (a leading `~` in `formatUsage`). A custom provider should leave `GenerateResult.usage` unset rather than fill in zeros. +- **Cost.** `costUsd` is the sum of `estimateCost` per model. It is `undefined`, never a misleading partial sum, as soon as any model used has unknown pricing; `byModel` shows which ones are priced. +- **Delegation.** A delegated child's usage is added to the parent's totals and `byModel`, and is also shown on its own as `usage.delegated` (`{ inputTokens, outputTokens, totalTokens, costUsd, modelCalls, estimated, runs }`). +- **Resume.** A run resumed from a checkpoint, or after an approval, continues from the saved totals instead of restarting at zero. Checkpoints written by older versions start from their saved token counts with an unknown cost. +- **Events and traces.** The `finish` event (and every lifecycle event that carried `usage`) now carries the running totals; `text-complete` also has `stepUsage`, and `onLLMResponse` receives the call's usage as a third argument. The `chat` span's `gen_ai.usage.*` attributes use the same numbers, with `loushy.usage.estimated` set to `true` when they are estimates. +- **Streaming.** `agent.stream()` events carry the same accounting: `step.done` and `run.done` have `usage` with `inputTokens`, `outputTokens`, `estimated` and `costUsd` (run-level also `modelCalls`), alongside the older `promptTokens`/`completionTokens`. +- `promptTokens` and `completionTokens` on `usage` remain as deprecated aliases of `inputTokens` and `outputTokens`. + +Prices come from the model registry above, so to get a cost for a custom or fine-tuned model, register it under the id you pass as the model: + +```ts +import { registerModel } from '@loushy/build-ai-agent'; + +registerModel({ + id: 'ft:gpt-4o-mini:acme', + provider: 'openai', + contextWindow: 128000, + inputCostPerMTok: 0.3, + outputCostPerMTok: 1.2, +}); +``` + ## Flows, evals, observability and security - `FlowBuilder` / `FlowExecutor` - multi-step workflow graphs. diff --git a/src/execution/AgentExecutor.ts b/src/execution/AgentExecutor.ts index d423d5267..3acbdfd0f 100644 --- a/src/execution/AgentExecutor.ts +++ b/src/execution/AgentExecutor.ts @@ -12,6 +12,8 @@ import { ToolRegistry } from '../tools'; import { SandboxAdapter, NoopSandbox } from '../security/sandboxCore'; import { ApprovalStore, ExecutionSnapshot, PendingApproval } from './ApprovalGate'; import { CheckpointStore } from './checkpoint'; +import type { CallUsage, RunUsage, StepUsage } from '../models/usage'; +import { mergeDelegatedUsage } from './runUsage'; import { TraceExporter, withSpan } from './tracing'; import { agentRunSpanInit, @@ -23,6 +25,7 @@ import { HookRegistry } from './hooks'; import { isAbortError } from './errors'; import { PreparedToolCall, + ToolCallContext, ToolCallOutcome, parseToolArguments, runToolCall, @@ -37,6 +40,7 @@ import { import { buildTools, compactGenerateError, + GeneratedStep, generateInSpan, prepareGenerateRequest, providerErrorMessage, @@ -44,10 +48,10 @@ import { } from './generateStep'; import { AgentRunState, - addUsage, loadRunState, pushAbortedBatchResults, pushToolResult, + recordStep, saveStepCheckpoint, toExecutionResult, } from './agentRunState'; @@ -103,11 +107,10 @@ export interface ExecutionEvent { error?: string; }; finishReason?: ExecutionFinishReason; - usage?: { - promptTokens: number; - completionTokens: number; - totalTokens: number; - }; + /** Running usage of the whole run so far (LOU-V5); on `finish`, the final total. */ + usage?: RunUsage; + /** On `text-complete`: what the model call that produced the text spent (LOU-V5). */ + stepUsage?: StepUsage; error?: Error; /** * On an `abort` event: the `reason` of the aborted signal (a @@ -164,6 +167,12 @@ export interface ExecuteOptions { * 0 (a genuinely fresh run) when omitted. */ initialSteps?: number; + /** + * Usage already spent by earlier work this run continues (LOU-V5), used + * when there is no checkpoint to rehydrate it from - e.g. resume.ts + * resuming a run paused for approval. Ignored when a checkpoint is loaded. + */ + initialUsage?: RunUsage; /** * Tracing/observability hooks (LOU-E1/E2). These are invoked immediately * before/after the underlying provider.generate() call and each tool @@ -176,9 +185,14 @@ export interface ExecuteOptions { onLLMRequest?: (request: GenerateOptions) => void | Promise; /** * Invoked immediately after each provider.generate() call resolves, - * with the elapsed wall-clock time in milliseconds. + * with the elapsed wall-clock time in milliseconds and what the call + * spent (LOU-V5: reported usage, or an estimate flagged `estimated`). */ - onLLMResponse?: (response: GenerateResult, latencyMs: number) => void | Promise; + onLLMResponse?: ( + response: GenerateResult, + latencyMs: number, + usage: CallUsage + ) => void | Promise; /** Invoked immediately before each tool execution. */ onToolCall?: (toolCall: ToolCall) => void | Promise; /** @@ -368,11 +382,13 @@ export interface ExecutionResult { text: string; messages: Message[]; toolCalls: ToolCall[]; - usage: { - promptTokens: number; - completionTokens: number; - totalTokens: number; - }; + /** + * Tokens, cost and per-model breakdown of the whole run, including + * delegated children and steps before a resume (LOU-V5). See `formatUsage()`. + */ + usage: RunUsage; + /** One entry per model call of this process's run, in order (LOU-V5). */ + stepUsage?: StepUsage[]; finishReason: ExecutionFinishReason; steps: number; approvalId?: string; @@ -470,7 +486,7 @@ export class AgentExecutor { options: ExecuteOptions, agentSpanId: string ): Promise { - const { agent, toolRegistry, maxSteps = 10, onEvent, signal } = options; + const { agent, toolRegistry, onEvent } = options; // Emit start event this.emitEvent(onEvent, { @@ -485,6 +501,18 @@ export class AgentExecutor { const state = await loadRunState(options); + return this.runSteps(options, state, tools, agentSpanId); + } + + /** The generate -> tools loop of runAgentLoop(), run once `state` is loaded. */ + private static async runSteps( + options: ExecuteOptions, + state: AgentRunState, + tools: ToolDefinition[], + agentSpanId: string + ): Promise { + const { maxSteps = 10, signal } = options; + // Execution loop with tool calling. LOU-V1: the signal is checked // before every model call (here) and every tool call (runToolCalls()). while (state.steps < maxSteps) { @@ -550,17 +578,18 @@ export class AgentExecutor { tools: ToolDefinition[], agentSpanId: string ): Promise<'continue' | 'stop' | ExecutionResult> { - const result = await this.generateOrSurfaceError(options, state, tools, agentSpanId); - if (!result) { + const generatedStep = await this.generateOrSurfaceError(options, state, tools, agentSpanId); + if (!generatedStep) { return 'continue'; } + const { generated: result, measured } = generatedStep; // A turn produced a real result - any pending "the last thing that // happened was a provider failure" tracking no longer applies. state.lastSurfacedProviderError = undefined; - // Update usage - addUsage(state.usage, result.usage); + // Update usage (LOU-V5) + const stepUsage = recordStep(state, measured); // Handle text response if (result.text) { @@ -569,6 +598,7 @@ export class AgentExecutor { type: 'text-complete', timestamp: new Date(), text: result.text, + stepUsage, }); } @@ -613,7 +643,7 @@ export class AgentExecutor { state: AgentRunState, tools: ToolDefinition[], agentSpanId: string - ): Promise { + ): Promise { const generateRequest = await prepareGenerateRequest(options, state.messages, tools); try { @@ -787,7 +817,9 @@ export class AgentExecutor { sessionId, state.messages, signal, - onPrepared + onPrepared, + // LOU-V5: a delegated child's usage is added to this run's totals. + (child) => mergeDelegatedUsage(state.usage, child) ); const parsedArgs = executed.args === undefined @@ -842,6 +874,7 @@ export class AgentExecutor { pendingToolCall: pending, steps: state.steps, sessionId, + usage: structuredClone(state.usage), }; await approvalStore.save(pending, snapshot); @@ -952,7 +985,8 @@ export class AgentExecutor { sessionId?: string, messages: Message[] = [], signal?: AbortSignal, - onPrepared?: (prepared: PreparedToolCall) => void + onPrepared?: (prepared: PreparedToolCall) => void, + onDelegatedUsage?: ToolCallContext['onDelegatedUsage'] ): Promise { return runToolCall( toolCall, @@ -966,6 +1000,7 @@ export class AgentExecutor { sessionId, messages, signal, + onDelegatedUsage, }, onPrepared ); diff --git a/src/execution/ApprovalGate.ts b/src/execution/ApprovalGate.ts index 9ba0ebd3d..b6502e901 100644 --- a/src/execution/ApprovalGate.ts +++ b/src/execution/ApprovalGate.ts @@ -6,6 +6,7 @@ import { Message } from '../providers'; import { AgentConfig } from '../types'; import { StorageService } from '../storage'; +import type { RunUsage } from '../models/usage'; /** * A tool call that is waiting on a human decision before it can execute. @@ -48,6 +49,8 @@ export interface ExecutionSnapshot { * execution at all. */ sessionId?: string; + /** Usage the paused run had spent (LOU-V5), so the resumed run continues its totals. */ + usage?: RunUsage; } /** diff --git a/src/execution/DelegationTool.ts b/src/execution/DelegationTool.ts index 059c22d0c..d0fc82291 100644 --- a/src/execution/DelegationTool.ts +++ b/src/execution/DelegationTool.ts @@ -11,6 +11,7 @@ import { AgentExecutor, PropagatingToolError } from './AgentExecutor'; import { LLMProvider, Message } from '../providers'; import { AgentConfig, ToolDescriptor } from '../types'; import { ToolRegistry } from '../tools'; +import type { ToolRunContext } from './sandboxGuard'; /** * Thrown when a delegation chain exceeds the configured maxDepth without @@ -119,7 +120,7 @@ export function createDelegateTool(opts: DelegateAgentOptions): ToolDescriptor { task: string; context?: Message[]; }, - options?: { abortSignal?: AbortSignal } + options?: { abortSignal?: AbortSignal } & ToolRunContext ): Promise => { const currentDepth = delegationDepthStorage.getStore() ?? 0; @@ -143,9 +144,13 @@ export function createDelegateTool(opts: DelegateAgentOptions): ToolDescriptor { signal: options?.abortSignal, }); + // LOU-V5: the child's full usage rolls up into the parent run's + // totals (`usage.delegated`); the model only sees the token counts. + options?.onDelegatedUsage?.(result.usage); + const { promptTokens, completionTokens, totalTokens } = result.usage; return { text: result.text, - usage: result.usage, + usage: { promptTokens, completionTokens, totalTokens }, }; }); }, diff --git a/src/execution/agentEvents.ts b/src/execution/agentEvents.ts index edee099c6..c87bd8079 100644 --- a/src/execution/agentEvents.ts +++ b/src/execution/agentEvents.ts @@ -18,9 +18,19 @@ export const AGENT_EVENT_SCHEMA_VERSION = 1 as const; /** Token usage, as reported on `step.done` and `run.done`. */ export interface AgentEventUsage { + /** Alias of `inputTokens`. */ promptTokens: number; + /** Alias of `outputTokens`. */ completionTokens: number; totalTokens: number; + inputTokens: number; + outputTokens: number; + /** `true` when the provider reported no usage and the tokens are an `estimateTokens` estimate. */ + estimated: boolean; + /** USD, or absent when a model used has unknown pricing (see `registerModel`). */ + costUsd?: number; + /** On `run.done`: model calls of the whole run. */ + modelCalls?: number; } /** A JSON-safe error: the `name` and `message` of the original error. */ diff --git a/src/execution/agentRun.ts b/src/execution/agentRun.ts index b7875980c..598aae6ae 100644 --- a/src/execution/agentRun.ts +++ b/src/execution/agentRun.ts @@ -27,6 +27,8 @@ import { } from './agentEvents'; import { parseToolArguments } from './toolCallExecution'; import { canStream, generateViaStream } from './streamStep'; +import { measureUsage } from './runUsage'; +import type { Usage } from '../models/usage'; /** * The handle returned by `agent.stream()` and `AgentExecutor.stream()`. @@ -97,11 +99,25 @@ function toEventError(error: unknown): AgentEventError { }; } -function toEventUsage(usage: AgentEventUsage): AgentEventUsage { +/** Usage of one step, or (with `modelCalls`) of the whole run, as event usage. */ +function toEventUsage(usage: { + inputTokens: number; + outputTokens: number; + totalTokens: number; + estimated: boolean; + costUsd?: number; + modelCalls?: number; +}): AgentEventUsage { + const { inputTokens, outputTokens, totalTokens, estimated, costUsd, modelCalls } = usage; return { - promptTokens: usage.promptTokens, - completionTokens: usage.completionTokens, - totalTokens: usage.totalTokens, + promptTokens: inputTokens, + completionTokens: outputTokens, + totalTokens, + inputTokens, + outputTokens, + estimated, + ...(costUsd !== undefined && { costUsd }), + ...(modelCalls !== undefined && { modelCalls }), }; } @@ -125,7 +141,7 @@ class AgentRunImpl implements AgentRun { private closed = false; private iterated = false; private wake: (() => void) | undefined; - private stepResult: GenerateResult | undefined; + private stepResult: { finishReason: GenerateResult['finishReason']; usage: Usage; estimated: boolean; costUsd?: number } | undefined; private lastError: unknown; constructor(start: RunStarter, signal: AbortSignal | undefined) { @@ -265,12 +281,12 @@ class AgentRunImpl implements AgentRun { this.emit({ type: 'step.start', step }); }, stepDone: (step, finishReason) => { - const usage = this.stepResult?.usage; + const measured = this.stepResult; this.emit({ type: 'step.done', step, - finishReason: finishReason ?? this.stepResult?.finishReason ?? 'error', - ...(usage && { usage: toEventUsage(usage) }), + finishReason: finishReason ?? measured?.finishReason ?? 'error', + ...(measured && { usage: toEventUsage({ ...measured.usage, estimated: measured.estimated, costUsd: measured.costUsd }) }), }); }, approvalRequested: (pending) => @@ -282,8 +298,10 @@ class AgentRunImpl implements AgentRun { args: toJsonValue(pending.args) as Record, }), generate: async (provider, request) => { - this.stepResult = await this.generateStep(provider, request); - return this.stepResult; + const generated = await this.generateStep(provider, request); + const measured = measureUsage(request.model ?? provider.name, request.messages, generated); + this.stepResult = { finishReason: generated.finishReason, ...measured, usage: measured.usage }; + return generated; }, }; } diff --git a/src/execution/agentRunState.ts b/src/execution/agentRunState.ts index 18cad33f5..ffe4bd829 100644 --- a/src/execution/agentRunState.ts +++ b/src/execution/agentRunState.ts @@ -9,16 +9,18 @@ import { Message, ToolCall } from '../providers'; import { AgentConfig } from '../types'; import { Checkpoint } from './checkpoint'; import { CompactedLLMProviderError } from './errors'; +import type { CallUsage, RunUsage, StepUsage } from '../models/usage'; +import { emptyRunUsage, recordStepUsage, restoreRunUsage } from './runUsage'; import type { ExecuteOptions, ExecutionResult } from './AgentExecutor'; import type { ToolCallOutcome } from './toolCallExecution'; import type { UnrecordedToolCall } from './toolBatch'; -type TokenUsage = ExecutionResult['usage']; - export interface AgentRunState { messages: Message[]; toolCalls: ToolCall[]; - usage: TokenUsage; + usage: RunUsage; + /** One entry per model call (LOU-V5). */ + stepUsage: StepUsage[]; steps: number; businessState: unknown; finalText: string; @@ -71,7 +73,7 @@ function buildMessages( type InitialRunState = Pick< AgentRunState, - 'messages' | 'toolCalls' | 'usage' | 'steps' | 'businessState' + 'messages' | 'toolCalls' | 'usage' | 'stepUsage' | 'steps' | 'businessState' >; /** Rehydrates run state from a checkpoint left by an earlier process. */ @@ -88,7 +90,9 @@ function stateFromCheckpoint(checkpoint: Checkpoint, businessState: unknown): In return { messages: [...checkpoint.messages], toolCalls: [...(checkpoint.toolCalls as ToolCall[])], - usage: { ...checkpoint.usage }, + // LOU-V5: continue from the checkpointed totals (older checkpoints: token counts only). + usage: restoreRunUsage(checkpoint.usage), + stepUsage: [...(checkpoint.stepUsage ?? [])], steps: checkpoint.stepIndex, businessState: businessState === undefined ? checkpoint.businessState : businessState, }; @@ -100,11 +104,8 @@ function freshState(options: ExecuteOptions): InitialRunState { return { messages: buildMessages(agent, input, skipSystemPromptInjection), toolCalls: [], - usage: { - promptTokens: 0, - completionTokens: 0, - totalTokens: 0, - }, + usage: options.initialUsage ? restoreRunUsage(options.initialUsage) : emptyRunUsage(), + stepUsage: [], steps: initialSteps ?? 0, businessState, }; @@ -128,11 +129,12 @@ export async function loadRunState(options: ExecuteOptions): Promise { expect(events.find((e) => e.type === 'step.done')).toMatchObject({ finishReason: 'tool_calls', usage }); const result = await run.result; expect(result.text).toBe('Found it.'); - expect(result.usage).toEqual({ promptTokens: 8, completionTokens: 4, totalTokens: 12 }); + expect(result.usage).toMatchObject({ promptTokens: 8, completionTokens: 4, totalTokens: 12 }); }); }); diff --git a/src/execution/checkpoint.ts b/src/execution/checkpoint.ts index 1ceda3490..26457286e 100644 --- a/src/execution/checkpoint.ts +++ b/src/execution/checkpoint.ts @@ -5,6 +5,8 @@ import { Message } from '../providers'; import { StorageService } from '../storage'; +import type { StepUsage } from '../models/usage'; +import type { CheckpointUsage } from './runUsage'; /** * A snapshot of an in-progress agent run, saved after each tool result so @@ -16,11 +18,14 @@ export interface Checkpoint { stepIndex: number; messages: Message[]; toolCalls: unknown[]; - usage: { - promptTokens: number; - completionTokens: number; - totalTokens: number; - }; + /** + * Running usage of the run so far (LOU-V5). A resumed run continues from + * these totals. Checkpoints written before LOU-V5 hold only the three + * token counts and still load (their tokens are kept, cost is unknown). + */ + usage: CheckpointUsage; + /** Per-call usage so far (LOU-V5); absent on older checkpoints. */ + stepUsage?: StepUsage[]; finishReason?: string; /** * LOU-T1: opaque, consumer-owned business/domain state co-located with diff --git a/src/execution/genAiSpans.ts b/src/execution/genAiSpans.ts index 9dd864329..8737d2db3 100644 --- a/src/execution/genAiSpans.ts +++ b/src/execution/genAiSpans.ts @@ -15,8 +15,10 @@ import type { GenerateOptions, GenerateResult, LLMProvider, Message, ToolCall } from '../providers'; import type { ToolRegistry } from '../tools'; import type { AgentConfig } from '../types'; -import { ErrorAttr, GenAiAttr, GenAiOperation, LegacyAttr, TOOL_TYPE_FUNCTION, CAPTURE_CONTENT_ENV } from './semconv'; +import { ErrorAttr, GenAiAttr, GenAiOperation, LegacyAttr, SdkAttr, TOOL_TYPE_FUNCTION, CAPTURE_CONTENT_ENV } from './semconv'; import type { Span, SpanKind } from './tracing'; +import { normalizeUsage } from '../models/usage'; +import type { Usage } from '../models/usage'; /** Name, kind and starting attributes for a span about to be opened. */ interface SpanInit { @@ -158,8 +160,18 @@ function responseModel(generated: GenerateResult): string | undefined { return typeof model === 'string' ? model : undefined; } -/** Records usage, finish reason, response model (and output content) on a `chat` span. */ -export function recordLlmResult(span: Span, generated: GenerateResult, captureContent: boolean): void { +/** + * Records usage, finish reason, response model (and output content) on a `chat` span. + * Usage is the same normalized `Usage` the run totals use: pass the executor's + * `measured` figures (which may be estimates), else what the provider reported. + */ +export function recordLlmResult( + span: Span, + generated: GenerateResult, + captureContent: boolean, + measured?: { usage: Usage; estimated: boolean } +): void { + const usage = measured?.usage ?? normalizeUsage(generated.usage); const parts: unknown[] = generated.text ? [textPart(generated.text)] : []; parts.push(...(generated.toolCalls ?? []).map(toolCallPart)); // Token counts and finish reason are never redacted. @@ -168,11 +180,12 @@ export function recordLlmResult(span: Span, generated: GenerateResult, captureCo ...defined({ [GenAiAttr.RESPONSE_MODEL]: responseModel(generated), [GenAiAttr.RESPONSE_FINISH_REASONS]: [genAiFinishReason(generated.finishReason)], - [GenAiAttr.USAGE_INPUT_TOKENS]: generated.usage.promptTokens, - [GenAiAttr.USAGE_OUTPUT_TOKENS]: generated.usage.completionTokens, - [LegacyAttr.PROMPT_TOKENS]: generated.usage.promptTokens, - [LegacyAttr.COMPLETION_TOKENS]: generated.usage.completionTokens, - [LegacyAttr.TOTAL_TOKENS]: generated.usage.totalTokens, + [GenAiAttr.USAGE_INPUT_TOKENS]: usage?.inputTokens, + [GenAiAttr.USAGE_OUTPUT_TOKENS]: usage?.outputTokens, + [LegacyAttr.PROMPT_TOKENS]: usage?.inputTokens, + [LegacyAttr.COMPLETION_TOKENS]: usage?.outputTokens, + [LegacyAttr.TOTAL_TOKENS]: usage?.totalTokens, + [SdkAttr.USAGE_ESTIMATED]: measured?.estimated ? true : undefined, [LegacyAttr.FINISH_REASON]: generated.finishReason, [GenAiAttr.OUTPUT_MESSAGES]: captureContent ? JSON.stringify([{ role: 'assistant', parts }]) diff --git a/src/execution/generateStep.ts b/src/execution/generateStep.ts index 27966aee0..8406815fa 100644 --- a/src/execution/generateStep.ts +++ b/src/execution/generateStep.ts @@ -9,6 +9,8 @@ import { GenerateOptions, GenerateResult, Message, ToolDefinition } from '../pro import { AgentConfig } from '../types'; import { ToolRegistry } from '../tools'; import { withSpan } from './tracing'; +import type { CallUsage } from '../models/usage'; +import { measureUsage } from './runUsage'; import { llmSpanInit, recordLlmResult, resolveCaptureContent } from './genAiSpans'; import { GenerateHookContext } from './hooks'; import { @@ -69,6 +71,11 @@ function generateHookContext( }; } +/** The model a call is made with: agent.settings.model > the provider's configured model > unknown. */ +function resolveModel(options: ExecuteOptions): string | undefined { + return options.agent.settings?.model || options.provider.defaultModel; +} + /** * Builds the next provider.generate() request and runs the onLLMRequest * callback and preGenerate hooks for it. Errors thrown here are NOT @@ -79,12 +86,12 @@ export async function prepareGenerateRequest( messages: Message[], tools: ToolDefinition[] ): Promise { - const { agent, provider, temperature, maxTokens, onLLMRequest, hooks, signal } = options; + const { temperature, maxTokens, onLLMRequest, hooks, signal } = options; const generateRequest: GenerateOptions = { // agent.settings.model > the model the provider was configured with > // undefined (the provider then applies its own built-in default). - model: agent.settings?.model || provider.defaultModel, + model: resolveModel(options), messages, temperature, maxTokens, @@ -104,6 +111,12 @@ export async function prepareGenerateRequest( return generateRequest; } +/** A generate() reply with what the call spent (LOU-V5). */ +export interface GeneratedStep { + generated: GenerateResult; + measured: CallUsage; +} + /** * Runs provider.generate() inside a `chat {model}` span parented to the * run's `invoke_agent` span, followed by the onLLMResponse callback and @@ -114,7 +127,7 @@ export function generateInSpan( generateRequest: GenerateOptions, messages: Message[], agentSpanId: string -): Promise { +): Promise { const { provider, exporter, onLLMResponse, hooks, redactContent = false } = options; const captureContent = resolveCaptureContent(options.captureContent); const init = llmSpanInit(provider, generateRequest, { redactContent, captureContent }); @@ -133,10 +146,11 @@ export function generateInSpan( : await provider.generate(generateRequest); const llmLatencyMs = Date.now() - llmStart; - recordLlmResult(llmSpan, generated, captureContent); + const measured = measureUsage(resolveModel(options) ?? provider.name, messages, generated); + recordLlmResult(llmSpan, generated, captureContent, measured); if (onLLMResponse) { - await onLLMResponse(generated, llmLatencyMs); + await onLLMResponse(generated, llmLatencyMs, measured); } if (hooks) { @@ -146,7 +160,7 @@ export function generateInSpan( ); } - return generated; + return { generated, measured }; }, agentSpanId, init.kind diff --git a/src/execution/resume.ts b/src/execution/resume.ts index 6d9ed28ad..c450c6a18 100644 --- a/src/execution/resume.ts +++ b/src/execution/resume.ts @@ -57,6 +57,7 @@ export type ResumeExecuteOptions = Omit< | 'toolRegistry' | 'skipSystemPromptInjection' | 'initialSteps' + | 'initialUsage' | 'sessionId' | 'checkpointStore' >; @@ -383,5 +384,7 @@ function continueResumedRun( // snapshot.steps is the step count AgentExecutor.execute() had already // reached (see ExecutionSnapshot) at the moment it paused for approval. initialSteps: snapshot.steps, + // LOU-V5: and usage totals, rather than restarting them at zero. + initialUsage: snapshot.usage, }); } diff --git a/src/execution/runUsage.ts b/src/execution/runUsage.ts new file mode 100644 index 000000000..1a7483ba7 --- /dev/null +++ b/src/execution/runUsage.ts @@ -0,0 +1,154 @@ +/** + * Accumulation of a run's usage (LOU-V5): per-call measurement (reported or + * estimated), the running RunUsage, delegated-child roll-up, and + * rehydration from a checkpoint. + */ + +import type { GenerateResult, Message } from '../providers'; +import { estimateCost, estimateTokens } from '../models'; +import { normalizeUsage } from '../models/usage'; +import type { CallUsage, DelegatedUsage, RunUsage, Usage } from '../models/usage'; + +/** Model id that pre-LOU-V5 checkpoint tokens are attributed to. */ +const UNKNOWN_MODEL = 'unknown'; + +/** A fresh, empty run total. */ +export function emptyRunUsage(): RunUsage { + return { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + costUsd: 0, + modelCalls: 0, + estimated: false, + byModel: {}, + promptTokens: 0, + completionTokens: 0, + }; +} + +function estimatedUsage(model: string, messages: Message[], generated: GenerateResult): Usage { + const calls = (generated.toolCalls ?? []).map((c) => `${c.function.name} ${c.function.arguments}`); + const inputTokens = estimateTokens(messages, { model }); + const outputTokens = estimateTokens([generated.text, ...calls].join('\n'), { model }); + return { inputTokens, outputTokens, totalTokens: inputTokens + outputTokens }; +} + +/** + * What a call spent: the provider's reported usage, or - when it reported + * none - an `estimateTokens` estimate of the request and reply (flagged). + */ +export function measureUsage( + model: string, + requestMessages: Message[], + generated: GenerateResult +): CallUsage { + const reported = normalizeUsage(generated.usage); + const usage = reported ?? estimatedUsage(model, requestMessages, generated); + return { model, usage, estimated: !reported, costUsd: estimateCost(usage, model) }; +} + +type OptionalCount = 'cachedInputTokens' | 'reasoningTokens'; + +function addOptional(run: RunUsage, key: OptionalCount, add: number | undefined): void { + if (add !== undefined) run[key] = (run[key] ?? 0) + add; +} + +/** Recomputes per-model and total cost from the token counts, and the legacy aliases. */ +function refreshDerived(run: RunUsage): void { + let total: number | undefined = 0; + for (const [model, entry] of Object.entries(run.byModel)) { + entry.costUsd = estimateCost(entry, model); + total = total === undefined || entry.costUsd === undefined ? undefined : total + entry.costUsd; + } + run.costUsd = total; + run.promptTokens = run.inputTokens; + run.completionTokens = run.outputTokens; +} + +function addToModel( + run: RunUsage, + model: string, + tokens: { inputTokens: number; outputTokens: number }, + calls: number +): void { + const entry = (run.byModel[model] ??= { inputTokens: 0, outputTokens: 0, calls: 0 }); + entry.inputTokens += tokens.inputTokens; + entry.outputTokens += tokens.outputTokens; + entry.calls += calls; +} + +/** Adds one model call to the run total (mutates). */ +export function recordStepUsage(run: RunUsage, step: CallUsage): void { + const { usage } = step; + run.inputTokens += usage.inputTokens; + run.outputTokens += usage.outputTokens; + run.totalTokens += usage.totalTokens; + run.modelCalls += 1; + run.estimated ||= step.estimated; + addOptional(run, 'cachedInputTokens', usage.cachedInputTokens); + addOptional(run, 'reasoningTokens', usage.reasoningTokens); + addToModel(run, step.model, usage, 1); + refreshDerived(run); +} + +function delegatedTotals(before: DelegatedUsage | undefined, child: RunUsage): DelegatedUsage { + const prior = before ?? { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + costUsd: 0 as number | undefined, + modelCalls: 0, + estimated: false, + runs: 0, + }; + return { + inputTokens: prior.inputTokens + child.inputTokens, + outputTokens: prior.outputTokens + child.outputTokens, + totalTokens: prior.totalTokens + child.totalTokens, + costUsd: + prior.costUsd === undefined || child.costUsd === undefined + ? undefined + : prior.costUsd + child.costUsd, + modelCalls: prior.modelCalls + child.modelCalls, + estimated: prior.estimated || child.estimated, + runs: prior.runs + 1, + }; +} + +/** Rolls a finished child agent run's total into the parent's (mutates the parent). */ +export function mergeDelegatedUsage(run: RunUsage, child: RunUsage): void { + run.inputTokens += child.inputTokens; + run.outputTokens += child.outputTokens; + run.totalTokens += child.totalTokens; + run.modelCalls += child.modelCalls; + run.estimated ||= child.estimated; + addOptional(run, 'cachedInputTokens', child.cachedInputTokens); + addOptional(run, 'reasoningTokens', child.reasoningTokens); + for (const [model, entry] of Object.entries(child.byModel)) addToModel(run, model, entry, entry.calls); + run.delegated = delegatedTotals(run.delegated, child); + refreshDerived(run); +} + +/** What a checkpoint stores of a run's usage: always the pre-LOU-V5 fields, plus the rest of RunUsage on newer ones. */ +export type CheckpointUsage = Pick & + Partial; + +/** + * Rehydrates a checkpointed run total. Checkpoints written before LOU-V5 + * only carry token counts: those are kept (attributed to the `unknown` + * model, so `costUsd` is `undefined` rather than a misleading partial sum). + */ +export function restoreRunUsage(saved: CheckpointUsage): RunUsage { + const run = emptyRunUsage(); + if (saved.byModel && typeof saved.modelCalls === 'number') { + Object.assign(run, structuredClone(saved)); + } else { + run.inputTokens = saved.promptTokens; + run.outputTokens = saved.completionTokens; + run.totalTokens = saved.totalTokens; + if (saved.totalTokens > 0) addToModel(run, UNKNOWN_MODEL, run, 0); + } + refreshDerived(run); + return run; +} diff --git a/src/execution/sandboxGuard.ts b/src/execution/sandboxGuard.ts index ff321aadc..1c3343544 100644 --- a/src/execution/sandboxGuard.ts +++ b/src/execution/sandboxGuard.ts @@ -21,6 +21,13 @@ import type { ToolExecutionOptions } from 'ai'; import { ToolDescriptor } from '../types'; import { SandboxAdapter } from '../security/sandboxCore'; +import type { RunUsage } from '../models/usage'; + +/** Extra context the executor hands a tool next to the 'ai' SDK's own execute options (LOU-V5). */ +export interface ToolRunContext { + /** Called by the delegate tool with a finished child run's usage, so the parent run adds it to its totals. */ + onDelegatedUsage?: (usage: RunUsage) => void; +} /** * Execute `toolDesc` against `args`, honoring `requiresSandbox`: @@ -46,7 +53,8 @@ export async function executeToolWithSandboxGuard( toolDesc: ToolDescriptor, args: Record, sandbox: SandboxAdapter, - signal?: AbortSignal + signal?: AbortSignal, + runContext?: ToolRunContext ): Promise { if (toolDesc.requiresSandbox) { if (!toolDesc.sandboxExecute) { @@ -62,6 +70,6 @@ export async function executeToolWithSandboxGuard( // The 'ai' SDK types toolCallId/messages as required, but tools invoked // here are not part of an 'ai' SDK generation, so only abortSignal is set. - const executeOptions = { abortSignal: signal } as ToolExecutionOptions; + const executeOptions = { abortSignal: signal, ...runContext } as ToolExecutionOptions; return toolDesc.tool.execute ? toolDesc.tool.execute(args, executeOptions) : null; } diff --git a/src/execution/semconv.ts b/src/execution/semconv.ts index 7752d59c1..ef42b8b5e 100644 --- a/src/execution/semconv.ts +++ b/src/execution/semconv.ts @@ -66,6 +66,12 @@ export const ErrorAttr = { /** `gen_ai.tool.type` value for tools executed by this SDK's host process. */ export const TOOL_TYPE_FUNCTION = 'function'; +/** Attributes specific to this SDK (not part of the GenAI semantic conventions). */ +export const SdkAttr = { + /** `true` on a `chat` span whose token counts are estimates because the provider reported none (LOU-V5). */ + USAGE_ESTIMATED: 'loushy.usage.estimated', +} as const; + /** * The attribute names this SDK emitted before LOU-D9. They are still emitted * alongside the GenAI ones (dual-emit) and are DEPRECATED. diff --git a/src/execution/toolCallExecution.ts b/src/execution/toolCallExecution.ts index b0613e965..e02312a52 100644 --- a/src/execution/toolCallExecution.ts +++ b/src/execution/toolCallExecution.ts @@ -9,6 +9,7 @@ import { AgentConfig, ToolDescriptor } from '../types'; import { ToolRegistry } from '../tools'; import { SandboxAdapter } from '../security/sandboxCore'; import { executeToolWithSandboxGuard } from './sandboxGuard'; +import type { ToolRunContext } from './sandboxGuard'; import { HookRegistry, ToolCallHookContext } from './hooks'; import { toolErrorMessage, toolErrorResult } from './propagatingToolError'; import { ToolArgumentsValidationError, validateToolArguments } from './toolArgsValidation'; @@ -39,6 +40,8 @@ export interface ToolCallContext { messages: Message[]; /** LOU-V1: the run's signal, handed to the tool as `abortSignal`. */ signal?: AbortSignal; + /** LOU-V5: where a delegated child's usage is reported (see ToolRunContext). */ + onDelegatedUsage?: ToolRunContext['onDelegatedUsage']; } /** @@ -202,7 +205,9 @@ async function executePrepared( if (prepared.requiresApproval) { return approvalOutcome(prepared); } - return doExecuteToolCall(prepared.toolCall, ctx.toolRegistry, ctx.sandbox, prepared.args, ctx.signal); + return doExecuteToolCall(prepared.toolCall, ctx.toolRegistry, ctx.sandbox, prepared.args, ctx.signal, { + onDelegatedUsage: ctx.onDelegatedUsage, + }); } /** @@ -300,7 +305,8 @@ async function doExecuteToolCall( toolRegistry: ToolRegistry | undefined, sandbox: SandboxAdapter, overrideArgs?: Record, - signal?: AbortSignal + signal?: AbortSignal, + runContext?: ToolRunContext ): Promise { if (!toolRegistry) { return toolFailure(toolCall, 'No tool registry available'); @@ -331,7 +337,8 @@ async function doExecuteToolCall( toolDesc, args, sandbox, - signal + signal, + runContext ); return { diff --git a/src/execution/usage.test.ts b/src/execution/usage.test.ts new file mode 100644 index 000000000..489d47944 --- /dev/null +++ b/src/execution/usage.test.ts @@ -0,0 +1,421 @@ +import { describe, it, expect } from 'vitest'; +import { z } from 'zod'; +import { AgentExecutor } from './AgentExecutor'; +import type { ExecutionEvent } from './AgentExecutor'; +import { createDelegateTool } from './DelegationTool'; +import { resumeAfterApproval } from './resume'; +import type { AgentEvent } from './agentEvents'; +import type { ApprovalStore, ExecutionSnapshot, PendingApproval } from './ApprovalGate'; +import type { Span, TraceExporter } from './tracing'; +import type { Checkpoint, CheckpointStore } from './checkpoint'; +import { ToolRegistry } from '../tools'; +import { defineTool } from '../tools/defineTool'; +import { AgentType } from '../types'; +import type { AgentConfig } from '../types'; +import { mockModel } from '../testing'; +import { formatUsage, normalizeUsage } from '../models'; + +const echo = defineTool({ + name: 'echo', + description: 'echo', + input: z.object({}), + execute: async () => 'echoed', +}); + +const agent = (overrides: Partial = {}): AgentConfig => ({ + id: 'a', + name: 'Agent', + agentType: AgentType.SmartAssistant, + tools: { echo: { tool: 'echo' } }, + ...overrides, +}); + +const registry = (): ToolRegistry => { + const r = new ToolRegistry(); + r.registerMany([echo]); + return r; +}; + +const callEcho = { toolCalls: [{ name: 'echo' }] }; + +describe('run usage (LOU-V5)', () => { + it('accumulates tokens, calls and cost across steps of a priced model', async () => { + const provider = mockModel( + [ + { ...callEcho, usage: { inputTokens: 1000, outputTokens: 200 } }, + { text: 'done', usage: { inputTokens: 1500, outputTokens: 300 } }, + ], + { defaultModel: 'gpt-4o-mini' } + ); + + const result = await AgentExecutor.execute({ agent: agent(), input: 'go', provider, toolRegistry: registry() }); + + expect(result.usage).toMatchObject({ + inputTokens: 2500, + outputTokens: 500, + totalTokens: 3000, + modelCalls: 2, + estimated: false, + // legacy aliases still work + promptTokens: 2500, + completionTokens: 500, + }); + expect(result.usage.costUsd).toBeCloseTo((2500 * 0.15 + 500 * 0.6) / 1e6, 10); + expect(result.usage.byModel).toEqual({ + 'gpt-4o-mini': { inputTokens: 2500, outputTokens: 500, costUsd: result.usage.costUsd, calls: 2 }, + }); + expect(result.stepUsage).toEqual([ + expect.objectContaining({ step: 1, model: 'gpt-4o-mini', estimated: false, usage: { inputTokens: 1000, outputTokens: 200, totalTokens: 1200 } }), + expect.objectContaining({ step: 2, model: 'gpt-4o-mini', estimated: false, usage: { inputTokens: 1500, outputTokens: 300, totalTokens: 1800 } }), + ]); + expect(formatUsage(result.usage)).toBe('2,500 in / 500 out tokens · $0.0007 (2 model calls)'); + }); + + it('leaves costUsd undefined (not a partial sum) when any model has unknown pricing', async () => { + const provider = mockModel( + [ + { ...callEcho, usage: { inputTokens: 100, outputTokens: 10 } }, + { text: 'done', usage: { inputTokens: 100, outputTokens: 10 } }, + ], + { defaultModel: 'my-private-model' } + ); + + const result = await AgentExecutor.execute({ agent: agent(), input: 'go', provider, toolRegistry: registry() }); + + expect(result.usage.costUsd).toBeUndefined(); + expect(result.usage.byModel['my-private-model'].costUsd).toBeUndefined(); + expect(result.usage.totalTokens).toBe(220); + expect(formatUsage(result.usage)).toBe('200 in / 20 out tokens (2 model calls)'); + }); + + it('estimates tokens and flags `estimated` when the provider reports no usage', async () => { + const provider = mockModel(['hello there']); + + const result = await AgentExecutor.execute({ agent: agent({ prompt: 'be brief' }), input: 'hi', provider }); + + expect(result.usage.estimated).toBe(true); + expect(result.usage.inputTokens).toBeGreaterThan(0); + expect(result.usage.outputTokens).toBeGreaterThan(0); + expect(result.usage.modelCalls).toBe(1); + expect(result.stepUsage?.[0].estimated).toBe(true); + expect(formatUsage(result.usage).startsWith('~')).toBe(true); + }); + + it('only flags the run estimated when some step was estimated, and still sums the reported ones', async () => { + const provider = mockModel([ + { ...callEcho, usage: { inputTokens: 50, outputTokens: 5 } }, + 'no usage here', + ]); + + const result = await AgentExecutor.execute({ agent: agent(), input: 'go', provider, toolRegistry: registry() }); + + expect(result.stepUsage?.map((s) => s.estimated)).toEqual([false, true]); + expect(result.usage.estimated).toBe(true); + expect(result.usage.inputTokens).toBeGreaterThan(50); + }); + + it('breaks usage down by model when agents override settings.model', async () => { + const child = agent({ id: 'child', name: 'Child', tools: undefined, settings: { model: 'gpt-4o' } }); + const parentProvider = mockModel( + [ + { toolCalls: [{ name: 'ask', args: { task: 'sub' } }], usage: { inputTokens: 10, outputTokens: 5 } }, + { text: 'done', usage: { inputTokens: 20, outputTokens: 5 } }, + ], + { defaultModel: 'gpt-4o-mini' } + ); + const childProvider = mockModel([{ text: 'sub done', usage: { inputTokens: 100, outputTokens: 50 } }], { + defaultModel: 'gpt-4o-mini', + }); + const tools = new ToolRegistry(); + tools.register('ask', createDelegateTool({ agent: child, provider: childProvider })); + + const result = await AgentExecutor.execute({ + agent: agent({ tools: { ask: { tool: 'ask' } } }), + input: 'go', + provider: parentProvider, + toolRegistry: tools, + }); + + expect(childProvider.calls[0].model).toBe('gpt-4o'); + expect(Object.keys(result.usage.byModel).sort()).toEqual(['gpt-4o', 'gpt-4o-mini']); + expect(result.usage.byModel['gpt-4o']).toMatchObject({ inputTokens: 100, outputTokens: 50, calls: 1 }); + expect(result.usage.byModel['gpt-4o-mini']).toMatchObject({ inputTokens: 30, outputTokens: 10, calls: 2 }); + const expectedCost = (100 * 2.5 + 50 * 10 + 30 * 0.15 + 10 * 0.6) / 1e6; + expect(result.usage.costUsd).toBeCloseTo(expectedCost, 10); + }); + + it('rolls a delegated child run up into the parent totals and exposes it as usage.delegated', async () => { + const child = agent({ id: 'child', name: 'Child', tools: undefined }); + const parentProvider = mockModel( + [ + { toolCalls: [{ name: 'ask', args: { task: 'sub' } }], usage: { inputTokens: 10, outputTokens: 5 } }, + { text: 'done', usage: { inputTokens: 20, outputTokens: 5 } }, + ], + { defaultModel: 'gpt-4o-mini' } + ); + const childProvider = mockModel( + [ + { text: 'sub done', usage: { inputTokens: 100, outputTokens: 50 } }, + ], + { defaultModel: 'gpt-4o-mini' } + ); + const tools = new ToolRegistry(); + tools.register('ask', createDelegateTool({ agent: child, provider: childProvider })); + + const result = await AgentExecutor.execute({ + agent: agent({ tools: { ask: { tool: 'ask' } } }), + input: 'go', + provider: parentProvider, + toolRegistry: tools, + }); + + expect(result.usage).toMatchObject({ inputTokens: 130, outputTokens: 60, totalTokens: 190, modelCalls: 3 }); + expect(result.usage.delegated).toMatchObject({ + inputTokens: 100, + outputTokens: 50, + totalTokens: 150, + modelCalls: 1, + runs: 1, + estimated: false, + }); + expect(result.usage.delegated?.costUsd).toBeCloseTo((100 * 0.15 + 50 * 0.6) / 1e6, 10); + // the parent's own steps exclude the child + expect(result.stepUsage).toHaveLength(2); + // the model only sees token counts, not the whole breakdown + const toolMessage = result.messages.find((m) => m.role === 'tool'); + expect(JSON.parse(toolMessage!.content)).toEqual({ + text: 'sub done', + usage: { promptTokens: 100, completionTokens: 50, totalTokens: 150 }, + }); + }); + + it('continues totals from the checkpoint when a run resumes', async () => { + const store = checkpointStore(); + const controller = new AbortController(); + const stopper = defineTool({ + name: 'echo', + description: 'echo', + input: z.object({}), + execute: async () => { + controller.abort(); + return 'echoed'; + }, + }); + const tools = new ToolRegistry(); + tools.registerMany([stopper]); + + const first = await AgentExecutor.execute({ + agent: agent(), + input: 'go', + provider: mockModel([{ ...callEcho, usage: { inputTokens: 40, outputTokens: 4 } }], { defaultModel: 'gpt-4o-mini' }), + toolRegistry: tools, + sessionId: 's1', + checkpointStore: store, + signal: controller.signal, + }); + expect(first.finishReason).toBe('aborted'); + expect(first.usage.totalTokens).toBe(44); + + const resumed = await AgentExecutor.execute({ + agent: agent(), + input: 'go', + provider: mockModel([{ text: 'done', usage: { inputTokens: 60, outputTokens: 6 } }], { defaultModel: 'gpt-4o-mini' }), + toolRegistry: registry(), + sessionId: 's1', + checkpointStore: store, + }); + + expect(resumed.usage).toMatchObject({ inputTokens: 100, outputTokens: 10, totalTokens: 110, modelCalls: 2 }); + expect(resumed.usage.byModel['gpt-4o-mini']).toMatchObject({ inputTokens: 100, calls: 2 }); + expect(resumed.stepUsage).toHaveLength(2); + }); + + it('keeps the totals of a run paused for approval when it is resumed', async () => { + const records = new Map(); + const approvalStore: ApprovalStore = { + async save(pending, snapshot) { + records.set(pending.id, { pending, snapshot }); + }, + async resolve(id) { + const record = records.get(id) ?? null; + records.delete(id); + return record; + }, + }; + const gated = defineTool({ name: 'echo', description: 'echo', input: z.object({}), needsApproval: true, execute: async () => 'ok' }); + const tools = new ToolRegistry(); + tools.registerMany([gated]); + const provider = mockModel( + [ + { ...callEcho, usage: { inputTokens: 40, outputTokens: 4 } }, + { text: 'done', usage: { inputTokens: 60, outputTokens: 6 } }, + ], + { defaultModel: 'gpt-4o-mini' } + ); + + const paused = await AgentExecutor.execute({ agent: agent(), input: 'go', provider, toolRegistry: tools, approvalStore }); + expect(paused.finishReason).toBe('awaiting-approval'); + expect(paused.usage.totalTokens).toBe(44); + + const resumed = await resumeAfterApproval({ id: paused.approvalId!, approved: true }, approvalStore, tools, provider); + + expect(resumed.usage).toMatchObject({ inputTokens: 100, outputTokens: 10, modelCalls: 2 }); + }); + + it('starts from the old token counts when the checkpoint predates usage tracking', async () => { + const store = checkpointStore(); + await store.save('old', { + agentId: 'a', + sessionId: 'old', + stepIndex: 1, + messages: [{ role: 'user', content: 'go' }], + toolCalls: [], + usage: { promptTokens: 30, completionTokens: 3, totalTokens: 33 }, + }); + + const result = await AgentExecutor.execute({ + agent: agent(), + input: 'ignored', + provider: mockModel([{ text: 'done', usage: { inputTokens: 10, outputTokens: 1 } }], { defaultModel: 'gpt-4o-mini' }), + sessionId: 'old', + checkpointStore: store, + }); + + expect(result.usage).toMatchObject({ inputTokens: 40, outputTokens: 4, totalTokens: 44, modelCalls: 1 }); + // the earlier tokens have no known model, so no (partial) cost is claimed + expect(result.usage.costUsd).toBeUndefined(); + }); + + it('reports usage on events, onLLMResponse and the chat span from the same numbers', async () => { + const events: ExecutionEvent[] = []; + const seen: unknown[] = []; + const tracer = createTracer(); + const provider = mockModel([{ text: 'done', usage: { inputTokens: 7, outputTokens: 3 } }], { + defaultModel: 'gpt-4o-mini', + }); + + const result = await AgentExecutor.execute({ + agent: agent(), + input: 'go', + provider, + exporter: tracer.exporter, + onEvent: (e) => events.push(e), + onLLMResponse: (_response, _latency, usage) => { + seen.push(usage); + }, + }); + + const finish = events.find((e) => e.type === 'finish'); + expect(finish?.usage).toMatchObject({ inputTokens: 7, outputTokens: 3, modelCalls: 1 }); + expect(events.find((e) => e.type === 'text-complete')?.stepUsage).toEqual(result.stepUsage?.[0]); + expect(seen).toEqual([expect.objectContaining({ model: 'gpt-4o-mini', usage: { inputTokens: 7, outputTokens: 3, totalTokens: 10 } })]); + const chat = tracer.spans().find((s) => s.name.startsWith('chat')); + expect(chat?.attributes['gen_ai.usage.input_tokens']).toBe(7); + expect(chat?.attributes['gen_ai.usage.output_tokens']).toBe(3); + expect(chat?.attributes['loushy.usage.estimated']).toBeUndefined(); + }); + + it('marks the chat span estimated when usage was estimated', async () => { + const tracer = createTracer(); + await AgentExecutor.execute({ agent: agent(), input: 'go', provider: mockModel(['hi']), exporter: tracer.exporter }); + + const chat = tracer.spans().find((s) => s.name.startsWith('chat')); + expect(chat?.attributes['loushy.usage.estimated']).toBe(true); + expect(chat?.attributes['gen_ai.usage.input_tokens']).toBeGreaterThan(0); + }); +}); + +describe('usage on the event stream (LOU-V5)', () => { + it('puts per-step and run usage, cost and the estimated flag on step.done and run.done', async () => { + const provider = mockModel( + [ + { ...callEcho, usage: { inputTokens: 1000, outputTokens: 200 } }, + 'no usage reported', + ], + { defaultModel: 'gpt-4o-mini' } + ); + + const events: AgentEvent[] = []; + for await (const event of AgentExecutor.stream({ agent: agent(), input: 'go', provider, toolRegistry: registry() })) { + events.push(event); + } + + const [first, second] = events.filter((e) => e.type === 'step.done'); + expect(first.usage).toMatchObject({ inputTokens: 1000, outputTokens: 200, promptTokens: 1000, estimated: false }); + expect(first.usage?.costUsd).toBeCloseTo((1000 * 0.15 + 200 * 0.6) / 1e6, 10); + expect(second.usage?.estimated).toBe(true); + const done = events.find((e) => e.type === 'run.done'); + expect(done?.usage).toMatchObject({ modelCalls: 2, estimated: true }); + expect(done?.usage?.inputTokens).toBeGreaterThan(1000); + }); + + it('rolls delegated children up into a streamed run as well', async () => { + const child = agent({ id: 'child', name: 'Child', tools: undefined }); + const parentProvider = mockModel( + [ + { toolCalls: [{ name: 'ask', args: { task: 'sub' } }], usage: { inputTokens: 10, outputTokens: 5 } }, + { text: 'done', usage: { inputTokens: 20, outputTokens: 5 } }, + ], + { defaultModel: 'gpt-4o-mini' } + ); + const childProvider = mockModel([{ text: 'sub', usage: { inputTokens: 100, outputTokens: 50 } }], { defaultModel: 'gpt-4o-mini' }); + const tools = new ToolRegistry(); + tools.register('ask', createDelegateTool({ agent: child, provider: childProvider })); + + const events: AgentEvent[] = []; + for await (const event of AgentExecutor.stream({ + agent: agent({ tools: { ask: { tool: 'ask' } } }), + input: 'go', + provider: parentProvider, + toolRegistry: tools, + })) { + events.push(event); + } + + expect(events.find((e) => e.type === 'run.done')?.usage).toMatchObject({ inputTokens: 130, outputTokens: 60, modelCalls: 3 }); + }); +}); + +describe('normalizeUsage', () => { + it('maps provider names, keeps cache/reasoning tokens when present, and rejects NaN or missing usage', () => { + expect(normalizeUsage({ promptTokens: 10, completionTokens: 5, totalTokens: 15, cachedInputTokens: 4, reasoningTokens: 2 })).toEqual({ + inputTokens: 10, + outputTokens: 5, + totalTokens: 15, + cachedInputTokens: 4, + reasoningTokens: 2, + }); + expect(normalizeUsage({ promptTokens: NaN, completionTokens: NaN, totalTokens: NaN })).toBeUndefined(); + expect(normalizeUsage(undefined)).toBeUndefined(); + expect(normalizeUsage({ promptTokens: 1, completionTokens: 2, totalTokens: NaN })?.totalTokens).toBe(3); + }); +}); + +describe('formatUsage', () => { + it('formats singular calls, large costs and estimates', () => { + const base = { inputTokens: 1234, outputTokens: 567, costUsd: 12.345, modelCalls: 1, estimated: false }; + expect(formatUsage(base)).toBe('1,234 in / 567 out tokens · $12.35 (1 model call)'); + expect(formatUsage({ ...base, costUsd: undefined, estimated: true })).toBe('~1,234 in / 567 out tokens (1 model call)'); + }); +}); + +function createTracer() { + const ended: Span[] = []; + const exporter: TraceExporter = { onSpanStart: () => undefined, onSpanEnd: (span) => void ended.push({ ...span }) }; + return { exporter, spans: () => ended }; +} + +function checkpointStore(): CheckpointStore { + const checkpoints = new Map(); + return { + async save(sessionId, checkpoint) { + checkpoints.set(sessionId, JSON.parse(JSON.stringify(checkpoint))); + }, + async load(sessionId) { + return checkpoints.get(sessionId) ?? null; + }, + async delete(sessionId) { + checkpoints.delete(sessionId); + }, + }; +} diff --git a/src/models/index.ts b/src/models/index.ts index d58a59a91..5914844ee 100644 --- a/src/models/index.ts +++ b/src/models/index.ts @@ -5,3 +5,4 @@ export * from './estimateTokens'; export * from './registry'; +export * from './usage'; diff --git a/src/models/usage.ts b/src/models/usage.ts new file mode 100644 index 000000000..c3d17ba56 --- /dev/null +++ b/src/models/usage.ts @@ -0,0 +1,146 @@ +/** + * Usage and cost accounting types (LOU-V5). + * + * `Usage` is the one normalized shape for "tokens spent"; `RunUsage` is the + * running total an `ExecutionResult` carries. Providers report the 'ai' SDK's + * `promptTokens`/`completionTokens` naming (`ProviderUsage`); the executor + * converts it with {@link normalizeUsage}. + */ + +import type { ProviderUsage } from '../providers/llm'; + +/** Normalized token usage of one model call. */ +export interface Usage { + /** Tokens sent to the model (the prompt). */ + inputTokens: number; + /** Tokens the model generated. */ + outputTokens: number; + /** `inputTokens + outputTokens`, or the provider's own total. */ + totalTokens: number; + /** Input tokens served from the provider's prompt cache. Only set when the provider reports it. */ + cachedInputTokens?: number; + /** Output tokens spent on hidden reasoning. Only set when the provider reports it. */ + reasoningTokens?: number; +} + +/** Per-model slice of a {@link RunUsage}. */ +export interface ModelUsage { + inputTokens: number; + outputTokens: number; + /** USD for these tokens; `undefined` when the model's price is unknown (see `registerModel`). */ + costUsd?: number; + /** Number of model calls made with this model. */ + calls: number; +} + +/** Usage that child agents (delegated through `createDelegateTool`) spent; already included in the run totals. */ +export interface DelegatedUsage { + inputTokens: number; + outputTokens: number; + totalTokens: number; + /** `undefined` when any contributing model has unknown pricing. */ + costUsd: number | undefined; + modelCalls: number; + estimated: boolean; + /** Number of child runs. */ + runs: number; +} + +/** + * Running usage of an agent run, on `ExecutionResult.usage`. Totals include + * every step and every delegated child run. + * + * @example + * ```ts + * const result = await AgentExecutor.execute({ agent, input, provider }); + * console.log(formatUsage(result.usage)); // 1,234 in / 567 out tokens · $0.0042 (2 model calls) + * ``` + */ +export interface RunUsage { + inputTokens: number; + outputTokens: number; + totalTokens: number; + /** Sum of `cachedInputTokens` over the calls that reported it; absent when none did. */ + cachedInputTokens?: number; + /** Sum of `reasoningTokens` over the calls that reported it; absent when none did. */ + reasoningTokens?: number; + /** + * Estimated USD for the whole run. `undefined` (never a partial sum) when + * any contributing model has unknown pricing; the per-model figures in + * `byModel` still show which models are priced. + */ + costUsd: number | undefined; + /** Model calls made, including those of delegated children. */ + modelCalls: number; + /** `true` when at least one call reported no usage and its tokens were estimated with `estimateTokens`. */ + estimated: boolean; + /** Breakdown keyed by the model id each call was made with. */ + byModel: Record; + /** Present when the run delegated to child agents. */ + delegated?: DelegatedUsage; + /** @deprecated Alias of `inputTokens`, kept for code written before LOU-V5. */ + promptTokens: number; + /** @deprecated Alias of `outputTokens`, kept for code written before LOU-V5. */ + completionTokens: number; +} + +/** Usage of one model call within a run (`ExecutionResult.stepUsage`). */ +export interface StepUsage { + /** 1-based loop step the call belongs to. */ + step: number; + /** The model the call was made with. */ + model: string; + usage: Usage; + /** `true` when the provider reported nothing and `usage` is an `estimateTokens` estimate. */ + estimated: boolean; + /** USD for this call; `undefined` when the model's price is unknown. */ + costUsd?: number; +} + +/** What one model call spent: a {@link StepUsage} before it is assigned a loop step. */ +export type CallUsage = Omit; + +function isCount(value: unknown): value is number { + return typeof value === 'number' && Number.isFinite(value) && value >= 0; +} + +/** + * Convert what a provider reported into {@link Usage}. Returns `undefined` + * (never zeros) when the provider reported nothing usable, e.g. the 'ai' SDK's + * `NaN` for a backend that omits token counts. + * + * @example + * normalizeUsage({ promptTokens: 10, completionTokens: 5, totalTokens: 15 }); + * // { inputTokens: 10, outputTokens: 5, totalTokens: 15 } + */ +export function normalizeUsage(raw: ProviderUsage | undefined): Usage | undefined { + if (!raw || !isCount(raw.promptTokens) || !isCount(raw.completionTokens)) return undefined; + return { + inputTokens: raw.promptTokens, + outputTokens: raw.completionTokens, + totalTokens: isCount(raw.totalTokens) ? raw.totalTokens : raw.promptTokens + raw.completionTokens, + ...(isCount(raw.cachedInputTokens) ? { cachedInputTokens: raw.cachedInputTokens } : {}), + ...(isCount(raw.reasoningTokens) ? { reasoningTokens: raw.reasoningTokens } : {}), + }; +} + +function formatCost(costUsd: number): string { + return `$${costUsd.toFixed(costUsd >= 1 ? 2 : 4)}`; +} + +/** + * One-line summary of a run's usage for logs and CLIs. A leading `~` marks + * estimated token counts; the cost part is left out when it is unknown. + * + * @example + * formatUsage(result.usage); // "1,234 in / 567 out tokens · $0.0042 (2 model calls)" + */ +export function formatUsage( + usage: Pick +): string { + const n = (value: number): string => value.toLocaleString('en-US'); + const tokens = `${usage.estimated ? '~' : ''}${n(usage.inputTokens)} in / ${n(usage.outputTokens)} out tokens`; + const cost = usage.costUsd === undefined ? '' : ` · ${formatCost(usage.costUsd)}`; + const calls = `${usage.modelCalls} model call${usage.modelCalls === 1 ? '' : 's'}`; + return `${tokens}${cost} (${calls})`; +} diff --git a/src/providers/aiSdkProvider.test.ts b/src/providers/aiSdkProvider.test.ts index d7267929a..c4e210ed9 100644 --- a/src/providers/aiSdkProvider.test.ts +++ b/src/providers/aiSdkProvider.test.ts @@ -57,6 +57,25 @@ describe('AiSdkProvider', () => { expect(settings.messages).toEqual([{ role: 'user', content: 'hi' }]); }); + it('LOU-V5: reports usage as-is, reads cache/reasoning tokens from provider metadata, and no usage when counts are NaN', async () => { + const provider = new OpenAIProvider({ name: 'openai', apiKey: 'k' }); + + generateTextMock.mockResolvedValue({ + ...textResult('stop'), + providerMetadata: { openai: { cachedPromptTokens: 1, reasoningTokens: 'n/a' } }, + }); + expect((await provider.generate({ model: '', messages: [] })).usage).toEqual({ + ...usage, + cachedInputTokens: 1, + }); + + generateTextMock.mockResolvedValue({ + ...textResult('stop'), + usage: { promptTokens: NaN, completionTokens: NaN, totalTokens: NaN }, + }); + expect((await provider.generate({ model: '', messages: [] })).usage).toBeUndefined(); + }); + it('falls back to the provider default model when neither call nor config names one', async () => { generateTextMock.mockResolvedValue(textResult('stop')); const provider = new OpenAIProvider({ name: 'openai', apiKey: 'k' }); diff --git a/src/providers/aiSdkProvider.ts b/src/providers/aiSdkProvider.ts index 53c7b7cdf..2d05d5da1 100644 --- a/src/providers/aiSdkProvider.ts +++ b/src/providers/aiSdkProvider.ts @@ -36,6 +36,7 @@ import { Message, ToolCall, ToolDefinition, + ProviderUsage, } from './llm'; /** Config fields shared by every 'ai'-SDK-backed provider. */ @@ -43,7 +44,7 @@ export interface AiSdkProviderConfig extends LLMProviderConfig { defaultModel?: string; } -type TokenUsage = GenerateResult['usage']; +type TokenUsage = ProviderUsage; /** The subset of an 'ai' SDK tool call this module reads. */ interface AiSdkToolCall { @@ -181,6 +182,43 @@ function toUsage(usage: TokenUsage): TokenUsage { }; } +/** A finite, non-negative count from untyped provider metadata, else `undefined`. */ +function metadataCount(value: unknown): number | undefined { + return typeof value === 'number' && Number.isFinite(value) && value >= 0 ? value : undefined; +} + +/** The first documented cache/reasoning token field a provider's metadata carries. */ +function pickCount(metadata: Record | undefined, keys: string[]): number | undefined { + for (const provider of Object.values(metadata ?? {})) { + const fields = provider as Record | undefined; + for (const key of keys) { + const count = metadataCount(fields?.[key]); + if (count !== undefined) return count; + } + } + return undefined; +} + +/** + * The call's usage, or `undefined` when the backend reported none (the 'ai' + * SDK yields `NaN` counts then; never zeros). Cache and reasoning tokens are + * read from `providerMetadata` when the provider package documents them + * (OpenAI `cachedPromptTokens`/`reasoningTokens`, Anthropic `cacheReadInputTokens`). + */ +function toGenerateUsage( + usage: TokenUsage, + metadata: Record | undefined +): ProviderUsage | undefined { + if (!Number.isFinite(usage.promptTokens) || !Number.isFinite(usage.completionTokens)) return undefined; + const cachedInputTokens = pickCount(metadata, ['cachedPromptTokens', 'cacheReadInputTokens']); + const reasoningTokens = pickCount(metadata, ['reasoningTokens']); + return { + ...toUsage(usage), + ...(cachedInputTokens !== undefined ? { cachedInputTokens } : {}), + ...(reasoningTokens !== undefined ? { reasoningTokens } : {}), + }; +} + /** Yield every text delta, then a single finish chunk carrying usage stats. */ async function* toFullStream(result: AiSdkStreamResult): AsyncGenerator { for await (const delta of result.textStream) { @@ -202,7 +240,7 @@ async function* toFullStream(result: AiSdkStreamResult): AsyncGenerator result.text)(), - usage: (async () => toUsage(await result.usage))(), + usage: (async () => toGenerateUsage(await result.usage, undefined))(), finishReason: (async () => result.finishReason)(), toolCalls: (async () => convertToolCalls(await result.toolCalls))(), }; @@ -281,7 +319,7 @@ export abstract class AiSdkProvider impleme return { text: result.text, finishReason: mapFinishReason(result.finishReason), - usage: toUsage(result.usage), + usage: toGenerateUsage(result.usage, result.providerMetadata), toolCalls: result.toolCalls && convertToolCalls(result.toolCalls), rawResponse: result, }; diff --git a/src/providers/llm.ts b/src/providers/llm.ts index ef8892f81..8822b8739 100644 --- a/src/providers/llm.ts +++ b/src/providers/llm.ts @@ -95,17 +95,32 @@ export interface GenerateOptions { signal?: AbortSignal; } +/** + * Token usage as a provider reports it (the 'ai' SDK's naming). AgentExecutor + * converts it to the normalized `Usage` (`inputTokens`/`outputTokens`). + */ +export interface ProviderUsage { + promptTokens: number; + completionTokens: number; + totalTokens: number; + /** Prompt tokens served from the provider's cache, when the provider reports it. */ + cachedInputTokens?: number; + /** Tokens spent on hidden reasoning, when the provider reports it. */ + reasoningTokens?: number; +} + /** * Generation result */ export interface GenerateResult { text: string; finishReason: 'stop' | 'length' | 'tool_calls' | 'content_filter' | 'error'; - usage: { - promptTokens: number; - completionTokens: number; - totalTokens: number; - }; + /** + * Token usage of this call. Leave it `undefined` when the backend reports + * nothing - never fill in zeros: AgentExecutor then estimates the tokens + * and flags the run's usage as `estimated`. + */ + usage?: ProviderUsage; toolCalls?: ToolCall[]; rawResponse?: any; } @@ -132,11 +147,8 @@ export interface StreamChunk { result: any; }; finishReason?: string; - usage?: { - promptTokens: number; - completionTokens: number; - totalTokens: number; - }; + /** Usage of the call, on the `finish` chunk; omit when the backend reports none (see GenerateResult.usage). */ + usage?: ProviderUsage; error?: Error; } @@ -147,11 +159,8 @@ export interface StreamResult { textStream: AsyncIterable; fullStream: AsyncIterable; text: Promise; - usage: Promise<{ - promptTokens: number; - completionTokens: number; - totalTokens: number; - }>; + /** Resolves to `undefined` when the backend reports no usage (see GenerateResult.usage). */ + usage: Promise; finishReason: Promise; toolCalls: Promise; } diff --git a/src/testing/cassette.ts b/src/testing/cassette.ts index 0eb0d2cd3..d623e4340 100644 --- a/src/testing/cassette.ts +++ b/src/testing/cassette.ts @@ -55,7 +55,8 @@ const requestSchema = z.object({ const responseSchema = z.object({ text: z.string(), finishReason: z.string(), - usage: usageSchema, + /** Absent when the provider reported no usage for the call. */ + usage: usageSchema.optional(), toolCalls: z.array(toolCallSchema).optional(), /** Present for `stream()` entries: the chunks and the delay before each. */ chunks: z.array(z.object({ delayMs: z.number(), chunk: chunkSchema })).optional(), diff --git a/src/testing/mockModel.test.ts b/src/testing/mockModel.test.ts index c9e8e3bfa..f7b6f24d0 100644 --- a/src/testing/mockModel.test.ts +++ b/src/testing/mockModel.test.ts @@ -7,7 +7,7 @@ import type { GenerateOptions, StreamChunk } from '../providers/llm'; const request = (content = 'hi'): GenerateOptions => ({ messages: [{ role: 'user', content }] }); describe('mockModel turns', () => { - it('treats a bare string as { text } and reports zero usage by default', async () => { + it('treats a bare string as { text } and reports no usage by default', async () => { const model = mockModel(['Hello!']); const result = await model.generate(request()); @@ -15,8 +15,8 @@ describe('mockModel turns', () => { expect(result).toEqual({ text: 'Hello!', finishReason: 'stop', - usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 }, }); + expect(result.usage).toBeUndefined(); }); it('emits tool calls with stable generated ids and JSON-encoded args', async () => { diff --git a/src/testing/mockModel.ts b/src/testing/mockModel.ts index 36ff6b93b..951cdcb19 100644 --- a/src/testing/mockModel.ts +++ b/src/testing/mockModel.ts @@ -61,7 +61,10 @@ export interface MockTurnObject { toolCalls?: readonly MockToolCall[]; /** Make `generate()` / `stream()` reject with this error instead of replying. */ error?: Error; - /** Token usage to report. Defaults to zero. */ + /** + * Token usage to report. When omitted the model reports no usage at all, like a + * backend that omits token counts (the executor then estimates and flags it). + */ usage?: { inputTokens: number; outputTokens: number }; /** Finish reason to report. Defaults to `'tool_calls'` when there are tool calls, else `'stop'`. */ finishReason?: GenerateResult['finishReason']; @@ -280,16 +283,14 @@ class ScriptedMockModel implements MockModel { private toResolved(turn: MockTurnObject): ResolvedTurn { const toolCalls = (turn.toolCalls ?? []).map((call) => this.toToolCall(call)); - const inputTokens = turn.usage?.inputTokens ?? 0; - const outputTokens = turn.usage?.outputTokens ?? 0; return { text: turn.text ?? '', toolCalls, finishReason: turn.finishReason ?? (toolCalls.length > 0 ? 'tool_calls' : 'stop'), - usage: { - promptTokens: inputTokens, - completionTokens: outputTokens, - totalTokens: inputTokens + outputTokens, + usage: turn.usage && { + promptTokens: turn.usage.inputTokens, + completionTokens: turn.usage.outputTokens, + totalTokens: turn.usage.inputTokens + turn.usage.outputTokens, }, }; } diff --git a/src/testing/recordReplay.ts b/src/testing/recordReplay.ts index 3f7aff979..70b1c76a7 100644 --- a/src/testing/recordReplay.ts +++ b/src/testing/recordReplay.ts @@ -9,6 +9,7 @@ import type { GenerateOptions, GenerateResult, LLMProvider, + ProviderUsage, StreamChunk, StreamResult, } from '../providers/llm'; @@ -135,6 +136,13 @@ function clone(value: T): T { return JSON.parse(JSON.stringify(value)) as T; } +/** Only the token counts are stored; a call that reported no usage records none. */ +function recordedUsage(usage: ProviderUsage | undefined): { usage?: ProviderUsage } { + if (!usage) return {}; + const { promptTokens, completionTokens, totalTokens } = usage; + return { usage: { promptTokens, completionTokens, totalTokens } }; +} + function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } @@ -157,7 +165,7 @@ function replayStream(response: CassetteResponse, replayTiming: boolean): Stream fullStream: fullStream(), textStream: textStream(), text: Promise.resolve(response.text), - usage: Promise.resolve(clone(response.usage)), + usage: Promise.resolve(response.usage && clone(response.usage)), finishReason: Promise.resolve(response.finishReason), toolCalls: Promise.resolve(clone(response.toolCalls ?? [])), }; @@ -167,7 +175,7 @@ function toGenerateResult(response: CassetteResponse): GenerateResult { return { text: response.text, finishReason: response.finishReason as GenerateResult['finishReason'], - usage: clone(response.usage), + ...(response.usage ? { usage: clone(response.usage) } : {}), ...(response.toolCalls?.length ? { toolCalls: clone(response.toolCalls) } : {}), }; } @@ -257,11 +265,7 @@ class Recorder implements RecordReplayProvider { return this.sanitizer.redact({ text: result.text, finishReason: result.finishReason, - usage: { - promptTokens: result.usage.promptTokens, - completionTokens: result.usage.completionTokens, - totalTokens: result.usage.totalTokens, - }, + ...recordedUsage(result.usage), ...(result.toolCalls?.length ? { toolCalls: result.toolCalls } : {}), }); } @@ -283,11 +287,7 @@ class Recorder implements RecordReplayProvider { return this.sanitizer.redact({ text, finishReason, - usage: { - promptTokens: usage.promptTokens, - completionTokens: usage.completionTokens, - totalTokens: usage.totalTokens, - }, + ...recordedUsage(usage), ...(toolCalls.length ? { toolCalls } : {}), chunks, });