From b0867ccb8a46b571b910df5ae2409e4981bf2df3 Mon Sep 17 00:00:00 2001 From: can1357 Date: Fri, 15 May 2026 17:47:42 +0200 Subject: [PATCH] fix(ai/providers): recomputed mock totals and scoped tool-call ID counter to instance - mergeUsage now recomputes totalTokens (and cost.total when cost components are supplied) when omitted by a Partial, so mock-backed telemetry/run-summary tests don't under-count. - Tool-call ID counter moved into MockState; reset() restores it, so two createMockModel() instances with identical scripted tool calls produce identical IDs regardless of test ordering. --- packages/ai/CHANGELOG.md | 2 + packages/ai/src/providers/mock.ts | 34 +++++++++----- packages/ai/test/mock-provider.test.ts | 62 ++++++++++++++++++++++++++ 3 files changed, 87 insertions(+), 11 deletions(-) diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 576192369..e42cde873 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -45,6 +45,8 @@ - Fixed Zod-emitted wire schemas dropping null-valued unknown root fields before `preserveUnknownRootFields` could snapshot them, so callers like `task.simple` no longer lose a `schema: null` argument and downstream rejection paths fire as intended. - Fixed `buildOpenAiNativeHistory` to serialize freeform tool calls (e.g. `apply_patch`) as `custom_tool_call` / `custom_tool_call_output` items matching the live Responses provider, instead of demoting them to `function_call` and breaking remote-compaction replay - Threaded an `AbortSignal` through `requestOpenAiRemoteCompaction` and `requestRemoteCompaction` so canceled compaction requests no longer keep the underlying `fetch` alive +- Fixed mock provider partial `Usage` to recompute `totalTokens` (and `cost.total` when cost components are supplied) when omitted, instead of reporting 0 +- Fixed mock provider auto-generated tool-call IDs to use a per-instance counter (now reset by `reset()`), so test order no longer affects IDs across `createMockModel()` instances ## [15.0.2] - 2026-05-15 ### Fixed diff --git a/packages/ai/src/providers/mock.ts b/packages/ai/src/providers/mock.ts index 96128e770..5dae1b980 100644 --- a/packages/ai/src/providers/mock.ts +++ b/packages/ai/src/providers/mock.ts @@ -81,8 +81,8 @@ export interface MockResponse { content?: ReadonlyArray; /** Stop reason. Defaults to `"toolUse"` when content has tool calls, else `"stop"`. */ stopReason?: StopReason; - /** Usage stats. Missing fields default to 0. */ - usage?: Partial; + /** Usage stats. Missing fields default to 0; missing `cost.total` is recomputed from components. */ + usage?: Partial> & { cost?: Partial }; /** Pre-set responseId. */ responseId?: string; /** If set, the stream emits a terminal error event instead of completing. */ @@ -159,6 +159,7 @@ interface MockState { readonly extras: MockHandler[]; fallback?: MockHandler; readonly calls: MockCall[]; + toolCallCounter: number; } const STATE_BY_MODEL = new WeakMap, MockState>(); @@ -196,6 +197,7 @@ export function createMockModel(options: MockModelOptions = {}): MockModelHandle extras: [], fallback: options.handler, calls: [], + toolCallCounter: 0, }; STATE_BY_MODEL.set(model, state); @@ -209,6 +211,7 @@ export function createMockModel(options: MockModelOptions = {}): MockModelHandle reset() { state.extras.length = 0; state.calls.length = 0; + state.toolCallCounter = 0; }, }; } @@ -334,7 +337,7 @@ async function runMock( stream.push({ type: "start", partial }); for (const input of response.content ?? []) { - const block = normalizeContent(input); + const block = normalizeContent(input, state); blocks.push(block); const contentIndex = blocks.length - 1; @@ -372,14 +375,14 @@ async function runMock( stream.push({ type: "done", reason: reason as "stop" | "length" | "toolUse", message: partial }); } -function normalizeContent(input: MockContent): TextContent | ThinkingContent | ToolCall { +function normalizeContent(input: MockContent, state: MockState): TextContent | ThinkingContent | ToolCall { if (typeof input === "string") { return { type: "text", text: input }; } if (input.type === "toolCall") { return { type: "toolCall", - id: input.id ?? generateToolCallId(), + id: input.id ?? generateToolCallId(state), name: input.name, arguments: typeof input.arguments === "string" ? input.arguments : { ...input.arguments }, } as ToolCall; @@ -398,13 +401,23 @@ function emptyUsage(): Usage { } as Usage; } -function mergeUsage(partial?: Partial): Usage { +function mergeUsage(partial?: Partial> & { cost?: Partial }): Usage { const base = emptyUsage(); if (!partial) return base; const merged = { ...base, ...partial } as Usage; - if (partial.cost) { + const costProvided = partial.cost !== undefined; + if (costProvided) { merged.cost = { ...base.cost, ...partial.cost } as Usage["cost"]; } + // Recompute totalTokens when not explicitly provided (canonical formula matches types.ts: + // input + output + cacheRead + cacheWrite). + if (partial.totalTokens === undefined) { + merged.totalTokens = merged.input + merged.output + merged.cacheRead + merged.cacheWrite; + } + // Recompute cost.total when cost components were supplied without an explicit total. + if (costProvided && partial.cost?.total === undefined) { + merged.cost.total = merged.cost.input + merged.cost.output + merged.cost.cacheRead + merged.cost.cacheWrite; + } return merged; } @@ -450,8 +463,7 @@ function sleep(ms: number, signal?: AbortSignal): Promise { return promise; } -let toolCallCounter = 0; -function generateToolCallId(): string { - toolCallCounter += 1; - return `mock-tc-${toolCallCounter}`; +function generateToolCallId(state: MockState): string { + state.toolCallCounter += 1; + return `mock-tc-${state.toolCallCounter}`; } diff --git a/packages/ai/test/mock-provider.test.ts b/packages/ai/test/mock-provider.test.ts index a051e38a9..f81258b90 100644 --- a/packages/ai/test/mock-provider.test.ts +++ b/packages/ai/test/mock-provider.test.ts @@ -190,6 +190,68 @@ describe("AssistantMessage shape", () => { expect(result.usage.input).toBe(5); expect(result.usage.output).toBe(2); }); + + test("partial usage without totalTokens recomputes the total from components", async () => { + const mock = createMockModel({ + responses: [{ content: ["x"], usage: { input: 5, output: 2 } }], + }); + const result = await mock.stream(mock.model, emptyContext()).result(); + expect(result.usage.totalTokens).toBe(7); + }); + + test("partial usage with explicit totalTokens is respected", async () => { + const mock = createMockModel({ + responses: [{ content: ["x"], usage: { input: 5, output: 2, totalTokens: 999 } }], + }); + const result = await mock.stream(mock.model, emptyContext()).result(); + expect(result.usage.totalTokens).toBe(999); + }); + + test("partial cost components recompute cost.total when total is omitted", async () => { + const mock = createMockModel({ + responses: [ + { + content: ["x"], + usage: { input: 5, output: 2, cost: { input: 0.5, output: 0.25, cacheRead: 0, cacheWrite: 0 } }, + }, + ], + }); + const result = await mock.stream(mock.model, emptyContext()).result(); + expect(result.usage.cost.total).toBeCloseTo(0.75, 10); + }); + + test("tool-call ID counter is scoped per mock instance and resets with reset()", async () => { + const makeMock = () => + createMockModel({ + responses: [ + { content: [{ type: "toolCall", name: "read", arguments: { path: "/x" } }] }, + { content: [{ type: "toolCall", name: "read", arguments: { path: "/y" } }] }, + ], + }); + + const a = makeMock(); + const b = makeMock(); + + const a1 = await a.stream(a.model, emptyContext()).result(); + const a2 = await a.stream(a.model, emptyContext()).result(); + const b1 = await b.stream(b.model, emptyContext()).result(); + const b2 = await b.stream(b.model, emptyContext()).result(); + + const idOf = (m: AssistantMessage): string => { + const tc = m.content.find(c => c.type === "toolCall") as ToolCall; + return tc.id; + }; + + expect(idOf(a1)).toBe("mock-tc-1"); + expect(idOf(a2)).toBe("mock-tc-2"); + expect(idOf(b1)).toBe("mock-tc-1"); + expect(idOf(b2)).toBe("mock-tc-2"); + + a.reset(); + a.push({ content: [{ type: "toolCall", name: "read", arguments: { path: "/z" } }] }); + const a3 = await a.stream(a.model, emptyContext()).result(); + expect(idOf(a3)).toBe("mock-tc-1"); + }); }); describe("mock provider — async-iterable response sources", () => {