From 2628ca822fb7984b2ac7864cfc7f7514ae59b447 Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 3 Nov 2026 11:14:48 +0000 Subject: [PATCH 1/8] feat(sdk): count AI SDK v5, v6 and v7 messages and tool schemas Read tool input schemas through the AI SDK's asSchema() instead of walking zod 4 internals, so zod 3, zod 4, jsonSchema() and (v6+) Standard Schema tools all count. A schema count() cannot read synchronously counts 0 and is listed in tools.unresolved. Count the payload a provider sees for v7 file parts, reasoning, reasoning-file and every tool-result output shape, with one countMediaTokens function for all images and files. Widen the peer range to ai 5, 6 and 7 and develop against ai 7 with an explicit zod dev dependency. Change-Id: I735b100758607503ccda61b172690e8d3946460c Signed-off-by: Thomas Kosiewski --- bun.lock | 27 +- package.json | 7 +- src/sdk.compat.test.ts | 919 +++++++++++++++++++++++++++++++++++++++++ src/sdk.ts | 498 ++++++++++++++++------ 4 files changed, 1312 insertions(+), 139 deletions(-) create mode 100644 src/sdk.compat.test.ts diff --git a/bun.lock b/bun.lock index b523c57..7ebce12 100644 --- a/bun.lock +++ b/bun.lock @@ -6,16 +6,17 @@ "name": "ai-tokenizer", "devDependencies": { "@types/bun": "latest", - "ai": "^5.0.105", + "ai": "^7.0.162", "base64-js": "^1.5.1", "gpt-tokenizer": "^3.0.1", "mitata": "^1.0.34", "prettier": "^3.6.2", "tiktoken": "^1.0.22", "tsdown": "^0.15.6", + "zod": "^4.6.7", }, "peerDependencies": { - "ai": "^5.0.61", + "ai": "^5.0.0 || ^6.0.0 || ^7.0.0", }, "optionalPeers": [ "ai", @@ -23,11 +24,11 @@ }, }, "packages": { - "@ai-sdk/gateway": ["@ai-sdk/gateway@2.0.17", "", { "dependencies": { "@ai-sdk/provider": "2.0.0", "@ai-sdk/provider-utils": "3.0.18", "@vercel/oidc": "3.0.5" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-oVAG6q72KsjKlrYdLhWjRO7rcqAR8CjokAbYuyVZoCO4Uh2PH/VzZoxZav71w2ipwlXhHCNaInGYWNs889MMDA=="], + "@ai-sdk/gateway": ["@ai-sdk/gateway@4.0.135", "", { "dependencies": { "@ai-sdk/provider": "4.0.35", "@ai-sdk/provider-utils": "5.0.74", "@vercel/oidc": "3.2.0" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-gnMpZP42WIhy24fpnbbtV7SBqEhYY8ZSMsfgoNRh5r41D1HwnJTm1k7d4uYSIpiaEb5J8IPTHVSzwgWmGoSFEA=="], - "@ai-sdk/provider": ["@ai-sdk/provider@2.0.0", "", { "dependencies": { "json-schema": "^0.4.0" } }, "sha512-6o7Y2SeO9vFKB8lArHXehNuusnpddKPk7xqL7T2/b+OvXMRIXUO1rR4wcv1hAFUAT9avGZshty3Wlua/XA7TvA=="], + "@ai-sdk/provider": ["@ai-sdk/provider@4.0.35", "", { "dependencies": { "json-schema": "^0.4.0" } }, "sha512-yBz6fmLS8MSCIzy7NiHEUAq56ZLD4ZA1GR6bpXBxgbwUacP1La7LuwAr/xy9jSzthem4F8Xu2OrhvU/JrE9fYA=="], - "@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@3.0.18", "", { "dependencies": { "@ai-sdk/provider": "2.0.0", "@standard-schema/spec": "^1.0.0", "eventsource-parser": "^3.0.6" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-ypv1xXMsgGcNKUP+hglKqtdDuMg68nWHucPPAhIENrbFAI+xCHiqPVN8Zllxyv1TNZwGWUghPxJXU+Mqps0YRQ=="], + "@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@5.0.74", "", { "dependencies": { "@ai-sdk/provider": "4.0.35", "@standard-schema/spec": "^1.1.0", "@workflow/serde": "4.1.0", "eventsource-parser": "^3.0.8", "undici": "^7.29.1" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-42R/bzs520oy3PMj8FI1KGegjLEjaUsNcaEltb6yxN0KnmhihWvlgMb19pEF8CKKqv6Bq6+qMVji8mtyrMubSQ=="], "@babel/generator": ["@babel/generator@7.28.3", "", { "dependencies": { "@babel/parser": "^7.28.3", "@babel/types": "^7.28.2", "@jridgewell/gen-mapping": "^0.3.12", "@jridgewell/trace-mapping": "^0.3.28", "jsesc": "^3.0.2" } }, "sha512-3lSpxGgvnmZznmBkCRnVREPUFJv2wrv9iAoFDvADJc0ypmdOxdUtcLeBgBJ6zE0PMeTKnxeQzyk0xTBq4Ep7zw=="], @@ -55,8 +56,6 @@ "@napi-rs/wasm-runtime": ["@napi-rs/wasm-runtime@1.0.6", "", { "dependencies": { "@emnapi/core": "^1.5.0", "@emnapi/runtime": "^1.5.0", "@tybys/wasm-util": "^0.10.1" } }, "sha512-DXj75ewm11LIWUk198QSKUTxjyRjsBwk09MuMk5DGK+GDUtyPhhEHOGP/Xwwj3DjQXXkivoBirmOnKrLfc0+9g=="], - "@opentelemetry/api": ["@opentelemetry/api@1.9.0", "", {}, "sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg=="], - "@oxc-project/types": ["@oxc-project/types@0.94.0", "", {}, "sha512-+UgQT/4o59cZfH6Cp7G0hwmqEQ0wE+AdIwhikdwnhWI9Dp8CgSY081+Q3O67/wq3VJu8mgUEB93J9EHHn70fOw=="], "@quansync/fs": ["@quansync/fs@0.1.5", "", { "dependencies": { "quansync": "^0.2.11" } }, "sha512-lNS9hL2aS2NZgNW7BBj+6EBl4rOf8l+tQ0eRY6JWCI8jI2kc53gSoqbjojU0OnAWhzoXiOjFyGsHcDGePB3lhA=="], @@ -91,7 +90,7 @@ "@rolldown/pluginutils": ["@rolldown/pluginutils@1.0.0-beta.42", "", {}, "sha512-N7pQzk9CyE7q0bBN/q0J8s6Db279r5kUZc6d7/wWRe9/zXqC52HQovVyu6iXPIDY4BEzzgbVLhVFXrOuGJ22ZQ=="], - "@standard-schema/spec": ["@standard-schema/spec@1.0.0", "", {}, "sha512-m2bOd0f2RT9k8QJx1JN85cZYyH1RqFBdlwtkSlf4tBDYLCiiZnv1fIIwacK6cqwXavOydf0NPToMQgpKq+dVlA=="], + "@standard-schema/spec": ["@standard-schema/spec@1.1.0", "", {}, "sha512-l2aFy5jALhniG5HgqrD6jXLi/rUWrKvqN/qJx6yoJsgKhblVd+iqqU4RCXavm/jPityDo5TCvKMnpjKnOriy0w=="], "@tybys/wasm-util": ["@tybys/wasm-util@0.10.1", "", { "dependencies": { "tslib": "^2.4.0" } }, "sha512-9tTaPJLSiejZKx+Bmog4uSubteqTvFrVrURwkmHixBo0G4seD0zUxp98E1DzUBJxLQ3NPwXrGKDiVjwx/DpPsg=="], @@ -101,9 +100,11 @@ "@types/react": ["@types/react@19.2.2", "", { "dependencies": { "csstype": "^3.0.2" } }, "sha512-6mDvHUFSjyT2B2yeNx2nUgMxh9LtOWvkhIU3uePn2I2oyNymUAX1NIsdgviM4CH+JSrp2D2hsMvJOkxY+0wNRA=="], - "@vercel/oidc": ["@vercel/oidc@3.0.5", "", {}, "sha512-fnYhv671l+eTTp48gB4zEsTW/YtRgRPnkI2nT7x6qw5rkI1Lq2hTmQIpHPgyThI0znLK+vX2n9XxKdXZ7BUbbw=="], + "@vercel/oidc": ["@vercel/oidc@3.2.0", "", {}, "sha512-UycprH3T6n3jH0k44NHMa7pnFHGu/N05MjojYr+Mc6I7obkoLIJujSWwin1pCvdy/eOxrI/l3uDLQsmcrOb4ug=="], + + "@workflow/serde": ["@workflow/serde@4.1.0", "", {}, "sha512-pav4F2BoirECWR7Nf1TKt+2eETcBj7jj4cBefQ8VXQCA6NPkaKeLfj/zMgi+3zYV5ZIBT4GuUiphsj0/b9hPQQ=="], - "ai": ["ai@5.0.105", "", { "dependencies": { "@ai-sdk/gateway": "2.0.17", "@ai-sdk/provider": "2.0.0", "@ai-sdk/provider-utils": "3.0.18", "@opentelemetry/api": "1.9.0" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-waQZAvv44KYzys6S3l25ti2jcSuJnkyWFTliSKy3swASL6w6ttPxJTm80d+v9sLWoIxrqE3OwhTJbweNp065fg=="], + "ai": ["ai@7.0.162", "", { "dependencies": { "@ai-sdk/gateway": "4.0.135", "@ai-sdk/provider": "4.0.35", "@ai-sdk/provider-utils": "5.0.74" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-NPzL6/GmjjPwhzM1kGP9O60TfJ3uXV4Z4A/ztljxcJFAmRAFAjgk5mAHaikk3sE8NQo4NbRCiS+DgHinE2d0sQ=="], "ansis": ["ansis@4.2.0", "", {}, "sha512-HqZ5rWlFjGiV0tDm3UxxgNRqsOTniqoKZu0pIAfh7TZQMGuZK+hH0drySty0si0QXj1ieop4+SkSfPZBPPkHig=="], @@ -131,7 +132,7 @@ "empathic": ["empathic@2.0.0", "", {}, "sha512-i6UzDscO/XfAcNYD75CfICkmfLedpyPDdozrLMmQc5ORaQcdMoc21OnlEylMIqI7U8eniKrPMxxtj8k0vhmJhA=="], - "eventsource-parser": ["eventsource-parser@3.0.6", "", {}, "sha512-Vo1ab+QXPzZ4tCa8SwIHJFaSzy4R6SHf7BY79rFBDf0idraZWAkYrDjDj8uWaSm3S2TK+hJ7/t1CEmZ7jXw+pg=="], + "eventsource-parser": ["eventsource-parser@3.1.1", "", {}, "sha512-EKN1vKAMcZ8MlYMpaNuxN6R9yakzH6uajHcHVTqWJzvu5pWw9DyhbP35HH8MVBQ+dZjAfDxk+A8NiR9KWaXiyQ=="], "fdir": ["fdir@6.5.0", "", { "peerDependencies": { "picomatch": "^3 || ^4" }, "optionalPeers": ["picomatch"] }, "sha512-tIbYtZbucOs0BRGqPJkshJUYdL+SDH7dVM8gjy+ERp3WAUjLEFJE+02kanyHtwjWOnwrKYBiwAmM0p4kLJAnXg=="], @@ -187,8 +188,10 @@ "unconfig": ["unconfig@7.3.3", "", { "dependencies": { "@quansync/fs": "^0.1.5", "defu": "^6.1.4", "jiti": "^2.5.1", "quansync": "^0.2.11" } }, "sha512-QCkQoOnJF8L107gxfHL0uavn7WD9b3dpBcFX6HtfQYmjw2YzWxGuFQ0N0J6tE9oguCBJn9KOvfqYDCMPHIZrBA=="], + "undici": ["undici@7.30.1", "", {}, "sha512-03KrMbjIFDWQR5RNPrSw3bd1r12JpyFQl9/yxjUbJtqkjmllqC5lFjZh6ujeuJCKr8EOQvRW872Oy/az+nLYxg=="], + "undici-types": ["undici-types@7.14.0", "", {}, "sha512-QQiYxHuyZ9gQUIrmPo3IA+hUl4KYk8uSA7cHrcKd/l3p1OTpZcM0Tbp9x7FAtXdAYhlasd60ncPpgu6ihG6TOA=="], - "zod": ["zod@4.1.12", "", {}, "sha512-JInaHOamG8pt5+Ey8kGmdcAcg3OL9reK8ltczgHTAwNhMys/6ThXHityHxVV2p3fkw/c+MAvBHFVYHFZDmjMCQ=="], + "zod": ["zod@4.6.7", "", {}, "sha512-g8W2D2kU6a8untc053fX1uVBKzM7yqAxW/3zwk83+6RyF2+SKOuV+c0wBDhUwzV1PaGapmZH+rmcUzPf+3CGwg=="], } } diff --git a/package.json b/package.json index 6af94fd..ad5b475 100644 --- a/package.json +++ b/package.json @@ -8,16 +8,17 @@ }, "devDependencies": { "@types/bun": "latest", - "ai": "^5.0.105", + "ai": "^7.0.162", "base64-js": "^1.5.1", "gpt-tokenizer": "^3.0.1", "mitata": "^1.0.34", "prettier": "^3.6.2", "tiktoken": "^1.0.22", - "tsdown": "^0.15.6" + "tsdown": "^0.15.6", + "zod": "^4.6.7" }, "peerDependencies": { - "ai": "^5.0.0" + "ai": "^5.0.0 || ^6.0.0 || ^7.0.0" }, "peerDependenciesMeta": { "ai": { diff --git a/src/sdk.compat.test.ts b/src/sdk.compat.test.ts new file mode 100644 index 0000000..a127cbd --- /dev/null +++ b/src/sdk.compat.test.ts @@ -0,0 +1,919 @@ +/** + * AI SDK compatibility tests. + * + * These run against whichever `ai` and `zod` are installed. CI installs + * ai@5, ai@6 and ai@7 with zod 3 and zod 4 (see scripts/test-sdk-matrix.sh). + * Message fixtures are plain objects, so every fixture runs in every cell: + * `since` only records the first AI SDK major that accepts the shape. + */ + +import { describe, test, expect } from "bun:test"; +import { createRequire } from "node:module"; +import { jsonSchema, type ModelMessage, type ToolSet } from "ai"; +import * as z3 from "zod/v3"; +import * as z4 from "zod/v4"; +import { count } from "./sdk"; +import { Tokenizer, type Model, type ModelTokens } from "./index"; +import * as o200k from "./encoding/o200k_base"; + +const aiMajor = Number( + ( + createRequire(import.meta.url)("ai/package.json") as { version: string } + ).version.split(".")[0] +); + +// Every overhead field has its own value, so a total can only +// match when each field was applied to the right thing. +const tokens: ModelTokens = { + contentMultiplier: 1.5, + baseOverhead: 7, + perMessage: 3, + toolsExist: 311, + perTool: 53, + perDesc: 5, + perFirstProp: 13, + perAdditionalProp: 17, + perPropDesc: 19, + perEnum: 23, + perNestedObject: 29, + perArrayOfObjects: 31, +}; +const model = { tokens } as Model; +const tokenizer = new Tokenizer(o200k); +const enc = (text: string) => tokenizer.count(text); + +function countTools(tools: Record) { + return count({ + tokenizer, + model, + messages: [], + tools: tools as unknown as ToolSet, + }).tools; +} + +function schemaTokens(inputSchema: unknown): number { + return countTools({ t: { inputSchema } }).definitions.t!.inputSchema; +} + +describe("tool schemas", () => { + // Shapes from the zod 4 `_def` walker that this release replaced. The + // numbers were produced by that walker (src/sdk.ts at 107814b) with the + // `tokens` above, so a change here is a change in calibrated totals. + const unchanged = (z: any) => ({ + flat: [z.object({ a: z.string(), b: z.number(), c: z.boolean() }), 50], + described: [ + z.object({ + city: z.string().describe("The city name"), + days: z.number().describe("Number of days"), + bare: z.boolean(), + }), + 94, + ], + enumNoDesc: [z.object({ units: z.enum(["celsius", "fahrenheit"]) }), 41], + enumDesc: [ + z.object({ + units: z.enum(["celsius", "fahrenheit", "kelvin"]).describe("Unit"), + }), + 63, + ], + nested: [ + z.object({ + filters: z + .object({ + category: z.string().describe("Category"), + range: z + .object({ start: z.string(), end: z.string().describe("End") }) + .describe("Range"), + }) + .describe("Filters"), + top: z.string(), + }), + 234, + ], + nestedNoDesc: [ + z.object({ outer: z.object({ inner: z.object({ leaf: z.string() }) }) }), + 100, + ], + arrayOfStrings: [ + z.object({ + tags: z.array(z.string()).describe("Tags"), + ids: z.array(z.number()), + }), + 52, + ], + arrayOfObjects: [ + z.object({ + items: z + .array( + z.object({ + name: z.string().describe("Item name"), + qty: z.number(), + }) + ) + .describe("Items"), + }), + 118, + ], + arrayInNested: [ + z.object({ + order: z.object({ + lines: z.array( + z.object({ sku: z.string(), kind: z.enum(["a", "b"]) }) + ), + }), + }), + 145, + ], + arrayOfArrays: [ + z.object({ grid: z.array(z.array(z.object({ x: z.number() }))) }), + 14, + ], + arrayOfEnums: [ + z.object({ modes: z.array(z.enum(["r", "w"])).describe("Modes") }), + 35, + ], + emptyObject: [z.object({}), 0], + strict: [z.object({ a: z.string().describe("A") }).strict(), 34], + passthrough: [z.object({ a: z.string() }).passthrough(), 14], + topLevelDescription: [ + z.object({ a: z.string() }).describe("Top level description"), + 14, + ], + literalUnionRecordAny: [ + z.object({ + lit: z.literal("x").describe("Lit"), + uni: z.union([z.string(), z.number()]).describe("Uni"), + rec: z.record(z.string(), z.string()), + any: z.any(), + }), + 108, + ], + unicodeKeys: [ + z.object({ größe: z.string().describe("Größe in cm"), 名前: z.string() }), + 55, + ], + defaultWithOuterDescription: [ + z.object({ + n: z.number().default(3).describe("N"), + s: z.string().default("x"), + }), + 52, + ], + optionalWithOuterDescription: [ + z.object({ + a: z.string().optional().describe("Desc on optional wrapper"), + }), + 37, + ], + nullableWithOuterDescription: [ + z.object({ + a: z.string().nullable().describe("Desc on nullable wrapper"), + }), + 37, + ], + unionOfObjects: [ + z.object({ + u: z + .union([z.object({ a: z.string() }), z.object({ b: z.string() })]) + .describe("U"), + }), + 34, + ], + }); + + // Shapes the old walker skipped because a wrapper hid the inner type. The + // provider sees the inner description, enum and properties, so they count. + // `legacy` is what the old walker returned. + const changed = (z: any) => ({ + optionalWithInnerDescription: [ + z.object({ a: z.string().describe("Desc inside").optional() }), + 35, + 14, + ], + nullableWithInnerDescription: [ + z.object({ a: z.string().describe("Desc inside").nullable() }), + 35, + 14, + ], + nullableEnum: [z.object({ u: z.enum(["a", "b"]).nullable() }), 39, 14], + nullableObject: [ + z.object({ o: z.object({ x: z.string() }).nullable() }), + 57, + 14, + ], + nullableObjectWithOuterDescription: [ + z.object({ + o: z + .object({ x: z.string().describe("X") }) + .describe("O") + .nullable(), + }), + 97, + 14, + ], + nullableArrayOfObjects: [ + z.object({ l: z.array(z.object({ x: z.string() })).nullable() }), + 59, + 14, + ], + nullableAndOptional: [ + z.object({ + q: z.string().describe("Q").nullable().optional(), + r: z.number().nullish(), + }), + 52, + 32, + ], + optionalEnum: [z.object({ u: z.enum(["a", "b"]).optional() }), 39, 14], + optionalEnumWithOuterDescription: [ + z.object({ u: z.enum(["a", "b"]).optional().describe("U") }), + 59, + 34, + ], + optionalObject: [ + z.object({ o: z.object({ x: z.string().describe("X") }).optional() }), + 77, + 14, + ], + optionalArrayOfObjects: [ + z.object({ l: z.array(z.object({ x: z.string() })).optional() }), + 59, + 14, + ], + defaultObject: [ + z.object({ o: z.object({ x: z.string() }).default({ x: "" }) }), + 57, + 14, + ], + readonlyEnum: [z.object({ e: z.enum(["a", "b"]).readonly() }), 39, 14], + mixedRealistic: [ + z.object({ + query: z.string().describe("Search query"), + limit: z.number().optional().describe("Max results"), + sort: z.enum(["asc", "desc"]).optional(), + filters: z + .object({ + since: z.string().optional(), + tags: z.array(z.string()).optional(), + }) + .optional(), + }), + 196, + 110, + ], + }); + + for (const [label, z] of [ + ["zod 3", z3], + ["zod 4", z4], + ] as const) { + describe(label, () => { + test.each(Object.entries(unchanged(z)))( + "%s counts as the zod 4 walker did", + (_name, [schema, expected]) => { + expect(schemaTokens(schema)).toBe(expected as number); + } + ); + + test.each(Object.entries(changed(z)))( + "%s counts the inner type behind a wrapper", + (_name, [schema, expected, legacy]) => { + expect(schemaTokens(schema)).toBe(expected as number); + expect(expected).toBeGreaterThan(legacy as number); + } + ); + }); + } + + test("numeric TypeScript enums count their values, not the reverse mapping", () => { + enum Level { + Low, + High, + } + expect(schemaTokens(z4.object({ e: z4.enum(Level) }))).toBe( + schemaTokens(z4.object({ e: z4.enum({ A: 1, B: 2 } as const) })) + ); + }); + + const sourceOfTruth = { + type: "object", + properties: { + query: { type: "string", description: "Search query" }, + units: { type: "string", enum: ["c", "f"] }, + filters: { + type: "object", + properties: { since: { type: "string" } }, + }, + items: { + type: "array", + items: { type: "object", properties: { id: { type: "number" } } }, + }, + }, + }; + const equivalent = (z: any) => + z.object({ + query: z.string().describe("Search query"), + units: z.enum(["c", "f"]), + filters: z.object({ since: z.string() }), + items: z.array(z.object({ id: z.number() })), + }); + + test("zod 3, zod 4 and jsonSchema() tools give one total", () => { + const expected = schemaTokens(jsonSchema(sourceOfTruth as never)); + expect(expected).toBeGreaterThan(0); + expect(schemaTokens(equivalent(z3))).toBe(expected); + expect(schemaTokens(equivalent(z4))).toBe(expected); + }); + + test("a jsonSchema() thunk is read once", () => { + let reads = 0; + const schema = jsonSchema((() => { + reads++; + return sourceOfTruth; + }) as never); + const first = schemaTokens(schema); + expect(schemaTokens(schema)).toBe(first); + expect(reads).toBe(1); + }); + + const standard = (convert: unknown) => ({ + "~standard": { + version: 1, + vendor: "test", + validate: (value: unknown) => ({ value }), + ...(convert ? { jsonSchema: { input: () => convert } } : {}), + }, + }); + + test.if(aiMajor >= 6)( + "a Standard Schema with JSON Schema conversion counts (ai v6+)", + () => { + expect(schemaTokens(standard(sourceOfTruth))).toBe( + schemaTokens(jsonSchema(sourceOfTruth as never)) + ); + } + ); + + test.if(aiMajor === 5)( + "a Standard Schema is unresolved on ai v5, which only knows zod", + () => { + const tools = countTools({ t: { inputSchema: standard(sourceOfTruth) } }); + expect(tools.unresolved).toHaveLength(1); + expect(tools.unresolved[0]).toMatchObject({ + tool: "t", + part: "inputSchema", + reason: "error", + }); + } + ); + + describe("schemas count cannot read synchronously count 0 and are listed", () => { + test("a Promise", async () => { + const promise = Promise.resolve(sourceOfTruth); + const tools = countTools({ + t: { inputSchema: jsonSchema(promise as never) }, + }); + expect(tools.definitions.t!.inputSchema).toBe(0); + expect(tools.unresolved).toEqual([ + { tool: "t", part: "inputSchema", reason: "async" }, + ]); + expect(tools.total).toBe(tokens.toolsExist + enc("t")); + await promise; + }); + + test("a thunk that returns a Promise", () => { + const tools = countTools({ + t: { inputSchema: jsonSchema((async () => sourceOfTruth) as never) }, + }); + expect(tools.unresolved).toEqual([ + { tool: "t", part: "inputSchema", reason: "async" }, + ]); + }); + + test("a rejecting Promise does not become an unhandled rejection", async () => { + const unhandled: unknown[] = []; + const onUnhandled = (reason: unknown) => unhandled.push(reason); + process.on("unhandledRejection", onUnhandled); + try { + const tools = countTools({ + t: { + inputSchema: jsonSchema(Promise.reject(new Error("boom")) as never), + }, + }); + expect(tools.unresolved[0]).toMatchObject({ reason: "async" }); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(unhandled).toEqual([]); + } finally { + process.off("unhandledRejection", onUnhandled); + } + }); + + test("a Standard Schema vendor without JSON Schema conversion", () => { + const tools = countTools({ t: { inputSchema: standard(undefined) } }); + expect(tools.unresolved).toHaveLength(1); + expect(tools.unresolved[0]).toMatchObject({ reason: "error" }); + expect(tools.unresolved[0]!.error).toBeInstanceOf(Error); + }); + + test("one unreadable tool leaves the other tools exact", () => { + const only = countTools({ + good: { inputSchema: z4.object({ a: z4.string() }) }, + }); + const both = countTools({ + bad: { inputSchema: standard(undefined) }, + good: { inputSchema: z4.object({ a: z4.string() }) }, + }); + expect(both.definitions.good).toEqual(only.definitions.good); + expect(both.unresolved.map((u) => u.tool)).toEqual(["bad"]); + }); + }); + + describe("descriptions", () => { + test("a v7 description function is called with an empty context", () => { + const tools = countTools({ + t: { + description: ({ context }: { context: object }) => + `ctx ${Object.keys(context).length}`, + }, + }); + expect(tools.definitions.t!.description).toBe( + tokens.perDesc + enc("ctx 0") + ); + expect(tools.unresolved).toEqual([]); + }); + + test("a description function that throws is listed, not thrown", () => { + const tools = countTools({ + t: { + description: ({ context }: { context: { repo: { name: string } } }) => + context.repo.name, + }, + }); + expect(tools.definitions.t!.description).toBe(0); + expect(tools.unresolved).toMatchObject([ + { tool: "t", part: "description", reason: "error" }, + ]); + }); + + test("a provider-defined tool has no schema and counts only its name", () => { + const tools = countTools({ + search: { type: "provider", id: "x.search", args: {} }, + }); + expect(tools.definitions.search).toEqual({ + name: enc("search"), + description: 0, + inputSchema: 0, + }); + expect(tools.unresolved).toEqual([]); + }); + }); +}); + +type Role = "system" | "user" | "assistant" | "tool"; + +interface Fixture { + name: string; + /** First AI SDK major that accepts the shape. */ + since: 5 | 6 | 7; + role: Role; + /** String content, or the one part under test. */ + content: unknown; + /** Texts the provider sees. Each is tokenized on its own. */ + text?: string[]; + /** Fixed media estimates, never multiplied. */ + media?: number; +} + +const IMAGE = 85; +const FILE = 100; +const B64 = "iVBORw0KGgo="; +const toolResult = (output: unknown) => ({ + type: "tool-result", + toolCallId: "call_1", + toolName: "t", + output, +}); +const denied = (output: unknown) => ({ role: "tool" as const, output }); + +const fixtures: Fixture[] = [ + { + name: "system string", + since: 5, + role: "system", + content: "Be brief.", + text: ["Be brief."], + }, + { + name: "user string", + since: 5, + role: "user", + content: "Hello!", + text: ["Hello!"], + }, + { + name: "text part", + since: 5, + role: "user", + content: { type: "text", text: "Hello there" }, + text: ["Hello there"], + }, + { + name: "image part, base64", + since: 5, + role: "user", + content: { type: "image", image: B64, mediaType: "image/png" }, + media: IMAGE, + }, + { + name: "image part, URL", + since: 5, + role: "user", + content: { type: "image", image: new URL("https://example.com/a.png") }, + media: IMAGE, + }, + { + name: "file part, pdf", + since: 5, + role: "user", + content: { type: "file", data: B64, mediaType: "application/pdf" }, + media: FILE, + }, + { + name: "file part with an image media type", + since: 5, + role: "user", + content: { type: "file", data: B64, mediaType: "image/png" }, + media: IMAGE, + }, + { + name: "assistant reasoning", + since: 5, + role: "assistant", + content: { type: "reasoning", text: "Let me think." }, + text: ["Let me think."], + }, + { + name: "assistant generated image", + since: 5, + role: "assistant", + content: { type: "file", data: B64, mediaType: "image/webp" }, + media: IMAGE, + }, + { + name: "output text", + since: 5, + role: "tool", + content: toolResult({ type: "text", value: "sunny" }), + text: ["sunny"], + }, + { + name: "output error-text", + since: 5, + role: "tool", + content: toolResult({ type: "error-text", value: "not found" }), + text: ["not found"], + }, + { + name: "output json counts the value, not the wrapper", + since: 5, + role: "tool", + content: toolResult({ type: "json", value: { temp: 21 } }), + text: ['{"temp":21}'], + }, + { + name: "output error-json", + since: 5, + role: "tool", + content: toolResult({ type: "error-json", value: { code: 404 } }), + text: ['{"code":404}'], + }, + { + name: "output json with a falsy value", + since: 5, + role: "tool", + content: toolResult({ type: "json", value: 0 }), + text: ["0"], + }, + { + name: "output content: text and media", + since: 5, + role: "tool", + content: toolResult({ + type: "content", + value: [ + { type: "text", text: "Here is the chart" }, + { type: "media", data: B64, mediaType: "image/png" }, + { type: "media", data: B64, mediaType: "application/pdf" }, + ], + }), + text: ["Here is the chart"], + media: IMAGE + FILE, + }, + { + name: "output as a bare string (legacy)", + since: 5, + role: "tool", + content: toolResult("plain"), + text: ["plain"], + }, + { + name: "unknown output shape falls back to JSON", + since: 5, + role: "tool", + content: toolResult({ foo: 1 }), + text: ['{"foo":1}'], + }, + { + name: "unknown message part is skipped", + since: 5, + role: "user", + content: { type: "future-part", text: "ignored" }, + }, + + { + name: "output execution-denied with a reason", + since: 6, + role: "tool", + content: toolResult({ type: "execution-denied", reason: "User said no" }), + text: ["User said no"], + }, + { + name: "output execution-denied without a reason", + since: 6, + role: "tool", + content: toolResult({ type: "execution-denied" }), + }, + { + name: "output content: image-* items are images", + since: 6, + role: "tool", + content: toolResult({ + type: "content", + value: [ + { type: "image-data", data: B64, mediaType: "image/png" }, + { type: "image-url", url: "https://example.com/a.png" }, + { type: "image-file-id", fileId: "file_1" }, + ], + }), + media: 3 * IMAGE, + }, + { + name: "output content: file-* items are files unless the media type is an image", + since: 6, + role: "tool", + content: toolResult({ + type: "content", + value: [ + { type: "file-data", data: B64, mediaType: "application/pdf" }, + { type: "file-url", url: "https://example.com/a.pdf" }, + { type: "file-id", fileId: "file_2" }, + { type: "file-data", data: B64, mediaType: "image/jpeg" }, + { type: "custom", providerOptions: {} }, + ], + }), + media: 3 * FILE + IMAGE, + }, + { + name: "tool-approval-request is not sent to the provider", + since: 6, + role: "assistant", + content: { + type: "tool-approval-request", + approvalId: "a", + toolCallId: "call_1", + }, + }, + { + name: "tool-approval-response is not sent to the provider", + since: 6, + role: "tool", + content: { + type: "tool-approval-response", + approvalId: "a", + approved: true, + }, + }, + + { + name: "file with media type 'image' and URL data", + since: 7, + role: "user", + content: { + type: "file", + mediaType: "image", + data: { type: "url", url: new URL("https://example.com/a.png") }, + }, + media: IMAGE, + }, + { + name: "file with image/* media type and tagged data", + since: 7, + role: "user", + content: { + type: "file", + mediaType: "image/png", + data: { type: "data", data: B64 }, + }, + media: IMAGE, + }, + { + name: "file with image media type and a provider reference", + since: 7, + role: "user", + content: { + type: "file", + mediaType: "image", + data: { type: "reference", reference: { openai: "file_1" } }, + }, + media: IMAGE, + }, + { + name: "file pdf with tagged data", + since: 7, + role: "user", + content: { + type: "file", + mediaType: "application/pdf", + data: { type: "data", data: new Uint8Array([1, 2, 3]) }, + }, + media: FILE, + }, + { + name: "file with inline text data is text, not a file", + since: 7, + role: "user", + content: { + type: "file", + mediaType: "text/plain", + data: { type: "text", text: "notes for the model" }, + }, + text: ["notes for the model"], + }, + { + name: "image part with a provider reference", + since: 7, + role: "user", + content: { type: "image", image: { openai: "file_1" } }, + media: IMAGE, + }, + { + name: "reasoning-file image", + since: 7, + role: "assistant", + content: { + type: "reasoning-file", + mediaType: "image/png", + data: { type: "data", data: B64 }, + }, + media: IMAGE, + }, + { + name: "reasoning-file other", + since: 7, + role: "assistant", + content: { + type: "reasoning-file", + mediaType: "application/octet-stream", + data: B64, + }, + media: FILE, + }, + { + name: "custom part", + since: 7, + role: "assistant", + content: { type: "custom", kind: "openai.compaction" }, + }, + { + name: "output content: file items with tagged data", + since: 7, + role: "tool", + content: toolResult({ + type: "content", + value: [ + { + type: "file", + mediaType: "image", + data: { type: "url", url: new URL("https://example.com/a.png") }, + }, + { + type: "file", + mediaType: "application/pdf", + data: { type: "data", data: B64 }, + }, + { + type: "file", + mediaType: "application/pdf", + data: { type: "reference", reference: { anthropic: "file_3" } }, + }, + { + type: "file", + mediaType: "text/plain", + data: { type: "text", text: "inline" }, + }, + ], + }), + text: ["inline"], + media: IMAGE + 2 * FILE, + }, +]; + +describe("message parts", () => { + test.each(fixtures)( + "v$since: $name", + ({ role, content, text = [], media = 0 }) => { + const message = { + role, + content: typeof content === "string" ? content : [content], + } as unknown as ModelMessage; + const [counted] = count({ + tokenizer, + model, + messages: [message], + }).messages; + + const textTokens = text.reduce((sum, t) => sum + enc(t), 0); + const expected = + Math.round(textTokens * tokens.contentMultiplier) + media; + const [part] = counted!.content; + if (!part) { + expect(expected).toBe(0); + expect(counted!.total).toBe(tokens.perMessage + enc(role)); + return; + } + expect(part.type === "tool-result" ? part.output : part.total).toBe( + expected + ); + expect(counted!.total).toBe(tokens.perMessage + enc(role) + part.total); + } + ); + + test("a tool result total adds the call id to the output", () => { + const [counted] = count({ + tokenizer, + model, + messages: [ + { + role: "tool", + content: [ + toolResult({ + type: "content", + value: [ + { type: "text", text: "ok" }, + { type: "media", data: B64, mediaType: "image/png" }, + ], + }), + ], + }, + ] as unknown as ModelMessage[], + }).messages; + expect(counted!.content[0]).toEqual({ + type: "tool-result", + total: + Math.round((enc("call_1") + enc("ok")) * tokens.contentMultiplier) + + IMAGE, + output: Math.round(enc("ok") * tokens.contentMultiplier) + IMAGE, + }); + }); + + test("a tool call counts its name and input", () => { + const input = { location: "Paris" }; + const [counted] = count({ + tokenizer, + model, + messages: [ + { + role: "assistant", + content: [ + { + type: "tool-call", + toolCallId: "call_1", + toolName: "getWeather", + input, + }, + ], + }, + ] as ModelMessage[], + }).messages; + expect(counted!.content[0]).toEqual({ + type: "tool-call", + total: Math.round( + (enc("getWeather") + enc(JSON.stringify(input))) * + tokens.contentMultiplier + ), + input: Math.round(enc(JSON.stringify(input)) * tokens.contentMultiplier), + }); + }); + + test("a message with every part type adds up", () => { + const messages = fixtures + .filter((f) => f.role === "user") + .map((f) => ({ + role: "user", + content: [f.content], + })) as unknown as ModelMessage[]; + const result = count({ tokenizer, model, messages }); + const sum = result.messages.reduce((total, m) => total + m.total, 0); + expect(result.total).toBe(tokens.baseOverhead + sum); + }); +}); diff --git a/src/sdk.ts b/src/sdk.ts index a273c88..30438a0 100644 --- a/src/sdk.ts +++ b/src/sdk.ts @@ -1,5 +1,5 @@ import Tokenizer from "./tokenizer"; -import type { ModelMessage, ToolSet } from "ai"; +import { asSchema, type ModelMessage, type ToolSet } from "ai"; import type { Model, ModelTokens } from "."; export interface CountOptions { @@ -45,6 +45,22 @@ export interface CountResultMessage { content: Array; } +/** + * A tool description or input schema that `count()` could not read + * synchronously. It contributes 0 tokens to the total. + */ +export interface UnresolvedToolPart { + tool: string; + part: "description" | "inputSchema"; + /** + * `async`: the schema is a Promise, or a thunk that returns one. + * `error`: reading it threw, for example a Standard Schema vendor without + * JSON Schema conversion. The thrown value is in `error`. + */ + reason: "async" | "error"; + error?: unknown; +} + export interface CountResultTools { total: number; definitions: Record< @@ -55,6 +71,8 @@ export interface CountResultTools { inputSchema: number; } >; + /** Empty when every tool definition was counted in full. */ + unresolved: UnresolvedToolPart[]; } export interface CountResult { @@ -130,6 +148,7 @@ function countToolsTokensDetailed( inputSchema: number; } >; + const unresolved: UnresolvedToolPart[] = []; let total = 0; if (tools && Object.keys(tools).length > 0) { @@ -138,26 +157,46 @@ function countToolsTokensDetailed( const toolEntries = Object.entries(tools); for (let i = 0; i < toolEntries.length; i++) { const [toolName, tool] = toolEntries[i]!; + const definition = tool as { + description?: unknown; + inputSchema?: unknown; + }; // Count name tokens - const nameTokens = tokenizer.encode(toolName).length; + const nameTokens = tokenizer.count(toolName); // Count description tokens let descriptionTokens = 0; - if (tool.description) { - descriptionTokens = - config.perDesc + tokenizer.encode(tool.description).length; + const description = readDescription(definition.description); + if ("error" in description) { + unresolved.push({ + tool: toolName, + part: "description", + reason: "error", + error: description.error, + }); + } else if (description.text) { + descriptionTokens = config.perDesc + tokenizer.count(description.text); } // Count inputSchema tokens let inputSchemaTokens = 0; - if (tool.inputSchema) { - const result = countZodSchemaProperties( - tool.inputSchema, - tokenizer, - config - ); - inputSchemaTokens = result.tokens; + if (definition.inputSchema) { + const schema = readInputSchema(definition.inputSchema); + if (schema.status === "ok") { + inputSchemaTokens = countSchemaProperties( + schema.jsonSchema, + tokenizer, + config + ); + } else { + unresolved.push({ + tool: toolName, + part: "inputSchema", + reason: schema.status, + ...(schema.status === "error" && { error: schema.error }), + }); + } } definitions[toolName as keyof TOOLS] = { @@ -175,7 +214,158 @@ function countToolsTokensDetailed( } } - return { total, definitions }; + return { total, definitions, unresolved }; +} + +/** + * AI SDK v7 allows a function that builds the description from the tool's + * runtime context. `count()` has no context, so the function gets an empty one. + */ +function readDescription( + description: unknown +): { text: string | undefined } | { error: unknown } { + if (typeof description === "function") { + try { + const text: unknown = description({ context: {} }); + return { text: typeof text === "string" ? text : undefined }; + } catch (error) { + return { error }; + } + } + return { text: typeof description === "string" ? description : undefined }; +} + +type InputSchemaRead = + | { status: "ok"; jsonSchema: unknown } + | { status: "async" } + | { status: "error"; error: unknown }; + +/** + * Reads the JSON Schema the provider receives for a tool, whichever schema + * library defined it. `asSchema` is the AI SDK's own conversion, so zod 3, + * zod 4, `jsonSchema()` and (v6+) Standard Schema tools all arrive here as + * JSON Schema. + */ +function readInputSchema(inputSchema: unknown): InputSchemaRead { + try { + const jsonSchema: unknown = asSchema( + inputSchema as Parameters[0] + ).jsonSchema; + if (isThenable(jsonSchema)) { + // Mark the rejection as handled: we are not going to await it. + jsonSchema.then(undefined, () => {}); + return { status: "async" }; + } + return { status: "ok", jsonSchema }; + } catch (error) { + return { status: "error", error }; + } +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; +} + +function isThenable(value: unknown): value is PromiseLike { + return isRecord(value) && typeof value.then === "function"; +} + +/** + * Object schemas are recognised by `properties`, not `type: "object"`: + * AI SDK v5 omits `type` on the empty schema it makes for an undefined input. + */ +function propertiesOf(schema: unknown): Record | undefined { + return isRecord(schema) && isRecord(schema.properties) + ? schema.properties + : undefined; +} + +/** + * `.nullable()` becomes `anyOf: [inner, { type: "null" }]`. Read through it, so + * a nullable property counts like the same property without `null`. A real + * union (two or more non-null variants) is left alone. + */ +function withoutNull(schema: Record): Record { + const variants = schema.anyOf ?? schema.oneOf; + if (!Array.isArray(variants)) { + return schema; + } + const rest = variants.filter((v) => !(isRecord(v) && v.type === "null")); + const inner = rest[0]; + if ( + rest.length !== 1 || + rest.length === variants.length || + !isRecord(inner) + ) { + return schema; + } + return schema.description === undefined + ? inner + : { ...inner, description: schema.description }; +} + +/** + * Count the properties of a JSON Schema object by walking its structure. + * Unions (`anyOf`, `oneOf`, `allOf`) and `$ref` are not followed. + */ +function countSchemaProperties( + schema: unknown, + tokenizer: Tokenizer, + config: ModelTokens, + ancestors: Set = new Set() +): number { + const properties = propertiesOf(schema); + if (!properties || ancestors.has(properties)) { + return 0; + } + ancestors.add(properties); + + let tokens = 0; + let isFirstProp = true; + + for (const [key, value] of Object.entries(properties)) { + // Count property name tokens + tokens += tokenizer.count(key); + + // Add property overhead (first vs additional) + tokens += isFirstProp ? config.perFirstProp : config.perAdditionalProp; + isFirstProp = false; + + if (!isRecord(value)) { + continue; + } + const propSchema = withoutNull(value); + + // Count property description if present + if (typeof propSchema.description === "string" && propSchema.description) { + tokens += config.perPropDesc; + tokens += tokenizer.count(propSchema.description); + } + + // Handle enums - count each enum value + overhead + if (Array.isArray(propSchema.enum)) { + tokens += config.perEnum; + for (const enumValue of propSchema.enum) { + tokens += tokenizer.count(String(enumValue)); + } + } + + if (propertiesOf(propSchema)) { + // Handle nested objects + tokens += config.perNestedObject || 0; + tokens += countSchemaProperties(propSchema, tokenizer, config, ancestors); + } else if (isRecord(propSchema.items)) { + // Handle arrays with object items + const items = withoutNull(propSchema.items); + if (propertiesOf(items)) { + tokens += config.perArrayOfObjects || 0; + } + tokens += countSchemaProperties(items, tokenizer, config, ancestors); + } + } + + ancestors.delete(properties); + return tokens; } /** @@ -190,7 +380,7 @@ function countMessageTokensDetailed( const content: CountResultMessageContent[] = []; // Count role - total += tokenizer.encode(message.role).length; + total += tokenizer.count(message.role); // Count content with detailed breakdown if (message.content) { @@ -218,65 +408,61 @@ function countMessageContentTokensDetailed( const results: CountResultMessageContent[] = []; if (typeof content === "string") { - const tokens = tokenizer.encode(content).length; + const tokens = tokenizer.count(content); const total = Math.round(tokens * multiplier); results.push({ type: "text", total }); return results; } if (Array.isArray(content)) { - for (const part of content) { - if (part.type === "text") { - const tokens = tokenizer.encode(part.text).length; - const total = Math.round(tokens * multiplier); - results.push({ type: "text", total }); - } else if (part.type === "tool-call") { - let toolCallTokens = 0; - let inputTokens = 0; + // Parts are read as plain records: their shape differs between AI SDK + // v5, v6 and v7, and unknown parts are skipped. + for (const part of content as unknown[]) { + if (!isRecord(part)) { + continue; + } + const tally = createTally(tokenizer); - if (part.toolName) { - toolCallTokens += tokenizer.encode(part.toolName).length; + if (part.type === "text" || part.type === "reasoning") { + if (typeof part.text === "string") { + tally.addText(part.text); + } + results.push({ + type: "text", + total: Math.round(tally.text * multiplier) + tally.media, + }); + } else if (part.type === "tool-call") { + let nameTokens = 0; + if (part.toolName && typeof part.toolName === "string") { + nameTokens = tokenizer.count(part.toolName); } if (part.input) { - inputTokens = tokenizer.encode(JSON.stringify(part.input)).length; - toolCallTokens += inputTokens; + tally.addText(JSON.stringify(part.input)); } - const total = Math.round(toolCallTokens * multiplier); results.push({ type: "tool-call", - total, - input: Math.round(inputTokens * multiplier), + total: Math.round((nameTokens + tally.text) * multiplier), + input: Math.round(tally.text * multiplier), }); } else if (part.type === "tool-result") { - let toolResultTokens = 0; - let outputTokens = 0; - - if (part.toolCallId) { - toolResultTokens += tokenizer.encode(part.toolCallId).length; - } - if (part.output) { - if (typeof part.output === "string") { - outputTokens = tokenizer.encode(part.output).length; - } else { - outputTokens = tokenizer.encode(JSON.stringify(part.output)).length; - } - toolResultTokens += outputTokens; + let idTokens = 0; + if (part.toolCallId && typeof part.toolCallId === "string") { + idTokens = tokenizer.count(part.toolCallId); } + tallyToolOutput(tally, part.output); - const total = Math.round(toolResultTokens * multiplier); results.push({ type: "tool-result", - total, - output: Math.round(outputTokens * multiplier), + total: Math.round((idTokens + tally.text) * multiplier) + tally.media, + output: Math.round(tally.text * multiplier) + tally.media, + }); + } else if (typeof part.type === "string" && MEDIA_KINDS.has(part.type)) { + tallyMedia(tally, part); + results.push({ + type: "text", + total: Math.round(tally.text * multiplier) + tally.media, }); - } else if (part.type === "image") { - // Images have a fixed token count (approximate for GPT-4o) - results.push({ type: "text", total: 85 }); - } else if (part.type === "file") { - // Files are treated similarly to text - // This is an approximation - actual token count may vary - results.push({ type: "text", total: 100 }); } } } @@ -285,89 +471,153 @@ function countMessageContentTokensDetailed( } /** - * Count properties in a Zod schema by walking its structure + * Tokens of one content part. `text` is raw tokenizer tokens and still gets + * `contentMultiplier`. `media` is a fixed estimate and never does. */ -function countZodSchemaProperties( - schema: any, - tokenizer: Tokenizer, - config: ModelTokens -): { tokens: number; propCount: number } { - let tokens = 0; - let propCount = 0; - - if (!schema || !schema._def) { - return { tokens, propCount }; - } +function createTally(tokenizer: Tokenizer) { + const tally = { + text: 0, + media: 0, + addText(text: string) { + tally.text += tokenizer.count(text); + }, + addMedia(media: Media) { + tally.media += countMediaTokens(media); + }, + }; + return tally; +} - const def = schema._def; +type Tally = ReturnType; - // Handle object schemas - if (def.type === "object") { - const shape = typeof def.shape === "function" ? def.shape() : def.shape; +type MediaKind = "image" | "file"; - if (!shape) { - return { tokens, propCount }; - } - - for (const [key, value] of Object.entries(shape)) { - const isFirstProp = propCount === 0; - propCount++; - const propSchema = value as any; +interface Media { + kind: MediaKind; + /** IANA media type, `"image"` (AI SDK v7 wildcard) or undefined. */ + mediaType: string | undefined; + /** + * The part's payload as the caller gave it: base64 string, bytes, URL, + * AI SDK v7 tagged file data, or a provider file reference. + */ + data: unknown; +} - // Count property name tokens - tokens += tokenizer.encode(key).length; +/** + * The single place that prices an image or a file. Every image-like and + * file-like part reaches it, whatever its AI SDK shape. + */ +function countMediaTokens(media: Media): number { + // Images have a fixed token count (approximate for GPT-4o) + // Files are treated similarly to text, which is an approximation too. + return media.kind === "image" ? 85 : 100; +} - // Add property overhead (first vs additional) - tokens += isFirstProp ? config.perFirstProp : config.perAdditionalProp; +/** + * Message parts and tool-result content items that carry an image or a file, + * with the kind to use when the media type does not say. The names are from + * AI SDK v5 (`media`), v6 (`image-*`, `file-*`) and v7 (`file`, + * `reasoning-file`). + */ +const MEDIA_KINDS = new Map([ + ["image", "image"], + ["image-data", "image"], + ["image-url", "image"], + ["image-file-id", "image"], + ["image-file-reference", "image"], + ["file", "file"], + ["reasoning-file", "file"], + ["media", "file"], + ["file-data", "file"], + ["file-url", "file"], + ["file-id", "file"], + ["file-reference", "file"], +]); + +function tallyMedia(tally: Tally, part: Record): void { + const data = + part.data ?? + part.image ?? + part.url ?? + part.fileId ?? + part.providerReference; + + // AI SDK v7 inline text documents are text, not a binary file. + if (isRecord(data) && data.type === "text" && typeof data.text === "string") { + tally.addText(data.text); + return; + } - // Count property description if present - if (propSchema.description) { - tokens += config.perPropDesc; - tokens += tokenizer.encode(propSchema.description).length; - } + const mediaType = + typeof part.mediaType === "string" ? part.mediaType : undefined; + const isImageMediaType = + mediaType !== undefined && + (mediaType === "image" || mediaType.startsWith("image/")); + tally.addMedia({ + kind: + isImageMediaType || MEDIA_KINDS.get(part.type as string) === "image" + ? "image" + : "file", + mediaType, + data, + }); +} - // Handle enums - count each enum value + overhead - if (propSchema._def && propSchema._def.type === "enum") { - const enumDef = propSchema._def; - const enumValues = - enumDef.values || Object.values(enumDef.entries || {}); - tokens += config.perEnum; // Enum overhead - for (const enumValue of enumValues) { - tokens += tokenizer.encode(String(enumValue)).length; +/** + * Count the payload the provider sees for a tool result, not the + * `{ type, value }` wrapper around it. Shapes we do not know are counted as + * JSON, as they were before. + */ +function tallyToolOutput(tally: Tally, output: unknown): void { + if (!output) { + return; + } + if (typeof output === "string") { + tally.addText(output); + return; + } + if (isRecord(output)) { + switch (output.type) { + case "text": + case "error-text": + if (typeof output.value === "string") { + tally.addText(output.value); + return; } - } - - // Handle nested objects - if (propSchema._def) { - const innerDef = propSchema._def; - if (innerDef.type === "object") { - // Add overhead for nested object if configured - tokens += config.perNestedObject || 0; - const nested = countZodSchemaProperties( - propSchema, - tokenizer, - config - ); - tokens += nested.tokens; - // Don't accumulate propCount across nesting levels - each level is independent - } else if (innerDef.type === "array" && innerDef.element) { - // Handle arrays with object/complex items - const elementDef = innerDef.element._def; - if (elementDef && elementDef.type === "object") { - // Add overhead for array of objects if configured - tokens += config.perArrayOfObjects || 0; + break; + case "json": + case "error-json": + if (output.value !== undefined) { + tally.addText(JSON.stringify(output.value)); + } + return; + case "execution-denied": + if (typeof output.reason === "string") { + tally.addText(output.reason); + } + return; + case "content": + if (Array.isArray(output.value)) { + for (const item of output.value as unknown[]) { + tallyToolContentItem(tally, item); } - const nested = countZodSchemaProperties( - innerDef.element, - tokenizer, - config - ); - tokens += nested.tokens; - // Don't accumulate propCount across nesting levels + return; } - } + break; } } + tally.addText(JSON.stringify(output)); +} - return { tokens, propCount }; +function tallyToolContentItem(tally: Tally, item: unknown): void { + if (!isRecord(item)) { + return; + } + if (item.type === "text" && typeof item.text === "string") { + tally.addText(item.text); + } else if (typeof item.type === "string" && MEDIA_KINDS.has(item.type)) { + tallyMedia(tally, item); + } else if (item.type !== "custom") { + tally.addText(JSON.stringify(item)); + } } From b0041f090ceaf945d568b62db77945f61a8d1138 Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 3 Nov 2026 11:14:48 +0000 Subject: [PATCH 2/8] ci: test the sdk against ai 5, 6, 7 and zod 3, 4 scripts/test-sdk-matrix.sh installs one pairing, type-checks and tests the sdk, builds dist and loads the CJS and ESM builds. CI runs the cells the script lists. Change-Id: I689441a6f0a6fe1c93ab32e44a7abf0089177991 Signed-off-by: Thomas Kosiewski --- .github/workflows/ci.yml | 35 ++++++++++++++ scripts/smoke-sdk-dist.mjs | 64 +++++++++++++++++++++++++ scripts/test-sdk-matrix.sh | 98 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 197 insertions(+) create mode 100644 scripts/smoke-sdk-dist.mjs create mode 100755 scripts/test-sdk-matrix.sh diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9322aaa..e52b708 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -33,3 +33,38 @@ jobs: - name: Build run: bun run build --dts=false + sdk-matrix: + runs-on: ubuntu-latest + outputs: + cells: ${{ steps.cells.outputs.cells }} + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: List AI SDK cells + id: cells + run: echo "cells=$(scripts/test-sdk-matrix.sh --list-json)" >> "$GITHUB_OUTPUT" + + sdk-compat: + name: sdk (ai@${{ matrix.ai }}, zod@${{ matrix.zod }}, node ${{ matrix.node }}) + needs: sdk-matrix + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: ${{ fromJSON(needs.sdk-matrix.outputs.cells) }} + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Setup Bun + uses: oven-sh/setup-bun@v2 + with: + bun-version: latest + + - name: Setup Node + uses: actions/setup-node@v4 + with: + node-version: ${{ matrix.node }} + + - name: Test the sdk against this AI SDK and zod + run: scripts/test-sdk-matrix.sh "${{ matrix.ai }}" "${{ matrix.zod }}" diff --git a/scripts/smoke-sdk-dist.mjs b/scripts/smoke-sdk-dist.mjs new file mode 100644 index 0000000..0664b83 --- /dev/null +++ b/scripts/smoke-sdk-dist.mjs @@ -0,0 +1,64 @@ +// Loads the built `ai-tokenizer/sdk` as CommonJS and as ESM against the +// installed `ai` and `zod`, and fails unless both give the same counts. +// Run after `bun run build --dts=false`. +import assert from "node:assert/strict"; +import { createRequire } from "node:module"; + +const require = createRequire(import.meta.url); +const dist = new URL("../dist/", import.meta.url); + +const aiVersion = require("ai/package.json").version; +const zodVersion = require("zod/package.json").version; + +async function smoke(format) { + const load = async (name) => + format === "cjs" + ? require(new URL(`${name}.cjs`, dist).pathname) + : import(new URL(`${name}.js`, dist).href); + const { count } = await load("sdk"); + const { Tokenizer, models } = await load("index"); + const encoding = await load("encoding/o200k_base"); + const ai = format === "cjs" ? require("ai") : await import("ai"); + const { z } = format === "cjs" ? require("zod") : await import("zod"); + + const result = count({ + tokenizer: new Tokenizer(encoding), + model: models["openai/gpt-5"], + messages: [ + { role: "system", content: "Be brief." }, + { + role: "user", + content: [ + { type: "text", text: "Weather in Paris?" }, + { type: "image", image: "iVBORw0KGgo=", mediaType: "image/png" }, + ], + }, + ], + tools: { + zodTool: { + description: "Get the weather", + inputSchema: z.object({ city: z.string().describe("City name") }), + }, + jsonTool: { + description: "Get the weather", + inputSchema: ai.jsonSchema({ + type: "object", + properties: { city: { type: "string", description: "City name" } }, + }), + }, + }, + }); + + assert.equal(result.tools.unresolved.length, 0, `${format}: unresolved tools`); + const { zodTool, jsonTool } = result.tools.definitions; + assert.ok(zodTool.inputSchema > 0, `${format}: zod schema counted`); + assert.equal(zodTool.inputSchema, jsonTool.inputSchema, `${format}: zod = jsonSchema`); + return result.total; +} + +const cjs = await smoke("cjs"); +const esm = await smoke("esm"); +assert.equal(cjs, esm, "CJS and ESM totals differ"); +console.log( + `smoke node=${process.version} ai=${aiVersion} zod=${zodVersion} cjs=${cjs} esm=${esm}` +); diff --git a/scripts/test-sdk-matrix.sh b/scripts/test-sdk-matrix.sh new file mode 100755 index 0000000..6acafa1 --- /dev/null +++ b/scripts/test-sdk-matrix.sh @@ -0,0 +1,98 @@ +#!/usr/bin/env bash +# Tests `ai-tokenizer/sdk` against one AI SDK + zod pairing, or against all of +# them. CI runs one cell per job. Every cell does four things: +# 1. installs the pairing in place of the dev dependencies, +# 2. type-checks src and runs the sdk tests, +# 3. builds dist, +# 4. loads the CJS and ESM builds with the installed `ai` and `zod`. +# +# scripts/test-sdk-matrix.sh +# Runs one cell in this checkout. It rewrites package.json, bun.lock and +# node_modules, so use it in a throwaway checkout such as a CI job. +# scripts/test-sdk-matrix.sh --all +# Runs every cell, each in a throwaway copy of the repo. +# scripts/test-sdk-matrix.sh --list-json +# Prints the cells as a GitHub Actions matrix. CI reads this list. +# +# Every cell prints one line starting with "CELL". +set -euo pipefail + +# ai version, zod version, Node major used in CI. Node 20 covers a v5/v6 user +# who loads the CJS build: v7 needs Node 22. +CELLS=( + "5 3.25.76 22" + "5 4 22" + "5 4 20" + "6 3.25.76 22" + "6 4 22" + "6 4 20" + "7 3.25.76 22" + "7 4 22" +) + +installed() { node -p "require('$1/package.json').version"; } + +run_cell() { + local ai="$1" zod="$2" log + # Bun resolves `ai` once for the dev dependency and the peer dependency and + # picks the peer range (^5 || ^6 || ^7), so drop the peer range here. + node -e ' + const fs = require("fs"); + const pkg = JSON.parse(fs.readFileSync("package.json", "utf8")); + pkg.devDependencies.ai = process.argv[1]; + pkg.devDependencies.zod = process.argv[2]; + delete pkg.peerDependencies; + delete pkg.peerDependenciesMeta; + fs.writeFileSync("package.json", JSON.stringify(pkg, null, 2) + "\n"); + ' "$ai" "$zod" + bun install >/dev/null + + # Only the sdk files are type-checked: scripts/ and bench/ use the newest AI SDK. + log="$(mktemp)" + local typecheck=ok tests + printf '{"extends":"./tsconfig.json","include":["src/sdk*.ts"]}' >tsconfig.sdk-cell.json + bunx tsc --noEmit -p tsconfig.sdk-cell.json >"$log" 2>&1 || { cat "$log"; typecheck=FAILED; } + rm -f tsconfig.sdk-cell.json + + bun test src/sdk.test.ts src/sdk.compat.test.ts >"$log" 2>&1 || { cat "$log"; exit 1; } + tests="$(grep -E '^ *[0-9]+ (pass|skip|fail)$' "$log" | tr -s ' ' | tr '\n' ' ')" + + bun run build --dts=false >/dev/null + local smoke + smoke="$(node scripts/smoke-sdk-dist.mjs)" + + echo "CELL ai@$(installed ai) zod@$(installed zod) typecheck=$typecheck tests:$tests| $smoke" + [ "$typecheck" = ok ] +} + +case "${1:-}" in + --list-json) + printf '{"include":[' + sep="" + for cell in "${CELLS[@]}"; do + read -r ai zod node <<<"$cell" + printf '%s{"ai":"%s","zod":"%s","node":"%s"}' "$sep" "$ai" "$zod" "$node" + sep="," + done + printf ']}\n' + ;; + --all) + root="$(cd "$(dirname "$0")/.." && pwd)" + failed=0 + for cell in "${CELLS[@]}"; do + read -r ai zod _node <<<"$cell" + tmp="$(mktemp -d)" + (cd "$root" && cp -r package.json bun.lock tsconfig.json tsdown.config.ts src scripts "$tmp") + (cd "$tmp" && bash scripts/test-sdk-matrix.sh "$ai" "$zod") || failed=1 + rm -rf "$tmp" + done + exit "$failed" + ;; + "") + echo "usage: $0 | --all | --list-json" >&2 + exit 2 + ;; + *) + run_cell "$1" "${2:?zod version required}" + ;; +esac From 26f927e725ee6700d5afbe9c2cccc8b2fffacc3d Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 3 Nov 2026 11:14:48 +0000 Subject: [PATCH 3/8] docs: document supported AI SDK versions and unresolved tools Change-Id: I2287f560194044fd968fbbe77277d04b38e6d7e8 Signed-off-by: Thomas Kosiewski --- README.md | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 82f0177..fffd252 100644 --- a/README.md +++ b/README.md @@ -69,7 +69,8 @@ const result = count({ // description: 10, // inputSchema: 15 // } -// } +// }, +// unresolved: [] // } // } @@ -78,6 +79,23 @@ const result = count({ const costUSD = result.total * model.pricing.input; ``` +`ai-tokenizer/sdk` works with AI SDK v5, v6 and v7 and needs `ai` installed +(v7 needs Node 22 or newer). Messages in any of those versions' shapes are +counted, and tool input schemas are read with the AI SDK's own `asSchema()`, so +zod 3, zod 4, `jsonSchema()` and (v6 and v7) Standard Schema tools all count. + +`count()` is synchronous. It cannot wait for a tool whose `jsonSchema` is a +Promise (or a function that returns one), and it cannot read a schema the AI SDK +fails to convert. It counts that part as 0 tokens, does not throw, and lists it +in `result.tools.unresolved`: + +```ts +for (const { tool, part, reason } of result.tools.unresolved) { + console.warn(`${tool} ${part} not counted (${reason})`); +} +// Fix: resolve the schema first, e.g. jsonSchema(await loadSchema()). +``` + Estimate tokens using an encoding: ```ts From ac247b17914f01f1c8ac12e2483d2b5a7abe8b28 Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 3 Nov 2026 11:38:59 +0000 Subject: [PATCH 4/8] fix(sdk): follow local $ref, skip provider tool schemas, count the default execution-denied text Review findings: - zod 4 .meta({ id }) and pydantic MCP schemas put nested objects behind a local $ref. Resolve #, #/definitions/X and #/$defs/X so those tools count as before. - Provider tools carry an inputSchema the provider never receives. Count their name and description only. - OpenAI and Anthropic send 'Tool call execution denied.' when there is no reason. Count it. - A thenable whose then() throws is still unreadable synchronously. - Drop two assertions that could not fail, and make the matrix script fail loudly when bun prints no test summary. Change-Id: Idd3aa4dda1c425c87c8a0a44c58ee2319170d209 Signed-off-by: Thomas Kosiewski --- scripts/test-sdk-matrix.sh | 6 +- src/sdk.compat.test.ts | 147 ++++++++++++++++++++++++++++++------- src/sdk.ts | 112 +++++++++++++++++++++++----- 3 files changed, 218 insertions(+), 47 deletions(-) diff --git a/scripts/test-sdk-matrix.sh b/scripts/test-sdk-matrix.sh index 6acafa1..4a24caf 100755 --- a/scripts/test-sdk-matrix.sh +++ b/scripts/test-sdk-matrix.sh @@ -10,7 +10,8 @@ # Runs one cell in this checkout. It rewrites package.json, bun.lock and # node_modules, so use it in a throwaway checkout such as a CI job. # scripts/test-sdk-matrix.sh --all -# Runs every cell, each in a throwaway copy of the repo. +# Runs every cell, each in a throwaway copy of the repo, on the Node +# that is installed here (the Node column only applies in CI). # scripts/test-sdk-matrix.sh --list-json # Prints the cells as a GitHub Actions matrix. CI reads this list. # @@ -55,7 +56,8 @@ run_cell() { rm -f tsconfig.sdk-cell.json bun test src/sdk.test.ts src/sdk.compat.test.ts >"$log" 2>&1 || { cat "$log"; exit 1; } - tests="$(grep -E '^ *[0-9]+ (pass|skip|fail)$' "$log" | tr -s ' ' | tr '\n' ' ')" + tests="$(grep -E '^ *[0-9]+ (pass|skip|fail)$' "$log" | tr -s ' ' | tr '\n' ' ' || true)" + [ -n "$tests" ] || { cat "$log"; echo "no test summary found" >&2; exit 1; } bun run build --dts=false >/dev/null local smoke diff --git a/src/sdk.compat.test.ts b/src/sdk.compat.test.ts index a127cbd..0dd2606 100644 --- a/src/sdk.compat.test.ts +++ b/src/sdk.compat.test.ts @@ -277,9 +277,8 @@ describe("tool schemas", () => { test.each(Object.entries(changed(z)))( "%s counts the inner type behind a wrapper", - (_name, [schema, expected, legacy]) => { + (_name, [schema, expected]) => { expect(schemaTokens(schema)).toBe(expected as number); - expect(expected).toBeGreaterThan(legacy as number); } ); }); @@ -408,7 +407,21 @@ describe("tool schemas", () => { } }); - test("a Standard Schema vendor without JSON Schema conversion", () => { + test("a thenable whose then() throws", () => { + const thenable = { + then() { + throw new Error("boom"); + }, + }; + const tools = countTools({ + t: { inputSchema: jsonSchema(thenable as never) }, + }); + expect(tools.unresolved).toEqual([ + { tool: "t", part: "inputSchema", reason: "async" }, + ]); + }); + + test("a Standard Schema the AI SDK cannot convert to JSON Schema", () => { const tools = countTools({ t: { inputSchema: standard(undefined) } }); expect(tools.unresolved).toHaveLength(1); expect(tools.unresolved[0]).toMatchObject({ reason: "error" }); @@ -428,6 +441,89 @@ describe("tool schemas", () => { }); }); + describe("local $ref", () => { + const address = z4 + .object({ + street: z4.string().describe("Street"), + city: z4.string(), + }) + .describe("Postal address"); + const inline = (schema: unknown) => schemaTokens(schema); + + test("a zod 4 schema registered with .meta({ id }) counts like the inline schema", () => { + const registered = address.meta({ id: "Address" }); + expect(inline(z4.object({ home: registered, work: registered }))).toBe( + inline(z4.object({ home: address, work: address })) + ); + expect(inline(z4.object({ stops: z4.array(registered) }))).toBe( + inline(z4.object({ stops: z4.array(address) })) + ); + }); + + test("a top-level $ref and a $ref inside nullable are followed", () => { + const definitions = { + Address: { + type: "object", + description: "Postal address", + properties: { street: { type: "string", description: "Street" } }, + }, + }; + const direct = { + type: "object", + properties: { home: { $ref: "#/definitions/Address" } }, + definitions, + }; + const nullable = { + type: "object", + properties: { + home: { + anyOf: [{ $ref: "#/definitions/Address" }, { type: "null" }], + }, + }, + definitions, + }; + const expected = inline( + jsonSchema({ + type: "object", + properties: { home: definitions.Address }, + } as never) + ); + expect(inline(jsonSchema(direct as never))).toBe(expected); + expect(inline(jsonSchema(nullable as never))).toBe(expected); + const rootRef = { + $ref: "#/$defs/Root", + $defs: { + Root: { + type: "object", + properties: { home: { $ref: "#/$defs/Address" } }, + }, + Address: definitions.Address, + }, + }; + expect(inline(jsonSchema(rootRef as never))).toBe(expected); + }); + + test("a recursive $ref ends", () => { + const tree = { + type: "object", + properties: { child: { $ref: "#" } }, + }; + expect(inline(jsonSchema(tree as never))).toBe( + enc("child") + tokens.perFirstProp + tokens.perNestedObject + ); + }); + + test("a $ref to somewhere else counts nothing", () => { + const remote = { + type: "object", + properties: { a: { $ref: "https://example.com/a.json" } }, + }; + expect(inline(jsonSchema(remote as never))).toBe( + enc("a") + tokens.perFirstProp + ); + }); + }); + describe("descriptions", () => { test("a v7 description function is called with an empty context", () => { const tools = countTools({ @@ -455,17 +551,25 @@ describe("tool schemas", () => { ]); }); - test("a provider-defined tool has no schema and counts only its name", () => { - const tools = countTools({ - search: { type: "provider", id: "x.search", args: {} }, - }); - expect(tools.definitions.search).toEqual({ - name: enc("search"), - description: 0, - inputSchema: 0, - }); - expect(tools.unresolved).toEqual([]); - }); + test.each(["provider", "provider-defined"])( + "a %s tool counts its name, not the schema the provider ignores", + (type) => { + const tools = countTools({ + search: { + type, + id: "x.search", + args: {}, + inputSchema: z4.object({ q: z4.string() }), + }, + }); + expect(tools.definitions.search).toEqual({ + name: enc("search"), + description: 0, + inputSchema: 0, + }); + expect(tools.unresolved).toEqual([]); + } + ); }); }); @@ -638,10 +742,11 @@ const fixtures: Fixture[] = [ text: ["User said no"], }, { - name: "output execution-denied without a reason", + name: "output execution-denied without a reason counts the default text providers send", since: 6, role: "tool", content: toolResult({ type: "execution-denied" }), + text: ["Tool call execution denied."], }, { name: "output content: image-* items are images", @@ -904,16 +1009,4 @@ describe("message parts", () => { input: Math.round(enc(JSON.stringify(input)) * tokens.contentMultiplier), }); }); - - test("a message with every part type adds up", () => { - const messages = fixtures - .filter((f) => f.role === "user") - .map((f) => ({ - role: "user", - content: [f.content], - })) as unknown as ModelMessage[]; - const result = count({ tokenizer, model, messages }); - const sum = result.messages.reduce((total, m) => total + m.total, 0); - expect(result.total).toBe(tokens.baseOverhead + sum); - }); }); diff --git a/src/sdk.ts b/src/sdk.ts index 30438a0..2004a72 100644 --- a/src/sdk.ts +++ b/src/sdk.ts @@ -158,6 +158,7 @@ function countToolsTokensDetailed( for (let i = 0; i < toolEntries.length; i++) { const [toolName, tool] = toolEntries[i]!; const definition = tool as { + type?: unknown; description?: unknown; inputSchema?: unknown; }; @@ -181,10 +182,10 @@ function countToolsTokensDetailed( // Count inputSchema tokens let inputSchemaTokens = 0; - if (definition.inputSchema) { + if (definition.inputSchema && !isProviderTool(definition.type)) { const schema = readInputSchema(definition.inputSchema); if (schema.status === "ok") { - inputSchemaTokens = countSchemaProperties( + inputSchemaTokens = countToolSchema( schema.jsonSchema, tokenizer, config @@ -235,6 +236,14 @@ function readDescription( return { text: typeof description === "string" ? description : undefined }; } +/** + * Provider tools (`provider-defined` in v5, `provider` in v6 and v7) still + * carry an `inputSchema`, but the provider sends its own definition instead. + */ +function isProviderTool(type: unknown): boolean { + return type === "provider" || type === "provider-defined"; +} + type InputSchemaRead = | { status: "ok"; jsonSchema: unknown } | { status: "async" } @@ -252,8 +261,12 @@ function readInputSchema(inputSchema: unknown): InputSchemaRead { inputSchema as Parameters[0] ).jsonSchema; if (isThenable(jsonSchema)) { - // Mark the rejection as handled: we are not going to await it. - jsonSchema.then(undefined, () => {}); + try { + // Mark the rejection as handled: we are not going to await it. + jsonSchema.then(undefined, () => {}); + } catch { + // A thenable whose `then` throws is still not readable synchronously. + } return { status: "async" }; } return { status: "ok", jsonSchema }; @@ -280,6 +293,39 @@ function propertiesOf(schema: unknown): Record | undefined { : undefined; } +/** + * Follows a local `$ref` ("#", "#/definitions/X", "#/$defs/X") to its target in + * the root schema. zod 4 emits one for a schema registered with `.meta({ id })` + * and pydantic-based MCP servers emit them for nested models. A `description` + * next to the `$ref` wins over the target's. + */ +function resolveRef( + schema: Record, + root: unknown +): Record { + const ref = schema.$ref; + if (typeof ref !== "string" || !ref.startsWith("#")) { + return schema; + } + let target: unknown = root; + try { + for (const segment of ref.slice(1).split("/").slice(1)) { + const key = decodeURIComponent(segment) + .replace(/~1/g, "/") + .replace(/~0/g, "~"); + target = isRecord(target) ? target[key] : undefined; + } + } catch { + return schema; + } + if (!isRecord(target)) { + return schema; + } + return schema.description === undefined + ? target + : { ...target, description: schema.description }; +} + /** * `.nullable()` becomes `anyOf: [inner, { type: "null" }]`. Read through it, so * a nullable property counts like the same property without `null`. A real @@ -304,16 +350,43 @@ function withoutNull(schema: Record): Record { : { ...inner, description: schema.description }; } +interface SchemaWalk { + tokenizer: Tokenizer; + config: ModelTokens; + /** The whole tool schema, which local `$ref`s point into. */ + root: unknown; + /** `properties` objects being walked, so a recursive schema ends. */ + ancestors: Set; +} + +/** The schema a property stands for once `$ref` and `null` wrappers are read through. */ +function unwrap(value: Record, walk: SchemaWalk) { + return resolveRef(withoutNull(resolveRef(value, walk.root)), walk.root); +} + +function countToolSchema( + jsonSchema: unknown, + tokenizer: Tokenizer, + config: ModelTokens +): number { + if (!isRecord(jsonSchema)) { + return 0; + } + const walk: SchemaWalk = { + tokenizer, + config, + root: jsonSchema, + ancestors: new Set(), + }; + return countSchemaProperties(resolveRef(jsonSchema, jsonSchema), walk); +} + /** * Count the properties of a JSON Schema object by walking its structure. - * Unions (`anyOf`, `oneOf`, `allOf`) and `$ref` are not followed. + * Unions (`anyOf`, `oneOf`, `allOf`) and non-local `$ref`s are not followed. */ -function countSchemaProperties( - schema: unknown, - tokenizer: Tokenizer, - config: ModelTokens, - ancestors: Set = new Set() -): number { +function countSchemaProperties(schema: unknown, walk: SchemaWalk): number { + const { tokenizer, config, ancestors } = walk; const properties = propertiesOf(schema); if (!properties || ancestors.has(properties)) { return 0; @@ -334,7 +407,7 @@ function countSchemaProperties( if (!isRecord(value)) { continue; } - const propSchema = withoutNull(value); + const propSchema = unwrap(value, walk); // Count property description if present if (typeof propSchema.description === "string" && propSchema.description) { @@ -353,14 +426,14 @@ function countSchemaProperties( if (propertiesOf(propSchema)) { // Handle nested objects tokens += config.perNestedObject || 0; - tokens += countSchemaProperties(propSchema, tokenizer, config, ancestors); + tokens += countSchemaProperties(propSchema, walk); } else if (isRecord(propSchema.items)) { // Handle arrays with object items - const items = withoutNull(propSchema.items); + const items = unwrap(propSchema.items, walk); if (propertiesOf(items)) { tokens += config.perArrayOfObjects || 0; } - tokens += countSchemaProperties(items, tokenizer, config, ancestors); + tokens += countSchemaProperties(items, walk); } } @@ -592,9 +665,12 @@ function tallyToolOutput(tally: Tally, output: unknown): void { } return; case "execution-denied": - if (typeof output.reason === "string") { - tally.addText(output.reason); - } + // OpenAI and Anthropic send this text when there is no reason. + tally.addText( + typeof output.reason === "string" + ? output.reason + : "Tool call execution denied." + ); return; case "content": if (Array.isArray(output.value)) { From 5db44f02d0756cf7f3432bfb047c823009e3e43c Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 3 Nov 2026 12:18:30 +0000 Subject: [PATCH 5/8] fix(sdk): keep tool-result totals equal to main, skip reasoning, bound $ref walks - Count a tool result as JSON.stringify(output), wrapper included, as main did. Payload-only counting under-counted every result by about 6 tokens. A content output still prices its images and files with countMediaTokens instead of serialising them. - Skip reasoning and reasoning-file parts again. Anthropic's count_tokens gives the same total with and without an earlier turn's thinking block. - Remember the count of each walked properties object, so definitions that share definitions are walked once instead of once per path (depth 22: 24.9 s to under 5 ms). - Resolve only JSON pointer $refs (# and #/...), not anchors like #A. Change-Id: I6b807d6398e6ce560ae42d4f3d678b9f667c91ac Signed-off-by: Thomas Kosiewski --- src/sdk.compat.test.ts | 115 ++++++++++++++++++++++++------ src/sdk.ts | 158 ++++++++++++++++++++++------------------- 2 files changed, 179 insertions(+), 94 deletions(-) diff --git a/src/sdk.compat.test.ts b/src/sdk.compat.test.ts index 0dd2606..e8c8c9f 100644 --- a/src/sdk.compat.test.ts +++ b/src/sdk.compat.test.ts @@ -513,6 +513,69 @@ describe("tool schemas", () => { ); }); + test("an anchor is not resolved", () => { + const anchored = { + type: "object", + properties: { a: { $ref: "#A" } }, + $defs: { A: { $anchor: "A", properties: { b: { type: "string" } } } }, + }; + expect(inline(jsonSchema(anchored as never))).toBe( + enc("a") + tokens.perFirstProp + ); + }); + + test("definitions that reference each other count from where the walk enters", () => { + const first = tokens.perFirstProp + tokens.perNestedObject; + const next = tokens.perAdditionalProp + tokens.perNestedObject; + const mutual = { + type: "object", + properties: { a: { $ref: "#/$defs/A" }, b: { $ref: "#/$defs/B" } }, + $defs: { + A: { type: "object", properties: { b: { $ref: "#/$defs/B" } } }, + B: { type: "object", properties: { a: { $ref: "#/$defs/A" } } }, + }, + }; + // Each branch stops where it meets a definition it is already inside. + const viaA = enc("b") + first + (enc("a") + first); + const viaB = enc("a") + first + (enc("b") + first); + expect(inline(jsonSchema(mutual as never))).toBe( + enc("a") + first + viaA + (enc("b") + next + viaB) + ); + }); + + test("definitions that share definitions are walked once, not once per path", () => { + const depth = 22; + const $defs: Record = { + [`D${depth}`]: { + type: "object", + properties: { x: { type: "string" } }, + }, + }; + let expected = enc("x") + tokens.perFirstProp; + for (let level = depth - 1; level >= 0; level--) { + const next = { $ref: `#/$defs/D${level + 1}` }; + $defs[`D${level}`] = { + type: "object", + properties: { a: next, b: next }, + }; + expected = + enc("a") + + tokens.perFirstProp + + tokens.perNestedObject + + expected + + (enc("b") + + tokens.perAdditionalProp + + tokens.perNestedObject + + expected); + } + const dag = { $ref: "#/$defs/D0", $defs }; + + const start = performance.now(); + const counted = inline(jsonSchema(dag as never)); + expect(performance.now() - start).toBeLessThan(500); + expect(counted).toBe(expected); + }); + test("a $ref to somewhere else counts nothing", () => { const remote = { type: "object", @@ -650,11 +713,10 @@ const fixtures: Fixture[] = [ media: IMAGE, }, { - name: "assistant reasoning", + name: "assistant reasoning is not counted", since: 5, role: "assistant", content: { type: "reasoning", text: "Let me think." }, - text: ["Let me think."], }, { name: "assistant generated image", @@ -664,42 +726,42 @@ const fixtures: Fixture[] = [ media: IMAGE, }, { - name: "output text", + name: "output text counts its wrapper", since: 5, role: "tool", content: toolResult({ type: "text", value: "sunny" }), - text: ["sunny"], + text: ['{"type":"text","value":"sunny"}'], }, { name: "output error-text", since: 5, role: "tool", content: toolResult({ type: "error-text", value: "not found" }), - text: ["not found"], + text: ['{"type":"error-text","value":"not found"}'], }, { - name: "output json counts the value, not the wrapper", + name: "output json counts its wrapper", since: 5, role: "tool", content: toolResult({ type: "json", value: { temp: 21 } }), - text: ['{"temp":21}'], + text: ['{"type":"json","value":{"temp":21}}'], }, { name: "output error-json", since: 5, role: "tool", content: toolResult({ type: "error-json", value: { code: 404 } }), - text: ['{"code":404}'], + text: ['{"type":"error-json","value":{"code":404}}'], }, { name: "output json with a falsy value", since: 5, role: "tool", content: toolResult({ type: "json", value: 0 }), - text: ["0"], + text: ['{"type":"json","value":0}'], }, { - name: "output content: text and media", + name: "output content: text counts its wrapper, media is priced", since: 5, role: "tool", content: toolResult({ @@ -710,7 +772,9 @@ const fixtures: Fixture[] = [ { type: "media", data: B64, mediaType: "application/pdf" }, ], }), - text: ["Here is the chart"], + text: [ + '{"type":"content","value":[{"type":"text","text":"Here is the chart"}]}', + ], media: IMAGE + FILE, }, { @@ -739,14 +803,14 @@ const fixtures: Fixture[] = [ since: 6, role: "tool", content: toolResult({ type: "execution-denied", reason: "User said no" }), - text: ["User said no"], + text: ['{"type":"execution-denied","reason":"User said no"}'], }, { - name: "output execution-denied without a reason counts the default text providers send", + name: "output execution-denied without a reason", since: 6, role: "tool", content: toolResult({ type: "execution-denied" }), - text: ["Tool call execution denied."], + text: ['{"type":"execution-denied"}'], }, { name: "output content: image-* items are images", @@ -760,6 +824,7 @@ const fixtures: Fixture[] = [ { type: "image-file-id", fileId: "file_1" }, ], }), + text: ['{"type":"content","value":[]}'], media: 3 * IMAGE, }, { @@ -776,6 +841,9 @@ const fixtures: Fixture[] = [ { type: "custom", providerOptions: {} }, ], }), + text: [ + '{"type":"content","value":[{"type":"custom","providerOptions":{}}]}', + ], media: 3 * FILE + IMAGE, }, { @@ -862,7 +930,7 @@ const fixtures: Fixture[] = [ media: IMAGE, }, { - name: "reasoning-file image", + name: "reasoning-file image is not counted", since: 7, role: "assistant", content: { @@ -870,10 +938,9 @@ const fixtures: Fixture[] = [ mediaType: "image/png", data: { type: "data", data: B64 }, }, - media: IMAGE, }, { - name: "reasoning-file other", + name: "reasoning-file other is not counted", since: 7, role: "assistant", content: { @@ -881,7 +948,6 @@ const fixtures: Fixture[] = [ mediaType: "application/octet-stream", data: B64, }, - media: FILE, }, { name: "custom part", @@ -918,7 +984,9 @@ const fixtures: Fixture[] = [ }, ], }), - text: ["inline"], + text: [ + '{"type":"content","value":[{"type":"file","mediaType":"text/plain","data":{"type":"text","text":"inline"}}]}', + ], media: IMAGE + 2 * FILE, }, ]; @@ -972,12 +1040,15 @@ describe("message parts", () => { }, ] as unknown as ModelMessage[], }).messages; + const serialized = + '{"type":"content","value":[{"type":"text","text":"ok"}]}'; expect(counted!.content[0]).toEqual({ type: "tool-result", total: - Math.round((enc("call_1") + enc("ok")) * tokens.contentMultiplier) + - IMAGE, - output: Math.round(enc("ok") * tokens.contentMultiplier) + IMAGE, + Math.round( + (enc("call_1") + enc(serialized)) * tokens.contentMultiplier + ) + IMAGE, + output: Math.round(enc(serialized) * tokens.contentMultiplier) + IMAGE, }); }); diff --git a/src/sdk.ts b/src/sdk.ts index 2004a72..547f9b5 100644 --- a/src/sdk.ts +++ b/src/sdk.ts @@ -294,8 +294,8 @@ function propertiesOf(schema: unknown): Record | undefined { } /** - * Follows a local `$ref` ("#", "#/definitions/X", "#/$defs/X") to its target in - * the root schema. zod 4 emits one for a schema registered with `.meta({ id })` + * Follows a local JSON pointer `$ref` ("#", "#/definitions/X", "#/$defs/X") to + * its target in the root schema. Anchors ("#name") are not resolved. zod 4 emits one for a schema registered with `.meta({ id })` * and pydantic-based MCP servers emit them for nested models. A `description` * next to the `$ref` wins over the target's. */ @@ -304,7 +304,7 @@ function resolveRef( root: unknown ): Record { const ref = schema.$ref; - if (typeof ref !== "string" || !ref.startsWith("#")) { + if (typeof ref !== "string" || (ref !== "#" && !ref.startsWith("#/"))) { return schema; } let target: unknown = root; @@ -357,6 +357,13 @@ interface SchemaWalk { root: unknown; /** `properties` objects being walked, so a recursive schema ends. */ ancestors: Set; + /** + * Tokens of each `properties` object already walked. A schema whose + * definitions share definitions would otherwise be walked once per path. + */ + counted: Map; + /** Times a recursive branch was skipped, to tell a complete count. */ + cycleCuts: number; } /** The schema a property stands for once `$ref` and `null` wrappers are read through. */ @@ -377,6 +384,8 @@ function countToolSchema( config, root: jsonSchema, ancestors: new Set(), + counted: new Map(), + cycleCuts: 0, }; return countSchemaProperties(resolveRef(jsonSchema, jsonSchema), walk); } @@ -388,10 +397,19 @@ function countToolSchema( function countSchemaProperties(schema: unknown, walk: SchemaWalk): number { const { tokenizer, config, ancestors } = walk; const properties = propertiesOf(schema); - if (!properties || ancestors.has(properties)) { + if (!properties) { + return 0; + } + const counted = walk.counted.get(properties); + if (counted !== undefined) { + return counted; + } + if (ancestors.has(properties)) { + walk.cycleCuts++; return 0; } ancestors.add(properties); + const cycleCutsBefore = walk.cycleCuts; let tokens = 0; let isFirstProp = true; @@ -438,6 +456,10 @@ function countSchemaProperties(schema: unknown, walk: SchemaWalk): number { } ancestors.delete(properties); + // A count that skipped a recursive branch depends on where the walk entered. + if (walk.cycleCuts === cycleCutsBefore) { + walk.counted.set(properties, tokens); + } return tokens; } @@ -489,14 +511,16 @@ function countMessageContentTokensDetailed( if (Array.isArray(content)) { // Parts are read as plain records: their shape differs between AI SDK - // v5, v6 and v7, and unknown parts are skipped. + // v5, v6 and v7, and unknown parts are skipped. That includes `reasoning` + // and `reasoning-file`: Anthropic's count_tokens gives the same total with + // and without the signed thinking block of an earlier assistant turn. for (const part of content as unknown[]) { if (!isRecord(part)) { continue; } const tally = createTally(tokenizer); - if (part.type === "text" || part.type === "reasoning") { + if (part.type === "text") { if (typeof part.text === "string") { tally.addText(part.text); } @@ -589,8 +613,7 @@ function countMediaTokens(media: Media): number { /** * Message parts and tool-result content items that carry an image or a file, * with the kind to use when the media type does not say. The names are from - * AI SDK v5 (`media`), v6 (`image-*`, `file-*`) and v7 (`file`, - * `reasoning-file`). + * AI SDK v5 (`media`), v6 (`image-*`, `file-*`) and v7 (`file`). */ const MEDIA_KINDS = new Map([ ["image", "image"], @@ -599,7 +622,6 @@ const MEDIA_KINDS = new Map([ ["image-file-id", "image"], ["image-file-reference", "image"], ["file", "file"], - ["reasoning-file", "file"], ["media", "file"], ["file-data", "file"], ["file-url", "file"], @@ -607,39 +629,61 @@ const MEDIA_KINDS = new Map([ ["file-reference", "file"], ]); -function tallyMedia(tally: Tally, part: Record): void { - const data = - part.data ?? - part.image ?? - part.url ?? - part.fileId ?? - part.providerReference; - - // AI SDK v7 inline text documents are text, not a binary file. - if (isRecord(data) && data.type === "text" && typeof data.text === "string") { - tally.addText(data.text); - return; - } +function mediaData(part: Record): unknown { + return ( + part.data ?? part.image ?? part.url ?? part.fileId ?? part.providerReference + ); +} + +/** AI SDK v7 inline text documents are text, not a binary file. */ +function inlineText(part: Record): string | undefined { + const data = mediaData(part); + return isRecord(data) && data.type === "text" && typeof data.text === "string" + ? data.text + : undefined; +} +function isBinaryMediaPart(part: Record): boolean { + return ( + typeof part.type === "string" && + MEDIA_KINDS.has(part.type) && + inlineText(part) === undefined + ); +} + +function toMedia(part: Record): Media { const mediaType = typeof part.mediaType === "string" ? part.mediaType : undefined; const isImageMediaType = mediaType !== undefined && (mediaType === "image" || mediaType.startsWith("image/")); - tally.addMedia({ + return { kind: isImageMediaType || MEDIA_KINDS.get(part.type as string) === "image" ? "image" : "file", mediaType, - data, - }); + data: mediaData(part), + }; +} + +function tallyMedia(tally: Tally, part: Record): void { + const text = inlineText(part); + if (text !== undefined) { + tally.addText(text); + } else { + tally.addMedia(toMedia(part)); + } } /** - * Count the payload the provider sees for a tool result, not the - * `{ type, value }` wrapper around it. Shapes we do not know are counted as - * JSON, as they were before. + * A tool result counts as `JSON.stringify(output)`, the `{ type, value }` + * wrapper included, as it always has. The wrapper stands in for the + * `tool_result` block overhead that `ModelTokens` has no field for yet, and + * counting only the payload under-counts every result by about 6 tokens. + * + * The one exception is a `content` output: its images and files are priced by + * `countMediaTokens` instead of serialising their base64 or URL. */ function tallyToolOutput(tally: Tally, output: unknown): void { if (!output) { @@ -649,51 +693,21 @@ function tallyToolOutput(tally: Tally, output: unknown): void { tally.addText(output); return; } - if (isRecord(output)) { - switch (output.type) { - case "text": - case "error-text": - if (typeof output.value === "string") { - tally.addText(output.value); - return; - } - break; - case "json": - case "error-json": - if (output.value !== undefined) { - tally.addText(JSON.stringify(output.value)); - } - return; - case "execution-denied": - // OpenAI and Anthropic send this text when there is no reason. - tally.addText( - typeof output.reason === "string" - ? output.reason - : "Tool call execution denied." - ); - return; - case "content": - if (Array.isArray(output.value)) { - for (const item of output.value as unknown[]) { - tallyToolContentItem(tally, item); - } - return; - } - break; + if ( + isRecord(output) && + output.type === "content" && + Array.isArray(output.value) + ) { + const rest: unknown[] = []; + for (const item of output.value as unknown[]) { + if (isRecord(item) && isBinaryMediaPart(item)) { + tally.addMedia(toMedia(item)); + } else { + rest.push(item); + } } - } - tally.addText(JSON.stringify(output)); -} - -function tallyToolContentItem(tally: Tally, item: unknown): void { - if (!isRecord(item)) { + tally.addText(JSON.stringify({ ...output, value: rest })); return; } - if (item.type === "text" && typeof item.text === "string") { - tally.addText(item.text); - } else if (typeof item.type === "string" && MEDIA_KINDS.has(item.type)) { - tallyMedia(tally, item); - } else if (item.type !== "custom") { - tally.addText(JSON.stringify(item)); - } + tally.addText(JSON.stringify(output)); } From 4a71ec6e573b9595548330dfefd79e34579374c9 Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 3 Nov 2026 12:35:13 +0000 Subject: [PATCH 6/8] fix(sdk): bound recursive schema walks, list unconvertible schemas, count $ref siblings - Count each definition once per tool, including definitions inside a cycle (the back edge counts 0 where the walk first reaches it). A clique of 10 mutually recursive definitions took 21.4 s and a depth-20 diamond with a back edge to the root took 7.9 s. Both now take about 1 ms. - A schema the AI SDK converts to nothing (arktype on ai v5) is listed in tools.unresolved with reason error instead of counting 0 silently. - Keywords next to a $ref (enum, description) win over the target's, so they count. - README: list what the walker does not follow. Change-Id: Ibafb90075be8586bb0c2e350ef4d484c8b89ce15 Signed-off-by: Thomas Kosiewski --- README.md | 7 ++++ src/sdk.compat.test.ts | 88 +++++++++++++++++++++++++++++++++++++++--- src/sdk.ts | 27 +++++++------ 3 files changed, 103 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index fffd252..e682431 100644 --- a/README.md +++ b/README.md @@ -96,6 +96,13 @@ for (const { tool, part, reason } of result.tools.unresolved) { // Fix: resolve the schema first, e.g. jsonSchema(await loadSchema()). ``` +Tool schemas are counted for properties, descriptions, enums, nested objects and +arrays of objects. Local `$ref`s (`#`, `#/definitions/X`, `#/$defs/X`) are +followed, and a definition inside a recursive schema counts once. These are not +walked: unions with two or more non-null variants (`anyOf`, `oneOf`), `allOf`, +`$ref`s that point outside the schema or use an anchor (`#name`), and `$ref` +alias chains longer than one hop. + Estimate tokens using an encoding: ```ts diff --git a/src/sdk.compat.test.ts b/src/sdk.compat.test.ts index e8c8c9f..d2e8f68 100644 --- a/src/sdk.compat.test.ts +++ b/src/sdk.compat.test.ts @@ -407,6 +407,16 @@ describe("tool schemas", () => { } }); + test("a schema the AI SDK converts to nothing (arktype on ai v5)", () => { + const tools = countTools({ + t: { inputSchema: jsonSchema(undefined as never) }, + }); + expect(tools.definitions.t!.inputSchema).toBe(0); + expect(tools.unresolved).toMatchObject([ + { tool: "t", part: "inputSchema", reason: "error" }, + ]); + }); + test("a thenable whose then() throws", () => { const thenable = { then() { @@ -524,7 +534,7 @@ describe("tool schemas", () => { ); }); - test("definitions that reference each other count from where the walk enters", () => { + test("a definition in a cycle counts once, with its back edge cut", () => { const first = tokens.perFirstProp + tokens.perNestedObject; const next = tokens.perAdditionalProp + tokens.perNestedObject; const mutual = { @@ -535,14 +545,62 @@ describe("tool schemas", () => { B: { type: "object", properties: { a: { $ref: "#/$defs/A" } } }, }, }; - // Each branch stops where it meets a definition it is already inside. - const viaA = enc("b") + first + (enc("a") + first); - const viaB = enc("a") + first + (enc("b") + first); + // A is walked first and reaches B, whose edge back to A is cut. The + // root's own `b` reuses that count of B. + const b = enc("a") + first; + const a = enc("b") + first + b; expect(inline(jsonSchema(mutual as never))).toBe( - enc("a") + first + viaA + (enc("b") + next + viaB) + enc("a") + first + a + (enc("b") + next + b) ); }); + test("definitions that all reference each other are bounded", () => { + const size = 10; + const $defs: Record = {}; + for (let i = 0; i < size; i++) { + const properties: Record = { v: { type: "string" } }; + for (let j = 0; j < size; j++) { + if (j !== i) properties[`k${j}`] = { $ref: `#/$defs/N${j}` }; + } + $defs[`N${i}`] = { type: "object", properties }; + } + const clique = { + type: "object", + properties: { n: { $ref: "#/$defs/N0" } }, + $defs, + }; + + const start = performance.now(); + expect(inline(jsonSchema(clique as never))).toBeGreaterThan(0); + expect(performance.now() - start).toBeLessThan(500); + }); + + test("definitions that share definitions below a cycle are bounded", () => { + const depth = 20; + const $defs: Record = { + [`D${depth}`]: { + type: "object", + properties: { x: { type: "string" }, back: { $ref: "#" } }, + }, + }; + for (let level = depth - 1; level >= 0; level--) { + const next = { $ref: `#/$defs/D${level + 1}` }; + $defs[`D${level}`] = { + type: "object", + properties: { l: next, r: next }, + }; + } + const diamond = { + type: "object", + properties: { a: { $ref: "#/$defs/D0" } }, + $defs, + }; + + const start = performance.now(); + expect(inline(jsonSchema(diamond as never))).toBeGreaterThan(0); + expect(performance.now() - start).toBeLessThan(500); + }); + test("definitions that share definitions are walked once, not once per path", () => { const depth = 22; const $defs: Record = { @@ -576,6 +634,26 @@ describe("tool schemas", () => { expect(counted).toBe(expected); }); + test("keywords next to a $ref count with the target's", () => { + const defs = { U: { type: "string" } }; + const withSiblings = { + type: "object", + properties: { + u: { $ref: "#/$defs/U", enum: ["a", "b"], description: "Unit" }, + }, + $defs: defs, + }; + const flat = { + type: "object", + properties: { + u: { type: "string", enum: ["a", "b"], description: "Unit" }, + }, + }; + expect(inline(jsonSchema(withSiblings as never))).toBe( + inline(jsonSchema(flat as never)) + ); + }); + test("a $ref to somewhere else counts nothing", () => { const remote = { type: "object", diff --git a/src/sdk.ts b/src/sdk.ts index 547f9b5..14fe5be 100644 --- a/src/sdk.ts +++ b/src/sdk.ts @@ -269,6 +269,10 @@ function readInputSchema(inputSchema: unknown): InputSchemaRead { } return { status: "async" }; } + if (!isRecord(jsonSchema)) { + // For example an arktype schema on ai v5, which converts only zod. + throw new TypeError("asSchema() returned no JSON Schema object"); + } return { status: "ok", jsonSchema }; } catch (error) { return { status: "error", error }; @@ -321,9 +325,8 @@ function resolveRef( if (!isRecord(target)) { return schema; } - return schema.description === undefined - ? target - : { ...target, description: schema.description }; + const { $ref: _ref, ...siblings } = schema; + return { ...target, ...siblings }; } /** @@ -358,12 +361,12 @@ interface SchemaWalk { /** `properties` objects being walked, so a recursive schema ends. */ ancestors: Set; /** - * Tokens of each `properties` object already walked. A schema whose - * definitions share definitions would otherwise be walked once per path. + * Tokens of each `properties` object already walked. Every definition is + * walked once per tool, which bounds the walk by the size of the schema. + * Without it, definitions that share definitions would be walked once per + * path, and recursive ones once per cycle. */ counted: Map; - /** Times a recursive branch was skipped, to tell a complete count. */ - cycleCuts: number; } /** The schema a property stands for once `$ref` and `null` wrappers are read through. */ @@ -385,7 +388,6 @@ function countToolSchema( root: jsonSchema, ancestors: new Set(), counted: new Map(), - cycleCuts: 0, }; return countSchemaProperties(resolveRef(jsonSchema, jsonSchema), walk); } @@ -405,11 +407,9 @@ function countSchemaProperties(schema: unknown, walk: SchemaWalk): number { return counted; } if (ancestors.has(properties)) { - walk.cycleCuts++; return 0; } ancestors.add(properties); - const cycleCutsBefore = walk.cycleCuts; let tokens = 0; let isFirstProp = true; @@ -456,10 +456,9 @@ function countSchemaProperties(schema: unknown, walk: SchemaWalk): number { } ancestors.delete(properties); - // A count that skipped a recursive branch depends on where the walk entered. - if (walk.cycleCuts === cycleCutsBefore) { - walk.counted.set(properties, tokens); - } + // A definition inside a cycle is counted with its back edge cut where the + // walk first reached it, and that count is reused wherever it is referenced. + walk.counted.set(properties, tokens); return tokens; } From 3342e13d941b9d693b2df7cbc8ba502d9bd7ff54 Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 3 Nov 2026 13:19:42 +0000 Subject: [PATCH 7/8] fix(sdk): skip undefined $ref siblings, merge sibling properties, convert each input schema once - Keywords next to a $ref whose value is undefined no longer hide the target's. - A sibling `properties` is added to the target's properties and wins on a clash, instead of replacing them. The merged object is cached per pair, so a recursive $ref with sibling properties still ends. - Remember the JSON Schema read from each tool input schema (WeakMap on the input schema), so repeated count() calls do not convert it again. Only complete reads are kept. Change-Id: I04a1dd03a78cf1a867459c70d923baff917af5f8 Signed-off-by: Thomas Kosiewski --- src/sdk.compat.test.ts | 89 ++++++++++++++++++++++++++++++++++++++++++ src/sdk.ts | 58 ++++++++++++++++++++++++++- 2 files changed, 145 insertions(+), 2 deletions(-) diff --git a/src/sdk.compat.test.ts b/src/sdk.compat.test.ts index d2e8f68..6b1508b 100644 --- a/src/sdk.compat.test.ts +++ b/src/sdk.compat.test.ts @@ -335,6 +335,23 @@ describe("tool schemas", () => { expect(reads).toBe(1); }); + test("tools that are counted again are not converted again", () => { + let conversions = 0; + const lazy = () => { + conversions++; + return jsonSchema(sourceOfTruth as never); + }; + const tools = { t: { inputSchema: lazy } }; + const first = countTools(tools).definitions.t!.inputSchema; + expect(first).toBeGreaterThan(0); + expect(countTools(tools).definitions.t!.inputSchema).toBe(first); + expect( + countTools({ other: { inputSchema: lazy } }).definitions.other! + .inputSchema + ).toBe(first); + expect(conversions).toBe(1); + }); + const standard = (convert: unknown) => ({ "~standard": { version: 1, @@ -654,6 +671,78 @@ describe("tool schemas", () => { ); }); + test("a sibling keyword that is undefined does not hide the target's", () => { + const withUndefined = { + type: "object", + properties: { u: { $ref: "#/$defs/U", description: undefined } }, + $defs: { U: { type: "string", description: "Unit" } }, + }; + const flat = { + type: "object", + properties: { u: { type: "string", description: "Unit" } }, + }; + expect(inline(jsonSchema(withUndefined as never))).toBe( + inline(jsonSchema(flat as never)) + ); + }); + + test("sibling properties are added to the target's, and win on a clash", () => { + const withSiblings = { + type: "object", + properties: { + base: { + $ref: "#/$defs/Base", + properties: { + name: { type: "string", description: "Display name" }, + extra: { type: "number" }, + }, + }, + }, + $defs: { + Base: { + type: "object", + properties: { + id: { type: "string", description: "Id" }, + name: { type: "string" }, + }, + }, + }, + }; + const flat = { + type: "object", + properties: { + base: { + type: "object", + properties: { + id: { type: "string", description: "Id" }, + name: { type: "string", description: "Display name" }, + extra: { type: "number" }, + }, + }, + }, + }; + expect(inline(jsonSchema(withSiblings as never))).toBe( + inline(jsonSchema(flat as never)) + ); + }); + + test("a recursive $ref with sibling properties ends", () => { + const first = tokens.perFirstProp + tokens.perNestedObject; + const tree = { + type: "object", + properties: { + child: { $ref: "#", properties: { note: { type: "string" } } }, + }, + }; + // `child` resolves to the root plus `note`. Inside that, `child` points + // back at the same merged object, so it stops there. + const merged = + enc("child") + first + (enc("note") + tokens.perAdditionalProp); + expect(inline(jsonSchema(tree as never))).toBe( + enc("child") + first + merged + ); + }); + test("a $ref to somewhere else counts nothing", () => { const remote = { type: "object", diff --git a/src/sdk.ts b/src/sdk.ts index 14fe5be..076dcd6 100644 --- a/src/sdk.ts +++ b/src/sdk.ts @@ -249,6 +249,19 @@ type InputSchemaRead = | { status: "async" } | { status: "error"; error: unknown }; +/** + * JSON Schema read from each tool input schema, so repeated `count()` calls on + * the same tools do not convert them again. Only complete reads are kept: the + * conversion depends on nothing but the input schema. + */ +const convertedSchemas = new WeakMap(); + +function isCacheKey(value: unknown): value is object { + return ( + (typeof value === "object" && value !== null) || typeof value === "function" + ); +} + /** * Reads the JSON Schema the provider receives for a tool, whichever schema * library defined it. `asSchema` is the AI SDK's own conversion, so zod 3, @@ -256,6 +269,10 @@ type InputSchemaRead = * JSON Schema. */ function readInputSchema(inputSchema: unknown): InputSchemaRead { + const cacheKey = isCacheKey(inputSchema) ? inputSchema : undefined; + if (cacheKey && convertedSchemas.has(cacheKey)) { + return { status: "ok", jsonSchema: convertedSchemas.get(cacheKey) }; + } try { const jsonSchema: unknown = asSchema( inputSchema as Parameters[0] @@ -273,6 +290,9 @@ function readInputSchema(inputSchema: unknown): InputSchemaRead { // For example an arktype schema on ai v5, which converts only zod. throw new TypeError("asSchema() returned no JSON Schema object"); } + if (cacheKey) { + convertedSchemas.set(cacheKey, jsonSchema); + } return { status: "ok", jsonSchema }; } catch (error) { return { status: "error", error }; @@ -297,6 +317,33 @@ function propertiesOf(schema: unknown): Record | undefined { : undefined; } +/** + * The properties of a `$ref` target plus the properties written next to the + * `$ref`, which win. The same pair always gives the same object, because the + * walk tells a definition it is already inside by object identity. + */ +const mergedProperties = new WeakMap< + object, + WeakMap> +>(); + +function mergeProperties( + target: Record, + siblings: Record +): Record { + let bySiblings = mergedProperties.get(target); + if (!bySiblings) { + bySiblings = new WeakMap(); + mergedProperties.set(target, bySiblings); + } + let merged = bySiblings.get(siblings); + if (!merged) { + merged = { ...target, ...siblings }; + bySiblings.set(siblings, merged); + } + return merged; +} + /** * Follows a local JSON pointer `$ref` ("#", "#/definitions/X", "#/$defs/X") to * its target in the root schema. Anchors ("#name") are not resolved. zod 4 emits one for a schema registered with `.meta({ id })` @@ -325,8 +372,15 @@ function resolveRef( if (!isRecord(target)) { return schema; } - const { $ref: _ref, ...siblings } = schema; - return { ...target, ...siblings }; + const { $ref: _ref, ...rest } = schema; + const siblings = Object.fromEntries( + Object.entries(rest).filter(([, value]) => value !== undefined) + ); + const merged = { ...target, ...siblings }; + if (isRecord(target.properties) && isRecord(siblings.properties)) { + merged.properties = mergeProperties(target.properties, siblings.properties); + } + return merged; } /** From 3e7ceec58ec51237f72cf9f1d1968d7961998303 Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 3 Nov 2026 13:51:40 +0000 Subject: [PATCH 8/8] fix(sdk): count a file part as a file whatever its media type Main counts every file part as 100. The rewrite counted a file part with an image media type as 85, so { type: "file", data, mediaType: "image/png" } on openai/gpt-5 went from 110 to 95. Keep main's totals: file parts count 100, image parts count 85, until the image-counting change replaces countMediaTokens(). Change-Id: I505b3572eea34e9ac298df0743ce86bd54f5020d Signed-off-by: Thomas Kosiewski --- src/sdk.compat.test.ts | 43 ++++++++++++++++++++++++++++++------------ src/sdk.ts | 5 ++++- 2 files changed, 35 insertions(+), 13 deletions(-) diff --git a/src/sdk.compat.test.ts b/src/sdk.compat.test.ts index 6b1508b..a9e57f9 100644 --- a/src/sdk.compat.test.ts +++ b/src/sdk.compat.test.ts @@ -13,7 +13,7 @@ import { jsonSchema, type ModelMessage, type ToolSet } from "ai"; import * as z3 from "zod/v3"; import * as z4 from "zod/v4"; import { count } from "./sdk"; -import { Tokenizer, type Model, type ModelTokens } from "./index"; +import { models, Tokenizer, type Model, type ModelTokens } from "./index"; import * as o200k from "./encoding/o200k_base"; const aiMajor = Number( @@ -873,11 +873,11 @@ const fixtures: Fixture[] = [ media: FILE, }, { - name: "file part with an image media type", + name: "file part with an image media type counts as a file, as before", since: 5, role: "user", content: { type: "file", data: B64, mediaType: "image/png" }, - media: IMAGE, + media: FILE, }, { name: "assistant reasoning is not counted", @@ -886,11 +886,11 @@ const fixtures: Fixture[] = [ content: { type: "reasoning", text: "Let me think." }, }, { - name: "assistant generated image", + name: "assistant generated image counts as a file, as before", since: 5, role: "assistant", content: { type: "file", data: B64, mediaType: "image/webp" }, - media: IMAGE, + media: FILE, }, { name: "output text counts its wrapper", @@ -1035,7 +1035,7 @@ const fixtures: Fixture[] = [ }, { - name: "file with media type 'image' and URL data", + name: "file with media type 'image' counts as a file, as before", since: 7, role: "user", content: { @@ -1043,10 +1043,10 @@ const fixtures: Fixture[] = [ mediaType: "image", data: { type: "url", url: new URL("https://example.com/a.png") }, }, - media: IMAGE, + media: FILE, }, { - name: "file with image/* media type and tagged data", + name: "file with image/* media type and tagged data counts as a file", since: 7, role: "user", content: { @@ -1054,10 +1054,10 @@ const fixtures: Fixture[] = [ mediaType: "image/png", data: { type: "data", data: B64 }, }, - media: IMAGE, + media: FILE, }, { - name: "file with image media type and a provider reference", + name: "file with image media type and a provider reference counts as a file", since: 7, role: "user", content: { @@ -1065,7 +1065,7 @@ const fixtures: Fixture[] = [ mediaType: "image", data: { type: "reference", reference: { openai: "file_1" } }, }, - media: IMAGE, + media: FILE, }, { name: "file pdf with tagged data", @@ -1154,7 +1154,7 @@ const fixtures: Fixture[] = [ text: [ '{"type":"content","value":[{"type":"file","mediaType":"text/plain","data":{"type":"text","text":"inline"}}]}', ], - media: IMAGE + 2 * FILE, + media: 3 * FILE, }, ]; @@ -1188,6 +1188,25 @@ describe("message parts", () => { } ); + test("a file part with an image media type counts 100 and an image part 85, as on main", () => { + const gpt5 = models["openai/gpt-5"] as Model; + const overhead = + gpt5.tokens.baseOverhead + gpt5.tokens.perMessage + enc("user"); + const total = (content: unknown) => + count({ + tokenizer, + model: gpt5, + messages: [{ role: "user", content: [content] }] as ModelMessage[], + }).total; + + expect( + total({ type: "file", data: "iVBORw0KGgo=", mediaType: "image/png" }) + ).toBe(overhead + FILE); + expect( + total({ type: "image", image: "iVBORw0KGgo=", mediaType: "image/png" }) + ).toBe(overhead + IMAGE); + }); + test("a tool result total adds the call id to the output", () => { const [counted] = count({ tokenizer, diff --git a/src/sdk.ts b/src/sdk.ts index 076dcd6..e1992b4 100644 --- a/src/sdk.ts +++ b/src/sdk.ts @@ -711,8 +711,11 @@ function toMedia(part: Record): Media { mediaType !== undefined && (mediaType === "image" || mediaType.startsWith("image/")); return { + // A `file` part counts as a file whatever its media type, as it always + // has. The image-counting change will decide this from `mediaType`. kind: - isImageMediaType || MEDIA_KINDS.get(part.type as string) === "image" + part.type !== "file" && + (isImageMediaType || MEDIA_KINDS.get(part.type as string) === "image") ? "image" : "file", mediaType,