From a02d47ad2cd2f54308a0c34dde1c59c084fd7052 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Fri, 17 Jul 2026 10:02:41 +0200 Subject: [PATCH 1/4] feat(auto-routing): add routing constraints to mirror payload contract Adds RoutingConstraintsSchema (requiredInputModalities, promptTokensEstimate) as an optional field on MirrorPayloadSchema, plus two new exports for capability-aware routing: * detectRequiredInputModalities - walks OpenAI chat, Anthropic, and Responses bodies to produce a deduped, sorted modality set. * estimateRoutingTokens - routing-only token estimator that excludes media payloads (base64, URLs, file data) and counts tool-call payload strings; output-token reservation is added; result is a positive integer. The MirrorPayloadSchema stays non-strict so a newer gateway sending constraints does not break an older worker. All new symbols are re-exported from the package root so S2 (gateway) and S3 (worker) can import them without deep paths. --- .../src/contracts.test.ts | 166 ++++++++++++ packages/auto-routing-contracts/src/index.ts | 26 +- .../src/normalize.test.ts | 248 ++++++++++++++++- .../auto-routing-contracts/src/normalize.ts | 255 ++++++++++++++++++ 4 files changed, 693 insertions(+), 2 deletions(-) diff --git a/packages/auto-routing-contracts/src/contracts.test.ts b/packages/auto-routing-contracts/src/contracts.test.ts index 04038fbd2b..6dbae647d8 100644 --- a/packages/auto-routing-contracts/src/contracts.test.ts +++ b/packages/auto-routing-contracts/src/contracts.test.ts @@ -4,8 +4,12 @@ import { AutoRoutingClassifierModelResponseSchema, AutoRoutingDecisionResponseSchema, MirrorPayloadSchema, + RoutingConstraintsSchema, UpdateClassifierModelRequestSchema, + detectRequiredInputModalities, + estimateRoutingTokens, } from './index'; +import type { RoutingConstraints } from './index'; import { BenchmarkConfigSchema, DEFAULT_BENCHMARK_ORG_ID, @@ -318,3 +322,165 @@ describe('BenchmarkConfigSchema duplicate model ids', () => { expect(result.success).toBe(true); }); }); + +describe('RoutingConstraintsSchema', () => { + it('accepts a fully populated constraints object', () => { + const result = RoutingConstraintsSchema.parse({ + requiredInputModalities: ['image', 'file'], + promptTokensEstimate: 12345, + }); + expect(result).toEqual({ + requiredInputModalities: ['image', 'file'], + promptTokensEstimate: 12345, + }); + }); + + it('accepts an empty object (all fields optional)', () => { + expect(RoutingConstraintsSchema.parse({})).toEqual({}); + }); + + it('rejects non-positive promptTokensEstimate', () => { + expect(() => RoutingConstraintsSchema.parse({ promptTokensEstimate: 0 })).toThrow(); + expect(() => RoutingConstraintsSchema.parse({ promptTokensEstimate: -1 })).toThrow(); + expect(() => RoutingConstraintsSchema.parse({ promptTokensEstimate: 1.5 })).toThrow(); + }); + + it('rejects empty/whitespace modality strings', () => { + expect(() => + RoutingConstraintsSchema.parse({ requiredInputModalities: ['image', ''] }) + ).toThrow(); + expect(() => + RoutingConstraintsSchema.parse({ requiredInputModalities: ['image', ' '] }) + ).toThrow(); + }); +}); + +describe('MirrorPayloadSchema with routing constraints', () => { + const baseNormalized = { + apiKind: 'chat_completions' as const, + requestedModel: 'kilo-auto/free', + systemPromptPrefix: 'You are Kilo Code.', + userPromptPrefix: 'Add parser tests.', + latestUserPromptPrefix: null, + messageCount: 2, + hasTools: false, + stream: true, + providerHints: { provider: null, providerOptions: null }, + }; + + const basePayload = { + input: baseNormalized, + userId: 'user-1', + sessionId: 'session-123', + machineId: 'machine-1', + clientRequestId: 'req-1', + mode: 'code', + userAgent: 'Kilo-Code/4.106.0', + bodyBytes: 1234, + }; + + it('parses successfully without constraints (no behavior change)', () => { + const parsed = MirrorPayloadSchema.parse(basePayload); + expect(parsed.constraints).toBeUndefined(); + }); + + it('parses successfully with constraints present', () => { + const parsed = MirrorPayloadSchema.parse({ + ...basePayload, + constraints: { + requiredInputModalities: ['image'], + promptTokensEstimate: 8000, + }, + }); + expect(parsed.constraints).toEqual({ + requiredInputModalities: ['image'], + promptTokensEstimate: 8000, + }); + }); + + it('strips unknown keys for backward compatibility with older workers', () => { + const parsed = MirrorPayloadSchema.parse({ + ...basePayload, + futureFlag: true, + nestedFuture: { a: 1 }, + }); + expect((parsed as Record).futureFlag).toBeUndefined(); + expect((parsed as Record).nestedFuture).toBeUndefined(); + }); + + it('accepts a payload whose promptTokensEstimate equals the estimator output', () => { + // Realistic multi-message body mixing text, a remote image URL, and a + // large base64 image — the estimator must exclude both image payloads + // and produce a positive integer that satisfies the schema. + const body = { + model: 'gpt-4o', + messages: [ + { role: 'system', content: 'You are Kilo Code.' }, + { + role: 'user', + content: [ + { type: 'text', text: 'What do you see in this image?' }, + { + type: 'image_url', + image_url: { url: 'https://example.com/' + 'x'.repeat(500) + '.png' }, + }, + { + type: 'image_url', + image_url: { url: 'data:image/png;base64,' + 'A'.repeat(20_000) }, + }, + ], + }, + { + role: 'assistant', + tool_calls: [ + { + id: 'call-1', + type: 'function', + function: { name: 'read_file', arguments: JSON.stringify({ path: '/tmp/a.ts' }) }, + }, + ], + }, + { + role: 'tool', + content: JSON.stringify({ result: 'file contents ' + 'y'.repeat(800) }), + }, + ], + max_tokens: 2000, + }; + + const estimate = estimateRoutingTokens(body); + expect(Number.isInteger(estimate)).toBe(true); + expect(estimate).toBeGreaterThanOrEqual(1); + + // The mirror payload's constraints must accept this estimate verbatim. + const constraints: RoutingConstraints = { + requiredInputModalities: detectRequiredInputModalities(body), + promptTokensEstimate: estimate, + }; + + const parsed = MirrorPayloadSchema.parse({ + ...basePayload, + constraints, + }); + + expect(parsed.constraints?.promptTokensEstimate).toBe(estimate); + expect(parsed.constraints?.requiredInputModalities).toEqual(['image']); + }); +}); + +describe('package root re-exports', () => { + // These imports come from the package entry point (./index) — proves + // S2 (gateway) and S3 (worker) can import them from the package root + // without reaching into deep paths. + it('re-exports detectRequiredInputModalities from the package root', () => { + expect(typeof detectRequiredInputModalities).toBe('function'); + }); + + it('re-exports estimateRoutingTokens from the package root', () => { + expect(typeof estimateRoutingTokens).toBe('function'); + }); + + it('re-exports RoutingConstraintsSchema from the package root', () => { + expect(RoutingConstraintsSchema.safeParse({}).success).toBe(true); + }); +}); diff --git a/packages/auto-routing-contracts/src/index.ts b/packages/auto-routing-contracts/src/index.ts index 467e36bc94..5bac183e40 100644 --- a/packages/auto-routing-contracts/src/index.ts +++ b/packages/auto-routing-contracts/src/index.ts @@ -18,6 +18,19 @@ export const AutoRoutingModeSchema = z.enum(['cost_per_accuracy', 'best_accuracy export type AutoRoutingMode = z.infer; export const DEFAULT_AUTO_ROUTING_MODE: AutoRoutingMode = 'cost_per_accuracy'; +// Capability-aware routing hints attached to the mirrored request payload. +// Absent `constraints` (or absent fields within it) means "no extra +// constraint": the worker must not narrow the candidate set beyond what the +// table says. Modality strings are intentionally unconstrained at the +// contract boundary — the routing table owns the canonical vocabulary +// (`image`, `file`, `audio`, ...) and this schema only guarantees they are +// non-empty after trimming so a malformed caller cannot send whitespace. +export const RoutingConstraintsSchema = z.object({ + requiredInputModalities: z.array(z.string().trim().min(1)).optional(), + promptTokensEstimate: z.number().int().positive().optional(), +}); +export type RoutingConstraints = z.infer; + export function isVirtualAutoModelId(model: string): boolean { return model.trim().toLowerCase().startsWith('kilo-auto/'); } @@ -48,6 +61,11 @@ export const MirrorPayloadSchema = z.object({ // Size of the original request body, kept as an analytics dimension now // that the body itself is no longer mirrored. bodyBytes: z.number().int().nonnegative(), + // Capability-aware routing hints. Optional and intentionally absent-by- + // default so today's behavior is unchanged when the gateway does not yet + // supply them. The schema is non-strict (plain z.object above) so an old + // worker can safely ignore unknown keys added by a newer gateway. + constraints: RoutingConstraintsSchema.optional(), }); export type MirrorPayload = z.infer; @@ -199,7 +217,13 @@ export type AutoRoutingClassifierAnalyticsResponse = z.infer< typeof AutoRoutingClassifierAnalyticsResponseSchema >; -export { normalizeClassifierInput, redactProviderHints, type ClassifierApiKind } from './normalize'; +export { + normalizeClassifierInput, + redactProviderHints, + detectRequiredInputModalities, + estimateRoutingTokens, + type ClassifierApiKind, +} from './normalize'; export * from './reasoning'; export * from './taxonomy'; diff --git a/packages/auto-routing-contracts/src/normalize.test.ts b/packages/auto-routing-contracts/src/normalize.test.ts index 60a554883b..0df613a867 100644 --- a/packages/auto-routing-contracts/src/normalize.test.ts +++ b/packages/auto-routing-contracts/src/normalize.test.ts @@ -1,5 +1,10 @@ import { describe, expect, it } from 'vitest'; -import { normalizeClassifierInput, redactProviderHints } from './normalize'; +import { + detectRequiredInputModalities, + estimateRoutingTokens, + normalizeClassifierInput, + redactProviderHints, +} from './normalize'; describe('classifier input normalization', () => { it('captures the first and latest user prompt text for long chat completion sessions', () => { @@ -159,3 +164,244 @@ describe('classifier input normalization', () => { }); }); }); + +describe('detectRequiredInputModalities', () => { + it('returns [] for text-only bodies', () => { + expect( + detectRequiredInputModalities({ + model: 'gpt-4o', + messages: [ + { role: 'system', content: 'You are helpful.' }, + { role: 'user', content: 'Hello' }, + ], + }) + ).toEqual([]); + }); + + it('returns [] for malformed or non-object input without throwing', () => { + expect(detectRequiredInputModalities(null)).toEqual([]); + expect(detectRequiredInputModalities(undefined)).toEqual([]); + expect(detectRequiredInputModalities('garbage')).toEqual([]); + expect(detectRequiredInputModalities(42)).toEqual([]); + expect(detectRequiredInputModalities({})).toEqual([]); + }); + + it('detects image_url in OpenAI chat completions', () => { + expect( + detectRequiredInputModalities({ + model: 'gpt-4o', + messages: [ + { + role: 'user', + content: [ + { type: 'text', text: 'What is this?' }, + { type: 'image_url', image_url: { url: 'https://example.com/cat.png' } }, + ], + }, + ], + }) + ).toEqual(['image']); + }); + + it('detects image and document parts in Anthropic messages', () => { + expect( + detectRequiredInputModalities({ + model: 'claude-sonnet-4-20250514', + system: 'You analyze images.', + messages: [ + { + role: 'user', + content: [ + { type: 'image', source: { type: 'base64', data: 'aGVsbG8=' } }, + { type: 'text', text: 'Describe it.' }, + ], + }, + { + role: 'user', + content: [{ type: 'document', source: { type: 'base64', data: 'ZG9j' } }], + }, + ], + }) + ).toEqual(['file', 'image']); + }); + + it('detects input_image and input_file in OpenAI Responses API', () => { + expect( + detectRequiredInputModalities({ + model: 'gpt-4o', + input: [ + { role: 'user', content: 'Look at this' }, + { + role: 'user', + content: [ + { type: 'input_image', image_url: 'https://example.com/x.png' }, + { type: 'input_file', file_data: 'data:application/pdf;base64,AAA' }, + ], + }, + ], + }) + ).toEqual(['file', 'image']); + }); + + it('deduplicates repeated modalities across messages', () => { + expect( + detectRequiredInputModalities({ + model: 'gpt-4o', + messages: [ + { role: 'user', content: [{ type: 'image_url', image_url: { url: 'a' } }] }, + { role: 'user', content: [{ type: 'image_url', image_url: { url: 'b' } }] }, + ], + }) + ).toEqual(['image']); + }); +}); + +describe('estimateRoutingTokens', () => { + it('returns 0 for non-object input', () => { + expect(estimateRoutingTokens(null)).toBe(0); + expect(estimateRoutingTokens(undefined)).toBe(0); + expect(estimateRoutingTokens('not an object')).toBe(0); + }); + + it('estimates text-only chat completion bodies via chars/4', () => { + // "Hello world, this is a test." = 28 chars => 28/4 = 7 + expect( + estimateRoutingTokens({ + model: 'gpt-4o', + messages: [{ role: 'user', content: 'Hello world, this is a test.' }], + }) + ).toBe(7); + }); + + it('excludes remote image URL strings from the estimate', () => { + const longUrl = 'https://example.com/' + 'a'.repeat(2000) + '.png'; + const estimate = estimateRoutingTokens({ + model: 'gpt-4o', + messages: [ + { + role: 'user', + content: [ + { type: 'text', text: 'Describe this image.' }, + { type: 'image_url', image_url: { url: longUrl } }, + ], + }, + ], + }); + // Only "Describe this image." (20 chars) / 4 = 5 + expect(estimate).toBe(5); + }); + + it('excludes large base64 image payloads from the estimate', () => { + const base64 = 'B'.repeat(50_000); + const estimate = estimateRoutingTokens({ + model: 'gpt-4o', + messages: [ + { + role: 'user', + content: [ + { type: 'text', text: 'What?' }, + { type: 'image_url', image_url: { url: `data:image/png;base64,${base64}` } }, + ], + }, + ], + }); + expect(estimate).toBe(1); // "What?" is 5 chars => 5/4 = 1.25 => round = 1 + }); + + it('counts tool_calls function arguments as text', () => { + const args = JSON.stringify({ + path: '/tmp/foo.ts', + content: 'x'.repeat(2000), + }); + const estimate = estimateRoutingTokens({ + model: 'gpt-4o', + messages: [ + { role: 'user', content: 'Edit file' }, + { + role: 'assistant', + tool_calls: [ + { + id: 'call-1', + type: 'function', + function: { name: 'write_file', arguments: args }, + }, + ], + }, + ], + }); + // "Edit file" (9) + args length (2000+ chars) / 4 + const expected = Math.round((9 + args.length) / 4); + expect(estimate).toBe(expected); + }); + + it('counts Anthropic tool_use input serialized as text', () => { + const toolUse = { + type: 'tool_use', + id: 'tool-1', + name: 'read_file', + input: { path: '/tmp/foo.ts', offset: 'x'.repeat(2000) }, + }; + const estimate = estimateRoutingTokens({ + model: 'claude-sonnet-4-20250514', + messages: [ + { role: 'user', content: 'Read the file' }, + { role: 'assistant', content: [toolUse] }, + ], + }); + // "Read the file" (13) + JSON.stringify(toolUse.input) length / 4 + const inputJsonLen = JSON.stringify(toolUse.input).length; + const expected = Math.round((13 + inputJsonLen) / 4); + expect(estimate).toBe(expected); + }); + + it('counts Responses API function_call arguments as text', () => { + const args = 'x'.repeat(4000); + const estimate = estimateRoutingTokens({ + model: 'gpt-4o', + input: [{ type: 'function_call', name: 'do_thing', arguments: args }], + }); + expect(estimate).toBe(Math.round(args.length / 4)); + }); + + it('adds max_tokens to the estimate', () => { + const estimate = estimateRoutingTokens({ + model: 'gpt-4o', + messages: [{ role: 'user', content: 'hi' }], + max_tokens: 500, + }); + // "hi" = 2 chars => 0.5, + 500 => 500.5 => round = 501 + expect(estimate).toBe(501); + }); + + it('adds max_completion_tokens to the estimate', () => { + const estimate = estimateRoutingTokens({ + model: 'gpt-4o', + messages: [{ role: 'user', content: 'hi' }], + max_completion_tokens: 1000, + }); + expect(estimate).toBe(1001); + }); + + it('adds max_output_tokens to the estimate', () => { + const estimate = estimateRoutingTokens({ + model: 'gpt-4o', + messages: [{ role: 'user', content: 'hi' }], + max_output_tokens: 256, + }); + expect(estimate).toBe(257); + }); + + it('returns 0 for a completely empty body with no reservation', () => { + expect(estimateRoutingTokens({})).toBe(0); + }); + + it('returns a positive integer (never fractional, never 0 when text exists)', () => { + // 1 char / 4 = 0.25 => would round to 0, must floor to 1 + const estimate = estimateRoutingTokens({ + model: 'gpt-4o', + messages: [{ role: 'user', content: 'a' }], + }); + expect(Number.isInteger(estimate)).toBe(true); + expect(estimate).toBeGreaterThanOrEqual(1); + }); +}); diff --git a/packages/auto-routing-contracts/src/normalize.ts b/packages/auto-routing-contracts/src/normalize.ts index 8687d0c01c..10a0e72fb0 100644 --- a/packages/auto-routing-contracts/src/normalize.ts +++ b/packages/auto-routing-contracts/src/normalize.ts @@ -333,3 +333,258 @@ function isSensitiveKey(key: string) { function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null; } + +// ---------- Capability-aware routing helpers ---------- + +// Walks OpenAI chat completions, OpenAI Responses, and Anthropic Messages +// bodies looking for multimodal content parts. Returns the deduped, sorted +// set of capability tokens the request demands. This is independent of +// `normalizeClassifierInput` because (a) it inspects the raw gateway body +// before the request is mutated, and (b) it does not care about text — +// only which modalities a model must support. Malformed/unknown shapes +// are silently ignored; the caller treats an empty result as "text only". +export function detectRequiredInputModalities(body: unknown): string[] { + const found = new Set(); + collectModalities(body, found); + + if (found.size === 0) { + return []; + } + + return [...found].sort(); +} + +function collectModalities(value: unknown, out: Set): void { + if (value === null || value === undefined) return; + + if (Array.isArray(value)) { + for (const item of value) { + collectModalities(item, out); + } + return; + } + + if (!isRecord(value)) return; + + // Walk known container fields so message arrays and content parts are + // visited. Bodies use different keys per provider: + // * OpenAI chat completions / Anthropic: `messages[]` + // * OpenAI Responses: `input` (string or array) + // * Gemini-style: `contents[].parts[]` + // Content parts live under `content` (chat / Anthropic) or `parts` + // (Gemini). We do NOT recurse into every key — that would over-walk + // tools, metadata, and provider hints — only into the known structural + // containers. + for (const key of MODALITY_CONTAINER_KEYS) { + collectModalities(value[key], out); + } + + // Typed content parts: OpenAI Responses `input_image` / `input_file`, + // OpenAI chat `image_url`, Anthropic `image` / `document`. + const type = value.type; + if (typeof type === 'string') { + if (type === 'image_url' || type === 'image' || type === 'input_image') { + out.add('image'); + } else if (type === 'file' || type === 'input_file' || type === 'document') { + out.add('file'); + } + } + + // Guards against callers that omit the `type` discriminator — the + // presence of a known media field is itself sufficient signal. + if ('image_url' in value || 'input_image' in value) { + out.add('image'); + } + if ('input_file' in value || 'file' in value) { + out.add('file'); + } +} + +const MODALITY_CONTAINER_KEYS = ['messages', 'input', 'content', 'parts', 'system']; + +// Routing-only token estimator. Deliberately distinct from the gateway's +// per-request `estimateTokenCount`, which mis-counts base64 image payloads +// as text. This estimator is a *capability gate* — its sole job is to +// produce a positive-integer hint that the worker can compare against +// model context limits. It must: +// * include only textual content (plain `content` strings, `text` / +// `input_text` part text, system strings, and tool-call payload +// strings — which dominate agentic traffic), +// * exclude media payload strings (image URLs, base64/data URLs, +// file/document data) regardless of their length, +// * add the body's output-token reservation when present, +// * never return 0 when there is any text at all, and +// * never return a fractional value (downstream schema is int). +export function estimateRoutingTokens(body: unknown): number { + if (!isRecord(body)) return 0; + + const textChars = sumBodyTextChars(body); + const reservation = readOutputReservation(body); + const raw = textChars / 4 + reservation; + + const rounded = Math.round(raw); + if (rounded <= 0 && raw > 0) return 1; + return rounded; +} + +// Walks known container structures (messages, input, system, instructions) +// and extracts text from content parts. Does NOT recurse into arbitrary +// object fields — that would count `model`, `role`, tool definitions, etc. +function sumBodyTextChars(body: Record): number { + let total = 0; + + // OpenAI chat completions / Anthropic messages: messages[] + if (Array.isArray(body.messages)) { + for (const msg of body.messages) { + total += sumMessageTextChars(msg); + } + } + + // Anthropic: top-level system field (string or parts array) + const system = body.system; + if (typeof system === 'string') { + total += system.length; + } else if (Array.isArray(system)) { + for (const part of system) { + total += sumContentPartTextChars(part); + } + } + + // Responses API: instructions (string) and input (string or array) + if (typeof body.instructions === 'string') { + total += body.instructions.length; + } + if ('input' in body) { + total += sumResponsesInputChars(body.input); + } + + return total; +} + +function sumMessageTextChars(msg: unknown): number { + if (!isRecord(msg)) return 0; + + let total = 0; + + // Content: string, parts array, or part object + const content = msg.content; + if (typeof content === 'string') { + total += content.length; + } else if (Array.isArray(content)) { + for (const part of content) { + total += sumContentPartTextChars(part); + } + } else if (isRecord(content)) { + total += sumContentPartTextChars(content); + } + + // OpenAI assistant tool_calls: count function.arguments strings + if (Array.isArray(msg.tool_calls)) { + for (const tc of msg.tool_calls) { + if (isRecord(tc) && isRecord(tc.function) && typeof tc.function.arguments === 'string') { + total += tc.function.arguments.length; + } + } + } + + return total; +} + +function sumContentPartTextChars(part: unknown): number { + if (typeof part === 'string') return part.length; + if (!isRecord(part)) return 0; + + const type = part.type; + + // Media parts: zero contribution regardless of payload size + if (typeof type === 'string') { + if ( + type === 'image_url' || + type === 'image' || + type === 'input_image' || + type === 'file' || + type === 'input_file' || + type === 'document' + ) { + return 0; + } + + // Text parts + if (type === 'text' || type === 'input_text') { + return typeof part.text === 'string' ? part.text.length : 0; + } + + // Anthropic tool_use: count serialized input (can dominate agentic traffic) + if (type === 'tool_use') { + if (part.input !== undefined && part.input !== null) { + return JSON.stringify(part.input).length; + } + return 0; + } + + // Tool/function result content strings + if (type === 'tool_result' || type === 'function_call_output' || type === 'tool_call_output') { + const content = part.content; + if (typeof content === 'string') return content.length; + if (Array.isArray(content)) { + return content.reduce((sum: number, p: unknown) => sum + sumContentPartTextChars(p), 0); + } + return 0; + } + } + + // Untyped part: try to extract text from common fields + if (typeof part.text === 'string') return part.text.length; + if (typeof part.content === 'string') return part.content.length; + if (Array.isArray(part.content)) { + return part.content.reduce((sum: number, p: unknown) => sum + sumContentPartTextChars(p), 0); + } + + return 0; +} + +// Responses API input: string, array of messages/parts, or object +function sumResponsesInputChars(input: unknown): number { + if (typeof input === 'string') return input.length; + if (Array.isArray(input)) { + let total = 0; + for (const item of input) { + if (typeof item === 'string') { + total += item.length; + } else if (isRecord(item)) { + const type = item.type; + // Responses API typed parts + if ( + typeof type === 'string' && + (type === 'function_call_output' || type === 'tool_call_output') + ) { + total += sumContentPartTextChars(item); + } else if (typeof type === 'string' && type === 'function_call') { + // Responses API function_call: count arguments string + if (typeof item.arguments === 'string') { + total += item.arguments.length; + } + } else { + // Message-like object with role and content + total += sumMessageTextChars(item); + } + } + } + return total; + } + if (isRecord(input)) { + // Single message-like object + return sumMessageTextChars(input); + } + return 0; +} + +function readOutputReservation(body: Record): number { + const candidates = [body.max_tokens, body.max_completion_tokens, body.max_output_tokens]; + for (const candidate of candidates) { + if (typeof candidate === 'number' && Number.isFinite(candidate) && candidate > 0) { + return candidate; + } + } + return 0; +} From cb8d7c5c841ff624327618608af053641e6e2708 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Fri, 17 Jul 2026 10:49:55 +0200 Subject: [PATCH 2/4] feat(ai-gateway): send modality and context-size facts to the efficient decider buildDecidePayload now computes constraints.requiredInputModalities and constraints.promptTokensEstimate from the original request body (before applyResolvedAutoModel mutates it) using the S1 contract helpers. Fields are omitted when empty/zero so an old-shaped, no-media, no-estimate body produces a byte-identical payload to today. Adds an invariant comment on BALANCED_QWEN_MODEL recording that the efficient static fallback must remain image-capable, since the worker-side capability filter relies on that guarantee when no capable benchmark candidate exists. --- .../src/lib/ai-gateway/auto-model/index.ts | 6 ++ .../ai-gateway/auto-routing-decision.test.ts | 75 +++++++++++++++++++ .../lib/ai-gateway/auto-routing-decision.ts | 19 +++++ 3 files changed, 100 insertions(+) diff --git a/apps/web/src/lib/ai-gateway/auto-model/index.ts b/apps/web/src/lib/ai-gateway/auto-model/index.ts index b5005f340c..a34c852eb1 100644 --- a/apps/web/src/lib/ai-gateway/auto-model/index.ts +++ b/apps/web/src/lib/ai-gateway/auto-model/index.ts @@ -87,6 +87,12 @@ export const BALANCED_CLAW_SETUP_MODEL: ResolvedAutoModel = { verbosity: 'high', }; +// INVARIANT: the efficient static fallback must remain image-capable. +// The capability-aware routing filter relies on this guarantee to make +// image requests succeed even when no benchmark candidate is capable. +// Whoever changes this model constant must re-verify image support +// (via live OpenRouter data or the `model_stats` table) before +// swapping it — do not assume parity with the prior value. export const BALANCED_QWEN_MODEL: ResolvedAutoModel = { model: QWEN37_PLUS_MODEL_ID, reasoning: { enabled: true }, diff --git a/apps/web/src/lib/ai-gateway/auto-routing-decision.test.ts b/apps/web/src/lib/ai-gateway/auto-routing-decision.test.ts index e438b6c680..97ef7c90fa 100644 --- a/apps/web/src/lib/ai-gateway/auto-routing-decision.test.ts +++ b/apps/web/src/lib/ai-gateway/auto-routing-decision.test.ts @@ -13,6 +13,10 @@ jest.mock('@/lib/utils.server', () => ({ import { fetchEfficientAutoDecision } from './auto-routing-decision'; import type { EfficientDecisionParams } from './auto-routing-decision'; +import { + detectRequiredInputModalities, + estimateRoutingTokens, +} from '@kilocode/auto-routing-contracts'; const originalFetch = globalThis.fetch; const mockedFetch = jest.fn() as jest.MockedFunction; @@ -191,4 +195,75 @@ describe('fetchEfficientAutoDecision', () => { expect(result).toEqual({ decision: null, costUsd: 0.001 }); }); + + it('forwards requiredInputModalities=["image"] into constraints for an image-bearing body', async () => { + mockedFetch.mockResolvedValueOnce(new Response(JSON.stringify(validResponse), { status: 200 })); + + const imageBody = { + model: 'kilo-auto/efficient', + stream: true, + messages: [ + { + role: 'user', + content: [ + { type: 'text', text: 'What is in this picture?' }, + { + type: 'image_url', + image_url: { url: 'https://example.com/cat.png' }, + }, + ], + }, + ], + }; + + await fetchEfficientAutoDecision({ ...makeParams(), body: imageBody }, options); + + const [, init] = mockedFetch.mock.calls[0]; + const body = JSON.parse(init?.body as string); + expect(body.constraints?.requiredInputModalities).toEqual(['image']); + expect(detectRequiredInputModalities(imageBody)).toEqual(['image']); + }); + + it('omits requiredInputModalities from constraints (and the whole constraints key when no other field is present) for a text-only body', async () => { + mockedFetch.mockResolvedValueOnce(new Response(JSON.stringify(validResponse), { status: 200 })); + + await fetchEfficientAutoDecision(makeParams(), options); + + const [, init] = mockedFetch.mock.calls[0]; + const body = JSON.parse(init?.body as string); + // The default `makeParams()` body is short text that may or may not + // produce a positive token estimate. We only assert the modality + // contract here: the key must never be `[]` (i.e. empty), and when + // there's no prompt-token estimate either the whole `constraints` + // key must be absent. + if ('constraints' in body) { + expect(body.constraints.requiredInputModalities).toBeUndefined(); + } else { + expect(body.constraints).toBeUndefined(); + } + }); + + it('forwards the promptTokensEstimate returned by estimateRoutingTokens unchanged', async () => { + mockedFetch.mockResolvedValueOnce(new Response(JSON.stringify(validResponse), { status: 200 })); + + const longText = + 'Please summarize this long document. ' + + 'The quick brown fox jumps over the lazy dog. '.repeat(200); + const body = { + model: 'kilo-auto/efficient', + stream: true, + messages: [ + { role: 'system', content: 'You are Kilo Code.' }, + { role: 'user', content: longText }, + ], + }; + + await fetchEfficientAutoDecision({ ...makeParams(), body }, options); + + const [, init] = mockedFetch.mock.calls[0]; + const sent = JSON.parse(init?.body as string); + const expected = estimateRoutingTokens(body); + expect(expected).toBeGreaterThan(0); + expect(sent.constraints?.promptTokensEstimate).toBe(expected); + }); }); diff --git a/apps/web/src/lib/ai-gateway/auto-routing-decision.ts b/apps/web/src/lib/ai-gateway/auto-routing-decision.ts index 427ef9d61e..842e89b411 100644 --- a/apps/web/src/lib/ai-gateway/auto-routing-decision.ts +++ b/apps/web/src/lib/ai-gateway/auto-routing-decision.ts @@ -1,5 +1,7 @@ import { AutoRoutingDecisionResponseSchema, + detectRequiredInputModalities, + estimateRoutingTokens, type AutoRoutingDecision, normalizeClassifierInput, } from '@kilocode/auto-routing-contracts'; @@ -39,11 +41,28 @@ function buildDecidePayload(params: EfficientDecisionParams): MirrorPayload | nu }); if (!normalizedInput) return null; + // Compute capability-aware routing hints from the original body (the + // caller mutates it after this thunk runs, so the full body is only + // available here). Omit each field when it carries no information, and + // omit `constraints` entirely when both would be absent, so today's + // payload shape is preserved byte-for-byte for text-only, sub-token + // requests. + const requiredInputModalities = detectRequiredInputModalities(params.body); + const promptTokensEstimate = estimateRoutingTokens(params.body); + const constraints: MirrorPayload['constraints'] = + requiredInputModalities.length > 0 || promptTokensEstimate > 0 + ? { + ...(requiredInputModalities.length > 0 ? { requiredInputModalities } : {}), + ...(promptTokensEstimate > 0 ? { promptTokensEstimate } : {}), + } + : undefined; + return { input: normalizedInput, ...(params.deniedModelIds?.length ? { routingPolicy: { deniedModelIds: [...params.deniedModelIds] } } : {}), + ...(constraints ? { constraints } : {}), userId: params.userId, organizationId: params.organizationId, sessionId: params.sessionId, From 8839ff1e8a45de7ca7b331f2ddddc6ebce12933e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Fri, 17 Jul 2026 10:50:07 +0200 Subject: [PATCH 3/4] feat(auto-routing): filter efficient candidates by modality and context window Adds services/auto-routing/src/model-capabilities.ts: a Hyperdrive-backed lookup of {input_modalities, context_length} from model_stats for the union of routing-table candidates and the coding-plan default model, cached with the same in-memory-TTL + KV-read-through pattern as the routing table (60s / 1h), with a 500ms sub-budget covering the entire lookup (routing-table fetch + cache/KV/DB) so a slow Hyperdrive can never consume the gateway's 2s /decide timeout or shift plain-text traffic to balanced. decision-engine.ts's computeDecision now accepts optional constraints and a capability map: an ENFORCED_MODALITIES constant (image always, file per OpenRouter's documented input-modality vocabulary) gates which required modalities are actually enforced (others are ignored, never fail-closed); unknown/missing capability data fails a candidate only when an enforced modality is required. Context eligibility excludes only provably-too-small candidates (known context < estimate); unknown-context candidates keep their existing rank so plain-text routing does not regress on missing metadata; an all-too-small route degrades to the max-known-context candidates. Both filters apply before fresh and sticky selection, so a sticky incumbent that becomes incapable or too small is replaced. decide.ts threads constraints + the capability map through every computeDecision call site (cached-classification, fresh-classification, heuristic-fallback) only when payload.constraints is present; an absent constraints field takes the exact byte-identical path as today, with zero capability fetch anywhere including the coding-plan short-circuit. The short-circuit itself now checks the fixed coding-plan model against the same modality/context rules before taking it, falling through to normal benchmark routing (subscription-billed to credit-billed) when the model violates a constraint. --- services/auto-routing/src/decide.ts | 102 ++++- .../auto-routing/src/decision-engine.test.ts | 410 ++++++++++++++++++ services/auto-routing/src/decision-engine.ts | 115 ++++- services/auto-routing/src/index.test.ts | 304 ++++++++++++- .../src/model-capabilities.test.ts | 357 +++++++++++++++ .../auto-routing/src/model-capabilities.ts | 258 +++++++++++ 6 files changed, 1529 insertions(+), 17 deletions(-) create mode 100644 services/auto-routing/src/model-capabilities.test.ts create mode 100644 services/auto-routing/src/model-capabilities.ts diff --git a/services/auto-routing/src/decide.ts b/services/auto-routing/src/decide.ts index fa5222db8a..12fc4aeb9f 100644 --- a/services/auto-routing/src/decide.ts +++ b/services/auto-routing/src/decide.ts @@ -4,6 +4,7 @@ import type { AutoRoutingDecisionResponse, MirrorPayload, NormalizedClassifierInput, + RoutingConstraints, } from '@kilocode/auto-routing-contracts'; import { formatError } from '@kilocode/worker-utils'; import type { Handler } from 'hono'; @@ -24,13 +25,44 @@ import { putCachedClassification, putStickyDecision, } from './decision-cache'; -import { computeDecision } from './decision-engine'; +import { computeDecision, ENFORCED_MODALITIES } from './decision-engine'; import { ClassifierRunError, classifyNormalizedInput } from './model-classifier'; import type { ClassifierRunResult } from './model-classifier'; import { getRoutingTable } from './routing-table'; import { getAutoRoutingMode } from './routing-mode'; import type { HonoEnv } from './hono-env'; import { codingPlanDefaultDecision, getCodingPlanPreference } from './coding-plan-preference'; +import { getModelCapabilities } from './model-capabilities'; +import type { ModelCapabilities, ModelCapabilitiesMap } from './model-capabilities'; + +// Check whether the coding-plan default model satisfies a constrained +// request. Mirrors the same `fail-closed when required` policy used in +// `decision-engine.ts`: +// * Unknown capability metadata fails when a required+enforced modality +// is set (the model might or might not support it, so we do not trust +// the short-circuit). +// * Unknown context length is treated as "still fits" (consistent with +// the unknown-keeps-rank policy in the engine). +// * A known context that is provably smaller than the estimate fails +// the check. +function codingPlanSatisfiesConstraints( + caps: ModelCapabilities | undefined, + constraints: RoutingConstraints +): boolean { + const required = constraints.requiredInputModalities ?? []; + const enforcedAndRequired = required.filter(m => ENFORCED_MODALITIES.includes(m)); + if (enforcedAndRequired.length > 0) { + if (!caps) return false; + for (const modality of enforcedAndRequired) { + if (!caps.inputModalities.has(modality)) return false; + } + } + const estimate = constraints.promptTokensEstimate; + if (typeof estimate === 'number' && caps && typeof caps.contextLength === 'number') { + if (caps.contextLength < estimate) return false; + } + return true; +} // Isolate-scoped request counter, used to correlate latency with isolate // warm-up in logs. @@ -286,17 +318,53 @@ export const decideHandler: Handler = async c => { const startedAt = performance.now(); const deniedModelIds = new Set(payload.routingPolicy?.deniedModelIds ?? []); const codingPlanPreference = await getCodingPlanPreference(c.env, payload.userId); - if (codingPlanPreference.active && !deniedModelIds.has(codingPlanPreference.modelId)) { - const decision = codingPlanDefaultDecision(codingPlanPreference); - writeClassifierMetricsDataPoint(c.env, { - status: 'coding_plan_default', - classifierModel: 'coding_plan_default', - requestedModel: payload.input.requestedModel, - classifierDurationMs: performance.now() - startedAt, - classifierCostCredits: 0, - cacheHit: false, + const codingPlanActive = + codingPlanPreference.active && !deniedModelIds.has(codingPlanPreference.modelId); + // Narrow once: `constraints` is only non-undefined inside the branches + // that already checked `hasConstraints`. This avoids a `!` non-null + // assertion across the closure. + const hasConstraints = payload.constraints !== undefined; + const constraints: RoutingConstraints | undefined = payload.constraints; + + // Capability-aware path: when the gateway attached a `constraints` field, + // we must consult capability data before either (a) taking the coding- + // plan short-circuit or (b) making a benchmark decision. The lookup has + // its own 500ms sub-budget; on failure we treat it as "no capability + // data" and the decision-engine fails closed on required modalities. + // + // When `constraints` is absent we MUST stay byte-identical to today: no + // capability fetch, no routing-table fetch, no benchmark hop on the + // coding-plan path. + let capabilities: ModelCapabilitiesMap = new Map(); + if (hasConstraints && constraints) { + capabilities = await getModelCapabilities(c.env, { + codingPlanModelId: codingPlanActive ? codingPlanPreference.modelId : null, }); - return c.json({ cost: 0, decision, classifierResult: null }); + } + + if (codingPlanActive) { + const canTakeShortCircuit = + hasConstraints && constraints + ? codingPlanSatisfiesConstraints( + capabilities.get(codingPlanPreference.modelId), + constraints + ) + : true; + if (canTakeShortCircuit) { + const decision = codingPlanDefaultDecision(codingPlanPreference); + writeClassifierMetricsDataPoint(c.env, { + status: 'coding_plan_default', + classifierModel: 'coding_plan_default', + requestedModel: payload.input.requestedModel, + classifierDurationMs: performance.now() - startedAt, + classifierCostCredits: 0, + cacheHit: false, + }); + return c.json({ cost: 0, decision, classifierResult: null }); + } + // Fall through to the normal benchmark flow because the coding-plan + // model cannot satisfy the constrained request. This moves the request + // from subscription-billed to credit-billed benchmark routing. } const [hashes, userIdHash, classifierModel, successSampleRate, routingTable, routingMode] = @@ -332,7 +400,11 @@ export const decideHandler: Handler = async c => { routingTable, stickyModel, deniedModelIds, - routingMode + routingMode, + { + constraints: payload.constraints, + capabilityMap: hasConstraints ? capabilities : undefined, + } ); if (decision) { c.executionCtx.waitUntil(putStickyDecision(c.env, ctx.conversationKey, decision.model)); @@ -368,7 +440,11 @@ export const decideHandler: Handler = async c => { routingTable, stickyModel, deniedModelIds, - routingMode + routingMode, + { + constraints: payload.constraints, + capabilityMap: hasConstraints ? capabilities : undefined, + } ); // Like the classification cache, sticky state only trusts real classifier // output: a heuristic fallback must not re-anchor the session's model. diff --git a/services/auto-routing/src/decision-engine.test.ts b/services/auto-routing/src/decision-engine.test.ts index a1bc13e279..b81cb29522 100644 --- a/services/auto-routing/src/decision-engine.test.ts +++ b/services/auto-routing/src/decision-engine.test.ts @@ -1,6 +1,20 @@ import { describe, expect, it } from 'vitest'; import type { ClassifierOutput, RoutingTable } from '@kilocode/auto-routing-contracts'; import { computeDecision } from './decision-engine'; +import type { ModelCapabilities, ModelCapabilitiesMap } from './model-capabilities'; + +function makeCaps( + rows: Record +): ModelCapabilitiesMap { + const map = new Map(); + for (const [id, row] of Object.entries(rows)) { + map.set(id, { + inputModalities: new Set(row.inputModalities ?? []), + contextLength: row.contextLength ?? null, + }); + } + return map; +} const classification: ClassifierOutput = { taskType: 'implementation', @@ -285,4 +299,400 @@ describe('computeDecision', () => { expect(decision).toMatchObject({ model: 'cheap/chat', sticky: false }); }); }); + + describe('capability filters', () => { + const visionTable: RoutingTable = { + ...table, + routes: { + ...table.routes, + 'implementation/code_generation': [ + { + model: 'text-only/chat', + accuracy: 0.95, + avgCostUsd: 0.001, + meetsThreshold: true, + }, + { + model: 'vision/chat', + accuracy: 0.85, + avgCostUsd: 0.002, + meetsThreshold: true, + }, + { + model: 'premium-vision/chat', + accuracy: 0.92, + avgCostUsd: 0.005, + meetsThreshold: true, + }, + ], + }, + }; + + it('skips a non-vision top-ranked candidate when an image is required', () => { + const caps = makeCaps({ + 'text-only/chat': { inputModalities: [] }, + 'vision/chat': { inputModalities: ['image'] }, + 'premium-vision/chat': { inputModalities: ['image'] }, + }); + const decision = computeDecision( + classification, + visionTable, + null, + new Set(), + 'cost_per_accuracy', + { + constraints: { requiredInputModalities: ['image'] }, + capabilityMap: caps, + } + ); + expect(decision).toMatchObject({ model: 'vision/chat', sticky: false }); + }); + + it('accepts a candidate whose capability map lists the image modality (folding happens upstream)', () => { + // Synonym folding (`image_url` -> `image`) lives in + // `model-capabilities.ts` and is tested there; here the engine just + // sees an already-folded capability set and accepts the candidate. + const caps = makeCaps({ + 'text-only/chat': { inputModalities: [] }, + 'vision/chat': { inputModalities: ['image'] }, + }); + const decision = computeDecision( + classification, + visionTable, + null, + new Set(), + 'cost_per_accuracy', + { + constraints: { requiredInputModalities: ['image'] }, + capabilityMap: caps, + } + ); + expect(decision).toMatchObject({ model: 'vision/chat', sticky: false }); + }); + + it('ignores a required modality outside ENFORCED_MODALITIES instead of failing closed', () => { + // 'audio' is not in ENFORCED_MODALITIES; the modality filter is a + // no-op for it, so every candidate still passes the modality check. + const caps = makeCaps({ + 'text-only/chat': { inputModalities: [] }, + 'vision/chat': { inputModalities: ['image'] }, + }); + const decision = computeDecision( + classification, + visionTable, + null, + new Set(), + 'cost_per_accuracy', + { + constraints: { requiredInputModalities: ['audio'] }, + capabilityMap: caps, + } + ); + expect(decision).toMatchObject({ model: 'text-only/chat', sticky: false }); + }); + + it('fails closed when every candidate is missing the required image modality', () => { + const caps = makeCaps({ + 'text-only/chat': { inputModalities: [] }, + 'vision/chat': { inputModalities: [] }, + }); + const decision = computeDecision( + classification, + visionTable, + null, + new Set(), + 'cost_per_accuracy', + { + constraints: { requiredInputModalities: ['image'] }, + capabilityMap: caps, + } + ); + expect(decision).toBeNull(); + }); + + it('fails closed when capabilityMap is missing and a required modality is set', () => { + const decision = computeDecision( + classification, + visionTable, + null, + new Set(), + 'cost_per_accuracy', + { + constraints: { requiredInputModalities: ['image'] }, + } + ); + expect(decision).toBeNull(); + }); + + it('replaces a non-vision sticky incumbent when the request gains an image requirement', () => { + // The text-only incumbent would normally be kept (cheap + accurate), + // but it lacks the image modality required by the new constraints, so + // the engine must pick a fresh vision candidate. + const caps = makeCaps({ + 'text-only/chat': { inputModalities: [] }, + 'vision/chat': { inputModalities: ['image'] }, + 'premium-vision/chat': { inputModalities: ['image'] }, + }); + const decision = computeDecision( + classification, + visionTable, + 'text-only/chat', + new Set(), + 'cost_per_accuracy', + { + constraints: { requiredInputModalities: ['image'] }, + capabilityMap: caps, + } + ); + expect(decision).toMatchObject({ model: 'vision/chat', sticky: false }); + }); + + it('a fitting lower-ranked candidate wins over a provably-too-small top candidate', () => { + const sizedTable: RoutingTable = { + ...table, + routes: { + ...table.routes, + 'implementation/code_generation': [ + { model: 'tiny/chat', accuracy: 0.95, avgCostUsd: 0.001, meetsThreshold: true }, + { model: 'large/chat', accuracy: 0.7, avgCostUsd: 0.003, meetsThreshold: true }, + ], + }, + }; + const caps = makeCaps({ + 'tiny/chat': { inputModalities: [], contextLength: 4_000 }, + 'large/chat': { inputModalities: [], contextLength: 1_000_000 }, + }); + const decision = computeDecision( + classification, + sizedTable, + null, + new Set(), + 'cost_per_accuracy', + { + constraints: { promptTokensEstimate: 50_000 }, + capabilityMap: caps, + } + ); + expect(decision).toMatchObject({ model: 'large/chat', sticky: false }); + }); + + it('keeps an unknown-context top candidate over a known-fitting lower candidate (no regression)', () => { + const sizedTable: RoutingTable = { + ...table, + routes: { + ...table.routes, + 'implementation/code_generation': [ + { model: 'unknown-ctx/chat', accuracy: 0.95, avgCostUsd: 0.001, meetsThreshold: true }, + { model: 'large/chat', accuracy: 0.7, avgCostUsd: 0.003, meetsThreshold: true }, + ], + }, + }; + const caps = makeCaps({ + 'unknown-ctx/chat': { inputModalities: [], contextLength: null }, + 'large/chat': { inputModalities: [], contextLength: 1_000_000 }, + }); + const decision = computeDecision( + classification, + sizedTable, + null, + new Set(), + 'cost_per_accuracy', + { + constraints: { promptTokensEstimate: 50_000 }, + capabilityMap: caps, + } + ); + expect(decision).toMatchObject({ model: 'unknown-ctx/chat', sticky: false }); + }); + + it('replaces a provably-too-small sticky incumbent with a fresh eligible pick', () => { + const sizedTable: RoutingTable = { + ...table, + routes: { + ...table.routes, + 'implementation/code_generation': [ + { model: 'large/chat', accuracy: 0.9, avgCostUsd: 0.002, meetsThreshold: true }, + { model: 'huge/chat', accuracy: 0.7, avgCostUsd: 0.003, meetsThreshold: true }, + ], + }, + }; + const caps = makeCaps({ + 'large/chat': { inputModalities: [], contextLength: 4_000 }, + 'huge/chat': { inputModalities: [], contextLength: 1_000_000 }, + }); + const decision = computeDecision( + classification, + sizedTable, + 'large/chat', + new Set(), + 'cost_per_accuracy', + { + constraints: { promptTokensEstimate: 50_000 }, + capabilityMap: caps, + } + ); + expect(decision).toMatchObject({ model: 'huge/chat', sticky: false }); + }); + + it('falls back to the max-known-context candidate when every known context is too small', () => { + const sizedTable: RoutingTable = { + ...table, + routes: { + ...table.routes, + 'implementation/code_generation': [ + { model: 'small/chat', accuracy: 0.95, avgCostUsd: 0.001, meetsThreshold: true }, + { model: 'medium/chat', accuracy: 0.9, avgCostUsd: 0.002, meetsThreshold: true }, + { model: 'unknown-ctx/chat', accuracy: 0.7, avgCostUsd: 0.003, meetsThreshold: true }, + ], + }, + }; + const caps = makeCaps({ + 'small/chat': { inputModalities: [], contextLength: 4_000 }, + 'medium/chat': { inputModalities: [], contextLength: 8_000 }, + 'unknown-ctx/chat': { inputModalities: [], contextLength: null }, + }); + // 50k tokens is bigger than even the largest known context; the + // unknown-context candidate keeps its rank (it is not provably too + // small) so it wins. + const decision = computeDecision( + classification, + sizedTable, + null, + new Set(), + 'cost_per_accuracy', + { + constraints: { promptTokensEstimate: 50_000 }, + capabilityMap: caps, + } + ); + expect(decision).toMatchObject({ model: 'unknown-ctx/chat', sticky: false }); + }); + + it('falls back to the max-known-context candidate when every known context is too small AND no unknown exists', () => { + const sizedTable: RoutingTable = { + ...table, + routes: { + ...table.routes, + 'implementation/code_generation': [ + { model: 'small/chat', accuracy: 0.95, avgCostUsd: 0.001, meetsThreshold: true }, + { model: 'medium/chat', accuracy: 0.9, avgCostUsd: 0.002, meetsThreshold: true }, + { model: 'largest/chat', accuracy: 0.7, avgCostUsd: 0.003, meetsThreshold: true }, + ], + }, + }; + const caps = makeCaps({ + 'small/chat': { inputModalities: [], contextLength: 4_000 }, + 'medium/chat': { inputModalities: [], contextLength: 8_000 }, + 'largest/chat': { inputModalities: [], contextLength: 32_000 }, + }); + const decision = computeDecision( + classification, + sizedTable, + null, + new Set(), + 'cost_per_accuracy', + { + constraints: { promptTokensEstimate: 50_000 }, + capabilityMap: caps, + } + ); + expect(decision).toMatchObject({ model: 'largest/chat', sticky: false }); + }); + + it('preserves existing ranking and sticky behaviour when all contexts are unknown', () => { + const sizedTable: RoutingTable = { + ...table, + routes: { + ...table.routes, + 'implementation/code_generation': [ + { model: 'a/chat', accuracy: 0.95, avgCostUsd: 0.001, meetsThreshold: true }, + { model: 'b/chat', accuracy: 0.9, avgCostUsd: 0.002, meetsThreshold: true }, + ], + }, + }; + const caps = makeCaps({ + 'a/chat': { inputModalities: [], contextLength: null }, + 'b/chat': { inputModalities: [], contextLength: null }, + }); + const decision = computeDecision( + classification, + sizedTable, + 'b/chat', + new Set(), + 'cost_per_accuracy', + { + constraints: { promptTokensEstimate: 50_000 }, + capabilityMap: caps, + } + ); + // b/chat is the incumbent but is more expensive than a/chat by less + // than 3x, so the sticky rule keeps it. + expect(decision).toMatchObject({ model: 'b/chat', sticky: true }); + }); + + it('a fitting text-only request with only a token estimate preserves the no-constraints winner', () => { + const caps = makeCaps({ + 'cheap/chat': { inputModalities: [], contextLength: 1_000_000 }, + 'mid/chat': { inputModalities: [], contextLength: 1_000_000 }, + 'pricey/chat': { inputModalities: [], contextLength: 1_000_000 }, + }); + const noConstraints = computeDecision(classification, table, null); + const withConstraints = computeDecision( + classification, + table, + null, + new Set(), + 'cost_per_accuracy', + { + constraints: { promptTokensEstimate: 1_000 }, + capabilityMap: caps, + } + ); + expect(withConstraints?.model).toBe(noConstraints?.model); + expect(withConstraints?.sticky).toBe(false); + }); + + it('a fitting text-only request with only a token estimate preserves the no-constraints winner in best_accuracy mode', () => { + const caps = makeCaps({ + 'pricey/chat': { inputModalities: [], contextLength: 1_000_000 }, + }); + const noConstraints = computeDecision( + classification, + table, + null, + new Set(), + 'best_accuracy' + ); + const withConstraints = computeDecision( + classification, + table, + null, + new Set(), + 'best_accuracy', + { + constraints: { promptTokensEstimate: 1_000 }, + capabilityMap: caps, + } + ); + expect(withConstraints?.model).toBe(noConstraints?.model); + }); + + it('treats constraints with no fields set as a no-op filter (regression guarantee)', () => { + // Spec: "if [constraints is] present with genuinely no fields set, + // behaviour should still reduce to a no-op filter per the no-op + // rules above". + const noConstraints = computeDecision(classification, table, null); + const emptyConstraints = computeDecision( + classification, + table, + null, + new Set(), + 'cost_per_accuracy', + { + constraints: {}, + } + ); + expect(emptyConstraints).toEqual(noConstraints); + }); + }); }); diff --git a/services/auto-routing/src/decision-engine.ts b/services/auto-routing/src/decision-engine.ts index 19ed140588..28b567030c 100644 --- a/services/auto-routing/src/decision-engine.ts +++ b/services/auto-routing/src/decision-engine.ts @@ -6,8 +6,20 @@ import { type AutoRoutingMode, type ClassifierOutput, type RankedCandidate, + type RoutingConstraints, type RoutingTable, } from '@kilocode/auto-routing-contracts'; +import type { ModelCapabilitiesMap } from './model-capabilities'; + +// Modalities the worker actively enforces against `model_stats.input_modalities`. +// Required modalities outside this set are intentionally ignored: they pass +// the filter today even though we have no way to confirm candidate support. +// Vocabulary evidence: `image` is folded from `image` / `image_url` per the +// existing web-side `modelSupportsImages` helper, and `file` is a confirmed +// OpenRouter `architecture.input_modalities` value (documented enum: +// `text | image | file | audio | video`), mirrored verbatim into +// `model_stats.inputModalities` (`apps/web/src/lib/model-stats/sync-openrouter.ts:77,95,124`). +export const ENFORCED_MODALITIES: ReadonlyArray = ['image', 'file']; function pickFreshCandidate( candidates: ReadonlyArray, @@ -29,19 +41,113 @@ function pickFreshCandidate( return candidate; } +// Apply the modality and context filters to the route candidates. +// +// * `ENFORCED_MODALITIES` is the only vocabulary we check: required +// modalities outside the set are ignored (no fail-closed for unknown +// vocabulary) so a future gateway sending `audio` does not break routing +// before the worker learns to honour it. +// * Missing capability data is treated the same as "no modalities" and +// fails the modality check; that matches the existing fail-closed web- +// side behaviour for image support. +// * Unknown context length is NOT proof of unfitness: a candidate whose +// row is missing `context_length` keeps its rank inside the eligible +// set. Only candidates with a known, provably-too-small context are +// excluded. +// * When every candidate's known context is provably too small, fall +// back to the candidates sharing the maximum known context so a +// large-but-still-too-small model is preferred over a slightly-smaller +// one we know cannot fit either. +function applyCapabilityFilters( + candidates: ReadonlyArray, + constraints: RoutingConstraints | undefined, + capabilityMap: ModelCapabilitiesMap | undefined +): { filtered: ReadonlyArray; reason: 'empty' | 'no_constraints' | 'ok' } { + if (!constraints) { + return { filtered: candidates, reason: 'no_constraints' }; + } + + const required = constraints.requiredInputModalities ?? []; + const enforcedAndRequired = required.filter(m => ENFORCED_MODALITIES.includes(m)); + + const modalityOk = (model: string): boolean => { + if (enforcedAndRequired.length === 0) return true; + const caps = capabilityMap?.get(model); + if (!caps) return false; + for (const modality of enforcedAndRequired) { + if (!caps.inputModalities.has(modality)) return false; + } + return true; + }; + + const afterModality = candidates.filter(c => modalityOk(c.model)); + if (afterModality.length === 0) { + return { filtered: [], reason: 'empty' }; + } + + const estimate = constraints.promptTokensEstimate; + if (typeof estimate !== 'number') { + return { filtered: afterModality, reason: 'ok' }; + } + + const eligible: RankedCandidate[] = []; + const provablyTooSmall: RankedCandidate[] = []; + for (const candidate of afterModality) { + const caps = capabilityMap?.get(candidate.model); + if (caps && typeof caps.contextLength === 'number' && caps.contextLength < estimate) { + provablyTooSmall.push(candidate); + } else { + eligible.push(candidate); + } + } + + if (eligible.length > 0) { + return { filtered: eligible, reason: 'ok' }; + } + + // Every candidate's known context is too small. Pick the candidates + // sharing the maximum known context so the largest-context option wins. + let maxKnown = -Infinity; + for (const candidate of provablyTooSmall) { + const caps = capabilityMap?.get(candidate.model); + if (caps && typeof caps.contextLength === 'number' && caps.contextLength > maxKnown) { + maxKnown = caps.contextLength; + } + } + const maxContextFallback = provablyTooSmall.filter(candidate => { + const caps = capabilityMap?.get(candidate.model); + return caps?.contextLength === maxKnown; + }); + return { filtered: maxContextFallback, reason: 'ok' }; +} + export function computeDecision( classification: ClassifierOutput, table: RoutingTable | null, incumbentModel: string | null, deniedModelIds: ReadonlySet = new Set(), - mode: AutoRoutingMode = DEFAULT_AUTO_ROUTING_MODE + mode: AutoRoutingMode = DEFAULT_AUTO_ROUTING_MODE, + options: { + constraints?: RoutingConstraints | undefined; + capabilityMap?: ModelCapabilitiesMap | undefined; + } = {} ): AutoRoutingDecision | null { if (!table) return null; const routeKey = taxonomyRouteKey(classification); - const candidates = table.routes[routeKey]?.filter( + const routeCandidates = table.routes[routeKey]?.filter( c => !deniedModelIds.has(c.model) && !isVirtualAutoModelId(c.model) ); - if (!candidates?.length) return null; + if (!routeCandidates?.length) return null; + + const { filtered: candidates, reason } = applyCapabilityFilters( + routeCandidates, + options.constraints, + options.capabilityMap + ); + if (reason === 'empty' || candidates.length === 0) { + return null; + } + const freshPick = pickFreshCandidate(candidates, mode); // Keep the session on its incumbent model when it is still good enough for @@ -50,6 +156,9 @@ export function computeDecision( // on a context that dominates agent-session spend — so a switch is only // worth it when the fresh pick's recurring per-turn savings clearly exceed // that one-time penalty, i.e. it is cheaper by more than switchCostFactor. + // Sticky lookup is performed against the filtered candidate set so an + // incumbent that is modality-incapable or provably too small is replaced + // by a fresh pick from the eligible set, not kept. const incumbent = incumbentModel === null ? undefined : candidates.find(c => c.model === incumbentModel); const stickyIncumbent = diff --git a/services/auto-routing/src/index.test.ts b/services/auto-routing/src/index.test.ts index ea8feed6f9..01a682309a 100644 --- a/services/auto-routing/src/index.test.ts +++ b/services/auto-routing/src/index.test.ts @@ -1,6 +1,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { clearClassifierConfigCache } from './classifier-config'; import { clearRoutingTableCache } from './routing-table'; +import { clearModelCapabilitiesCache } from './model-capabilities'; import { app } from './index'; import { ClassifierRunError } from './model-classifier'; import type * as DbModule from '@kilocode/db'; @@ -13,6 +14,8 @@ const dbFrom = vi.hoisted(() => vi.fn()); const dbInnerJoin = vi.hoisted(() => vi.fn()); const dbWhere = vi.hoisted(() => vi.fn()); const dbLimit = vi.hoisted(() => vi.fn()); +// Model-capabilities mock chain (select -> from -> where, no innerJoin/limit). +const dbWhereCaps = vi.hoisted(() => vi.fn()); vi.mock('./model-classifier', async importOriginal => { const actual = await importOriginal(); @@ -192,6 +195,7 @@ describe('auto routing worker', () => { beforeEach(() => { clearClassifierConfigCache(); clearRoutingTableCache(); + clearModelCapabilitiesCache(); classifyNormalizedInput.mockReset(); classifyNormalizedInput.mockResolvedValue(mockClassifierResult); getWorkerDb.mockReset(); @@ -199,13 +203,19 @@ describe('auto routing worker', () => { dbSelect.mockReset(); dbSelect.mockReturnValue({ from: dbFrom }); dbFrom.mockReset(); - dbFrom.mockReturnValue({ innerJoin: dbInnerJoin }); + // The coding-plan path goes through `innerJoin -> where -> limit`; the + // model-capabilities path goes straight to `where` and awaits a plain + // promise. Both are mounted on the same `from()` so a single test can + // exercise either chain without a separate mock harness. + dbFrom.mockReturnValue({ innerJoin: dbInnerJoin, where: dbWhereCaps }); dbInnerJoin.mockReset(); dbInnerJoin.mockReturnValue({ where: dbWhere }); dbWhere.mockReset(); dbWhere.mockReturnValue({ limit: dbLimit }); dbLimit.mockReset(); dbLimit.mockResolvedValue([]); + dbWhereCaps.mockReset(); + dbWhereCaps.mockResolvedValue([]); writeDataPoint.mockReset(); configGet.mockReset(); // Real KV returns null for missing keys; an undefined here would send the @@ -251,6 +261,298 @@ describe('auto routing worker', () => { vi.restoreAllMocks(); }); + describe('capability-aware routing', () => { + // A two-candidate route where the cheaper model is text-only and the + // second is image-capable. This lets a single fixture exercise both + // the fresh, cached, and fallback code paths in decide.ts. + const visionTable = { + ...benchmarkRoutingTable, + routes: { + ...benchmarkRoutingTable.routes, + 'implementation/feature_development': [ + { + model: 'text-only/chat', + accuracy: 0.95, + avgCostUsd: 0.001, + meetsThreshold: true, + reasoningEffort: null, + }, + { + model: 'vision/chat', + accuracy: 0.85, + avgCostUsd: 0.002, + meetsThreshold: true, + reasoningEffort: null, + }, + ], + }, + }; + + function setVisionBenchmark() { + benchmarkFetch.mockImplementation(async (url: string) => { + if (String(url).includes('/admin/classifier-winner')) { + return { ok: true, status: 200, json: async () => ({ winner: null }) }; + } + return { + ok: true, + status: 200, + json: async () => ({ + table: visionTable, + publishedAt: visionTable.generatedAt, + }), + }; + }); + } + + function setVisionCaps() { + dbWhereCaps.mockResolvedValue([ + { openrouterId: 'text-only/chat', inputModalities: [], contextLength: 1_000_000 }, + { openrouterId: 'vision/chat', inputModalities: ['image'], contextLength: 1_000_000 }, + ]); + } + + it('skips a non-vision top candidate on the fresh-classification path', async () => { + setVisionBenchmark(); + setVisionCaps(); + const response = await decideRequest( + mirrorPayload({ constraints: { requiredInputModalities: ['image'] } }) + ); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + decision: { model: 'vision/chat', sticky: false }, + }); + }); + + it('skips a non-vision top candidate on the cached-classification-hit path', async () => { + setVisionBenchmark(); + setVisionCaps(); + cacheGetEntry.mockResolvedValueOnce(mockClassification); + const response = await decideRequest( + mirrorPayload({ constraints: { requiredInputModalities: ['image'] } }) + ); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + decision: { model: 'vision/chat', sticky: false }, + classifierResult: { classification: mockClassification }, + }); + expect(classifyNormalizedInput).not.toHaveBeenCalled(); + }); + + it('skips a non-vision top candidate on the heuristic-fallback-classification path', async () => { + setVisionBenchmark(); + setVisionCaps(); + classifyNormalizedInput.mockResolvedValueOnce({ + ...mockClassifierResult, + classification: { ...mockClassification, confidence: 0 }, + fallback: { reason: 'invalid_output' }, + }); + const response = await decideRequest( + mirrorPayload({ constraints: { requiredInputModalities: ['image'] } }) + ); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + decision: { model: 'vision/chat', sticky: false }, + }); + // A fallback classification must not re-anchor the sticky model. + expect(cachePutEntry).not.toHaveBeenCalledWith('sticky', expect.anything()); + }); + + it('is byte-identical for an old-gateway payload (no constraints field)', async () => { + vi.spyOn(Math, 'random').mockReturnValue(0); + const response = await decideRequest(mirrorPayload()); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + decision: { model: 'google/gemini-2.5-flash-lite', sticky: false }, + }); + // No capability fetch on the old-gateway path: the DB chain is not + // touched by the capability lookup. + expect(dbWhereCaps).not.toHaveBeenCalled(); + }); + + it('proceeds unfiltered when capability lookup fails and constraints only carry an estimate', async () => { + vi.spyOn(Math, 'random').mockReturnValue(0); + dbWhereCaps.mockRejectedValue(new Error('db down')); + const response = await decideRequest( + mirrorPayload({ constraints: { promptTokensEstimate: 1_000 } }) + ); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + decision: { model: 'google/gemini-2.5-flash-lite', sticky: false }, + }); + }); + + it('returns null when capability lookup fails and constraints require an image', async () => { + dbWhereCaps.mockRejectedValue(new Error('db down')); + const response = await decideRequest( + mirrorPayload({ constraints: { requiredInputModalities: ['image'] } }) + ); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + decision: null, + classifierResult: { classification: mockClassification }, + }); + }); + + it('falls through the coding-plan short-circuit when the model lacks a required modality', async () => { + configGet.mockImplementation(async (key: string) => + key.startsWith('coding_plan_preference:') + ? JSON.stringify({ + active: true, + planId: 'minimax-token-plan-plus', + providerId: 'minimax', + modelId: 'minimax/minimax-m3', + }) + : null + ); + // The coding-plan default model has no image modality → short- + // circuit guard rejects it and the request falls through to a + // benchmark candidate. The benchmark table's top candidate is + // also text-only, so we need a vision-capable candidate to be + // available in the route. + setVisionBenchmark(); + dbWhereCaps.mockResolvedValue([ + { + openrouterId: 'minimax/minimax-m3', + inputModalities: ['text'], + contextLength: 1_000_000, + }, + { openrouterId: 'text-only/chat', inputModalities: [], contextLength: 1_000_000 }, + { openrouterId: 'vision/chat', inputModalities: ['image'], contextLength: 1_000_000 }, + ]); + const response = await decideRequest( + mirrorPayload({ constraints: { requiredInputModalities: ['image'] } }) + ); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + decision: { + model: 'vision/chat', + source: 'benchmark', + sticky: false, + }, + }); + // The coding-plan short-circuit did not fire: the benchmark path ran. + expect(classifyNormalizedInput).toHaveBeenCalledTimes(1); + }); + + it('falls through the coding-plan short-circuit when the estimate exceeds the model context', async () => { + configGet.mockImplementation(async (key: string) => + key.startsWith('coding_plan_preference:') + ? JSON.stringify({ + active: true, + planId: 'minimax-token-plan-plus', + providerId: 'minimax', + modelId: 'minimax/minimax-m3', + }) + : null + ); + setVisionBenchmark(); + dbWhereCaps.mockResolvedValue([ + { + openrouterId: 'minimax/minimax-m3', + inputModalities: ['image'], + contextLength: 8_000, + }, + { openrouterId: 'text-only/chat', inputModalities: [], contextLength: 4_000 }, + { openrouterId: 'vision/chat', inputModalities: ['image'], contextLength: 1_000_000 }, + ]); + const response = await decideRequest( + mirrorPayload({ + constraints: { promptTokensEstimate: 50_000 }, + }) + ); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + decision: { + model: 'vision/chat', + source: 'benchmark', + sticky: false, + }, + }); + expect(classifyNormalizedInput).toHaveBeenCalledTimes(1); + }); + + it('takes the coding-plan short-circuit when the model context is unknown', async () => { + configGet.mockImplementation(async (key: string) => + key.startsWith('coding_plan_preference:') + ? JSON.stringify({ + active: true, + planId: 'minimax-token-plan-plus', + providerId: 'minimax', + modelId: 'minimax/minimax-m3', + }) + : null + ); + // No image requirement; estimate present but context is null. The + // unknown-keeps-rank policy applies: short-circuit is still taken. + dbWhereCaps.mockResolvedValue([ + { + openrouterId: 'minimax/minimax-m3', + inputModalities: ['text'], + contextLength: null, + }, + ]); + const response = await decideRequest( + mirrorPayload({ + constraints: { promptTokensEstimate: 50_000 }, + }) + ); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + decision: { model: 'minimax/minimax-m3', source: 'coding_plan_default' }, + }); + expect(classifyNormalizedInput).not.toHaveBeenCalled(); + }); + + it('returns null when capability lookup fails on the coding-plan path with an image requirement', async () => { + configGet.mockImplementation(async (key: string) => + key.startsWith('coding_plan_preference:') + ? JSON.stringify({ + active: true, + planId: 'minimax-token-plan-plus', + providerId: 'minimax', + modelId: 'minimax/minimax-m3', + }) + : null + ); + // Lookup fails → no capability data for the coding-plan model → the + // short-circuit guard cannot confirm the model supports image, so + // it must fail closed rather than take the possibly-incapable + // short-circuit. + dbWhereCaps.mockRejectedValue(new Error('db down')); + const response = await decideRequest( + mirrorPayload({ constraints: { requiredInputModalities: ['image'] } }) + ); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + decision: null, + classifierResult: { classification: mockClassification }, + }); + expect(classifyNormalizedInput).toHaveBeenCalledTimes(1); + }); + + it('enforces file modality the same way it enforces image', async () => { + // 'file' is a first-class modality in the gateway's request-side + // detector and is in ENFORCED_MODALITIES. A candidate without + // 'file' in its known input_modalities must be excluded. + setVisionBenchmark(); + dbWhereCaps.mockResolvedValue([ + { openrouterId: 'text-only/chat', inputModalities: ['image'], contextLength: 1_000_000 }, + { + openrouterId: 'vision/chat', + inputModalities: ['image', 'file'], + contextLength: 1_000_000, + }, + ]); + const response = await decideRequest( + mirrorPayload({ constraints: { requiredInputModalities: ['file'] } }) + ); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + decision: { model: 'vision/chat', sticky: false }, + }); + }); + }); + it('returns health without requiring classifier payload fields', async () => { const response = await request('/health', { headers: { authorization: 'Bearer classifier-token' }, diff --git a/services/auto-routing/src/model-capabilities.test.ts b/services/auto-routing/src/model-capabilities.test.ts new file mode 100644 index 0000000000..c21193110a --- /dev/null +++ b/services/auto-routing/src/model-capabilities.test.ts @@ -0,0 +1,357 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { clearModelCapabilitiesCache, getModelCapabilities } from './model-capabilities'; +import { clearRoutingTableCache } from './routing-table'; +import type * as RoutingTableModule from './routing-table'; +import type * as DbModule from '@kilocode/db'; +import type { RoutingTable } from '@kilocode/auto-routing-contracts'; + +const getWorkerDb = vi.hoisted(() => vi.fn()); +const dbSelect = vi.hoisted(() => vi.fn()); +const dbFrom = vi.hoisted(() => vi.fn()); +const dbWhere = vi.hoisted(() => vi.fn()); +const mockGetRoutingTable = vi.hoisted(() => vi.fn()); + +vi.mock('@kilocode/db', async importOriginal => { + const actual = await importOriginal(); + return { ...actual, getWorkerDb }; +}); + +vi.mock('./routing-table', async importOriginal => { + const actual = await importOriginal(); + return { ...actual, getRoutingTable: mockGetRoutingTable }; +}); + +const SAMPLE_ROUTING_TABLE: RoutingTable = { + version: 'bench-1', + generatedAt: '2026-06-12T00:00:00.000Z', + minAccuracy: 0.7, + switchCostFactor: 3, + bestAccuracySwitchThreshold: 0.05, + source: 'benchmark', + routes: { + 'implementation/code_generation': [ + { model: 'a/chat', accuracy: 0.9, avgCostUsd: 0.001, meetsThreshold: true }, + { model: 'b/chat', accuracy: 0.85, avgCostUsd: 0.002, meetsThreshold: true }, + ], + }, +}; + +function makeEnv(kvValue: string | null): Env { + return { + AUTO_ROUTING_CONFIG: { + get: vi.fn(async () => kvValue), + put: vi.fn(async () => undefined), + } as unknown as KVNamespace, + HYPERDRIVE: { connectionString: 'postgres://worker' } as Hyperdrive, + BENCHMARK_SERVICE: { + fetch: vi.fn(async () => ({ + ok: true, + status: 200, + json: async () => ({ + table: SAMPLE_ROUTING_TABLE, + publishedAt: SAMPLE_ROUTING_TABLE.generatedAt, + }), + })), + } as unknown as Fetcher, + INTERNAL_API_SECRET_PROD: { get: async () => 'secret' } as unknown as SecretsStoreSecret, + } as unknown as Env; +} + +afterEach(() => { + clearModelCapabilitiesCache(); + clearRoutingTableCache(); +}); + +beforeEach(() => { + getWorkerDb.mockReset(); + getWorkerDb.mockReturnValue({ select: dbSelect }); + dbSelect.mockReset(); + dbSelect.mockReturnValue({ from: dbFrom }); + dbFrom.mockReset(); + dbFrom.mockReturnValue({ where: dbWhere }); + dbWhere.mockReset(); + dbWhere.mockImplementation(() => Promise.resolve([])); + mockGetRoutingTable.mockReset(); + mockGetRoutingTable.mockResolvedValue(SAMPLE_ROUTING_TABLE); +}); + +describe('getModelCapabilities', () => { + it('folds image_url to image in the capability set', async () => { + dbWhere.mockImplementation(() => + Promise.resolve([ + { openrouterId: 'a/chat', inputModalities: ['image_url'], contextLength: 8192 }, + ]) + ); + const env = makeEnv(null); + const result = await getModelCapabilities(env); + expect(result.get('a/chat')?.inputModalities.has('image')).toBe(true); + expect(result.get('a/chat')?.inputModalities.has('image_url')).toBe(false); + expect(result.get('a/chat')?.contextLength).toBe(8192); + }); + + it('folds confirmed real input modalities to their canonical forms', async () => { + dbWhere.mockImplementation(() => + Promise.resolve([ + { openrouterId: 'doc/chat', inputModalities: ['image_url', 'file'], contextLength: 32768 }, + ]) + ); + const env = makeEnv(null); + const result = await getModelCapabilities(env); + const set = result.get('doc/chat')?.inputModalities; + expect(set?.has('image')).toBe(true); // image_url folded to canonical image + expect(set?.has('file')).toBe(true); // file is a real input modality + expect(set?.has('image_url')).toBe(false); + }); + + it('treats null input_modalities as an empty modality set, not a failure', async () => { + dbWhere.mockImplementation(() => + Promise.resolve([{ openrouterId: 'a/chat', inputModalities: null, contextLength: 4096 }]) + ); + const env = makeEnv(null); + const result = await getModelCapabilities(env); + expect(result.get('a/chat')?.inputModalities.size).toBe(0); + expect(result.get('a/chat')?.contextLength).toBe(4096); + }); + + it('caches results in KV and avoids a second DB read on subsequent calls', async () => { + dbWhere.mockImplementation(() => + Promise.resolve([ + { openrouterId: 'a/chat', inputModalities: ['image'], contextLength: 8192 }, + { openrouterId: 'b/chat', inputModalities: ['text'], contextLength: 16384 }, + ]) + ); + const env = makeEnv(null); + const first = await getModelCapabilities(env); + const second = await getModelCapabilities(env); + expect(first.get('a/chat')?.inputModalities.has('image')).toBe(true); + expect(second.get('a/chat')?.inputModalities.has('image')).toBe(true); + // The DB is only hit on the first call; the second call satisfies from + // the in-memory cache (no DB read, no KV read). + expect(dbWhere).toHaveBeenCalledTimes(1); + }); + + it('reads from KV on in-memory miss and avoids the DB', async () => { + const cached = { + 'a/chat': { inputModalities: ['image'], contextLength: 8192 }, + 'b/chat': { inputModalities: ['text'], contextLength: 16384 }, + }; + const env = makeEnv(JSON.stringify(cached)); + const result = await getModelCapabilities(env); + expect(result.get('a/chat')?.inputModalities.has('image')).toBe(true); + expect(result.get('b/chat')?.contextLength).toBe(16384); + expect(dbWhere).not.toHaveBeenCalled(); + }); + + it('writes the queried rows to KV on a true miss with the configured expirationTtl', async () => { + const put = vi.fn(async () => undefined); + const env = { + AUTO_ROUTING_CONFIG: { + get: vi.fn(async () => null), + put, + } as unknown as KVNamespace, + HYPERDRIVE: { connectionString: 'postgres://worker' } as Hyperdrive, + BENCHMARK_SERVICE: { + fetch: vi.fn(async () => ({ + ok: true, + status: 200, + json: async () => ({ + table: SAMPLE_ROUTING_TABLE, + publishedAt: SAMPLE_ROUTING_TABLE.generatedAt, + }), + })), + } as unknown as Fetcher, + INTERNAL_API_SECRET_PROD: { get: async () => 'secret' } as unknown as SecretsStoreSecret, + } as unknown as Env; + dbWhere.mockImplementation(() => + Promise.resolve([{ openrouterId: 'a/chat', inputModalities: ['image'], contextLength: 8192 }]) + ); + + await getModelCapabilities(env); + + expect(put).toHaveBeenCalledWith('model_capabilities_v1', expect.stringContaining('"a/chat"'), { + expirationTtl: 3600, + }); + }); + + it('returns an empty map and does NOT write to KV when the DB throws', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}); + const put = vi.fn(async () => undefined); + const env = { + AUTO_ROUTING_CONFIG: { + get: vi.fn(async () => null), + put, + } as unknown as KVNamespace, + HYPERDRIVE: { connectionString: 'postgres://worker' } as Hyperdrive, + BENCHMARK_SERVICE: { + fetch: vi.fn(async () => ({ + ok: true, + status: 200, + json: async () => ({ + table: SAMPLE_ROUTING_TABLE, + publishedAt: SAMPLE_ROUTING_TABLE.generatedAt, + }), + })), + } as unknown as Fetcher, + INTERNAL_API_SECRET_PROD: { get: async () => 'secret' } as unknown as SecretsStoreSecret, + } as unknown as Env; + dbWhere.mockImplementation(() => Promise.reject(new Error('db down'))); + + const result = await getModelCapabilities(env); + expect(result.size).toBe(0); + // The model_capabilities_v1 key is never written; the routing-table + // lookup on the cache-miss path may write the routing_table_v1 key, + // and that is unrelated to capability data. + const capabilityPuts = put.mock.calls.filter( + (call: unknown[]) => call[0] === 'model_capabilities_v1' + ); + expect(capabilityPuts).toEqual([]); + expect(warn).toHaveBeenCalled(); + warn.mockRestore(); + }); + + it('returns an empty map promptly when the underlying load exceeds the sub-budget (named timing test)', async () => { + vi.useFakeTimers(); + try { + // Simulate a slow Hyperdrive: the DB promise never resolves in real + // time, so the 500ms sub-budget must trip first. + dbWhere.mockImplementation(() => new Promise(() => {}) as unknown as Promise); + const env = makeEnv(null); + + const resultP = getModelCapabilities(env); + // Advance the fake clock past the 500ms budget; the budget timer + // fires and rejects, which the wrapper converts to an empty Map. + await vi.advanceTimersByTimeAsync(600); + const result = await resultP; + expect(result.size).toBe(0); + } finally { + vi.useRealTimers(); + } + }); + + it('attaches a no-op swallow to the slow promise so no unhandled rejection escapes', async () => { + vi.useFakeTimers(); + const captured: unknown[] = []; + const onUnhandled = (reason: unknown) => { + captured.push(reason); + }; + process.on('unhandledRejection', onUnhandled); + try { + let rejectDb: (err: unknown) => void = () => {}; + dbWhere.mockImplementation( + () => + new Promise((_, reject) => { + rejectDb = reject; + }) as unknown as Promise + ); + const env = makeEnv(null); + + const resultP = getModelCapabilities(env); + await vi.advanceTimersByTimeAsync(600); + const result = await resultP; + expect(result.size).toBe(0); + + // Now reject the original DB promise; without a no-op catch it would + // surface as an unhandledRejection. + rejectDb(new Error('db failed after budget fired')); + // Let the rejection propagate; a tick is enough. + await Promise.resolve(); + await Promise.resolve(); + expect(captured).toEqual([]); + } finally { + process.off('unhandledRejection', onUnhandled); + vi.useRealTimers(); + } + }); + + it('returns an empty map promptly when the routing table fetch exceeds the sub-budget', async () => { + vi.useFakeTimers(); + try { + mockGetRoutingTable.mockImplementation( + () => new Promise(() => {}) as Promise + ); + const env = makeEnv(null); + + const resultP = getModelCapabilities(env); + await vi.advanceTimersByTimeAsync(600); + const result = await resultP; + expect(result.size).toBe(0); + } finally { + vi.useRealTimers(); + } + }); + + it('does not leak an unhandled rejection when the routing table fetch rejects after the budget', async () => { + vi.useFakeTimers(); + const captured: unknown[] = []; + const onUnhandled = (reason: unknown) => { + captured.push(reason); + }; + process.on('unhandledRejection', onUnhandled); + try { + let rejectRoutingTable: (err: unknown) => void = () => {}; + mockGetRoutingTable.mockImplementation( + () => + new Promise((_, reject) => { + rejectRoutingTable = reject; + }) as Promise + ); + const env = makeEnv(null); + + const resultP = getModelCapabilities(env); + await vi.advanceTimersByTimeAsync(600); + const result = await resultP; + expect(result.size).toBe(0); + + // Now reject the original routing-table promise; without a no-op catch + // it would surface as an unhandledRejection. + rejectRoutingTable(new Error('routing table failed after budget fired')); + await Promise.resolve(); + await Promise.resolve(); + expect(captured).toEqual([]); + } finally { + process.off('unhandledRejection', onUnhandled); + vi.useRealTimers(); + } + }); + + it('includes the coding-plan model id in the queried id set', async () => { + dbWhere.mockImplementation((..._args: unknown[]) => { + // First call is the in-cache-miss DB query (full id set), which will + // not happen because we are testing the partial-fill path. We still + // answer the partial-fill query for the coding-plan id. + return Promise.resolve([ + { openrouterId: 'coding-plan/chat', inputModalities: ['text'], contextLength: 200000 }, + ]); + }); + const env = makeEnv( + JSON.stringify({ + 'a/chat': { inputModalities: ['image'], contextLength: 8192 }, + 'b/chat': { inputModalities: ['text'], contextLength: 16384 }, + }) + ); + const result = await getModelCapabilities(env, { codingPlanModelId: 'coding-plan/chat' }); + expect(result.get('coding-plan/chat')?.contextLength).toBe(200000); + }); + + it('returns an empty map when the routing table is missing entirely', async () => { + mockGetRoutingTable.mockResolvedValue(null); + const env = { + AUTO_ROUTING_CONFIG: { + get: vi.fn(async () => null), + put: vi.fn(async () => undefined), + } as unknown as KVNamespace, + HYPERDRIVE: { connectionString: 'postgres://worker' } as Hyperdrive, + BENCHMARK_SERVICE: { + fetch: vi.fn(async () => ({ + ok: true, + status: 200, + json: async () => ({ table: null, publishedAt: null }), + })), + } as unknown as Fetcher, + INTERNAL_API_SECRET_PROD: { get: async () => 'secret' } as unknown as SecretsStoreSecret, + } as unknown as Env; + dbWhere.mockReset(); + const result = await getModelCapabilities(env); + expect(result.size).toBe(0); + }); +}); diff --git a/services/auto-routing/src/model-capabilities.ts b/services/auto-routing/src/model-capabilities.ts new file mode 100644 index 0000000000..7c9a112e07 --- /dev/null +++ b/services/auto-routing/src/model-capabilities.ts @@ -0,0 +1,258 @@ +import { formatError, ttlCached } from '@kilocode/worker-utils'; +import { getWorkerDb, modelStats } from '@kilocode/db'; +import { inArray } from 'drizzle-orm'; +import { kvReadThrough } from './kv-read-through'; +import { getRoutingTable } from './routing-table'; + +// Capability snapshot for a single model. `inputModalities` is the synonym- +// folded set (e.g. an `image_url` row is mapped to `image` so callers do not +// have to know the original vocabulary). `contextLength` is the published +// maximum input tokens, or `null` when the row is missing the column. +export type ModelCapabilities = { + inputModalities: ReadonlySet; + contextLength: number | null; +}; + +// An empty Map signals "no capability data" to callers: a request carrying +// `requiredInputModalities` fails closed, a request with only a token +// estimate proceeds unfiltered. A missing key for a specific model id +// carries the same meaning for that model. +export type ModelCapabilitiesMap = ReadonlyMap; + +// Modalities the worker actively enforces against `model_stats.input_modalities`. +// Vocabulary evidence: `image` / `image_url` folding mirrors +// `apps/web/src/lib/ai-gateway/providers/model-capabilities.ts:34`; `file` is a +// confirmed OpenRouter `architecture.input_modalities` value (documented enum: +// `text | image | file | audio | video`), and `model_stats.inputModalities` copies +// that field verbatim from the OpenRouter API +// (`apps/web/src/lib/model-stats/sync-openrouter.ts:77,95,124`). +const MODALITY_SYNONYMS: Readonly> = { + image: 'image', + image_url: 'image', + file: 'file', +}; + +function foldModalities(raw: ReadonlyArray | null | undefined): Set { + const out = new Set(); + if (!raw) return out; + for (const value of raw) { + const folded = MODALITY_SYNONYMS[value]; + if (folded !== undefined) { + out.add(folded); + } + } + return out; +} + +// CACHE LAYOUT +// +// `model_capabilities_v1` is a JSON object keyed by `openrouter_id` mapping +// to a `{ inputModalities: string[], contextLength: number | null }` row. +// The 1-hour KV TTL means a brand-new routing-table candidate can be +// fail-closed on constrained requests for up to an hour after publication; +// this is accepted as safe because the gateway's balanced fallback remains +// image-capable. The 60s in-memory TTL bounds the same fetch across +// requests within a warm isolate. +const MODEL_CAPABILITIES_KV_KEY = 'model_capabilities_v1'; +const MODEL_CAPABILITIES_IN_MEMORY_TTL_MS = 60_000; +const MODEL_CAPABILITIES_KV_TTL_SECONDS = 3_600; + +// Hard ceiling for the whole lookup (in-memory check + KV read + DB query). +// 500ms leaves headroom inside the gateway's 2s /decide budget when other +// steps are slow; the `statement_timeout: 2_000` on the Postgres side alone +// could otherwise let a slow-failing Hyperdrive connection eat the entire +// request budget. +const MODEL_CAPABILITIES_LOOKUP_BUDGET_MS = 500; + +type ModelCapabilitiesEnv = Pick< + Env, + 'AUTO_ROUTING_CONFIG' | 'HYPERDRIVE' | 'BENCHMARK_SERVICE' | 'INTERNAL_API_SECRET_PROD' +>; + +type ModelCapabilitiesCacheValue = Record< + string, + { inputModalities: string[]; contextLength: number | null } +>; + +function isCacheValue(value: unknown): value is ModelCapabilitiesCacheValue { + if (typeof value !== 'object' || value === null) return false; + for (const [key, entry] of Object.entries(value as Record)) { + if (typeof key !== 'string' || key.length === 0) return false; + if (typeof entry !== 'object' || entry === null) return false; + const row = entry as { inputModalities?: unknown; contextLength?: unknown }; + if (!Array.isArray(row.inputModalities)) return false; + if (row.contextLength !== null && typeof row.contextLength !== 'number') return false; + } + return true; +} + +async function queryModelCapabilities( + env: ModelCapabilitiesEnv, + modelIds: ReadonlyArray +): Promise { + if (modelIds.length === 0) return {}; + const db = getWorkerDb(env.HYPERDRIVE.connectionString, { statement_timeout: 2_000 }); + const rows = await db + .select({ + openrouterId: modelStats.openrouterId, + inputModalities: modelStats.inputModalities, + contextLength: modelStats.contextLength, + }) + .from(modelStats) + .where(inArray(modelStats.openrouterId, modelIds as string[])); + const out: ModelCapabilitiesCacheValue = {}; + for (const row of rows) { + if (typeof row.openrouterId !== 'string') continue; + out[row.openrouterId] = { + inputModalities: Array.isArray(row.inputModalities) ? row.inputModalities : [], + contextLength: typeof row.contextLength === 'number' ? row.contextLength : null, + }; + } + return out; +} + +const cache = ttlCached( + MODEL_CAPABILITIES_IN_MEMORY_TTL_MS, + async env => loadAll(env) +); + +function mergeInto( + target: Map, + source: Readonly +): void { + for (const [modelId, row] of Object.entries(source)) { + target.set(modelId, { + inputModalities: foldModalities(row.inputModalities), + contextLength: row.contextLength, + }); + } +} + +export function clearModelCapabilitiesCache(): void { + cache.clear(); +} + +// One-shot load that reads the full cached union of capability rows from +// KV, fills any missing entries from the DB, and returns the whole union +// (as a plain object so it is JSON-serialisable for the in-memory cache). +async function loadAll(env: ModelCapabilitiesEnv): Promise { + const fromKv = await kvReadThrough({ + kv: env.AUTO_ROUTING_CONFIG, + key: MODEL_CAPABILITIES_KV_KEY, + ttlSeconds: MODEL_CAPABILITIES_KV_TTL_SECONDS, + fetchOrigin: () => { + // Cache-miss path: ask the DB for every id we have ever needed. + // `loadAll` does not know the current id set, so it falls back to + // scanning the routing table for the canonical id set. + return queryAllIds(env); + }, + parse: (raw: string) => { + try { + const parsed: unknown = JSON.parse(raw); + if (!isCacheValue(parsed)) { + console.warn(JSON.stringify({ event: 'kv_model_capabilities_corrupt' })); + return null; + } + return parsed; + } catch (error) { + console.warn( + JSON.stringify({ event: 'kv_model_capabilities_corrupt', ...formatError(error) }) + ); + return null; + } + }, + }); + return fromKv ?? {}; +} + +async function queryAllIds(env: ModelCapabilitiesEnv): Promise { + const routingTable = await getRoutingTable(env); + const ids = new Set(); + if (routingTable) { + for (const route of Object.values(routingTable.routes)) { + for (const candidate of route) { + ids.add(candidate.model); + } + } + } + return queryModelCapabilities(env, Array.from(ids)); +} + +// Look up capability rows for the union of: every model in the published +// routing table, plus the coding-plan default model id when provided. The +// whole lookup (routing-table fetch + id derivation + in-memory check + KV +// read + DB query) is raced against a 500ms budget; on timeout or thrown +// error the returned Map is empty, which the caller treats as "no capability +// data". +export async function getModelCapabilities( + env: ModelCapabilitiesEnv, + options: { codingPlanModelId?: string | null } = {} +): Promise { + const load = async (): Promise> => { + // We derive the id set inside the module so the caller (decide.ts) does + // not have to wait on the routing-table fetch before kicking off the + // capability lookup. Keeping the fetch inside this closure means the + // 500ms sub-budget covers the routing-table read as well as the cache/DB + // lookups. `routing-table.ts`'s `ttlCached` dedups the concurrent in-flight + // call with whichever other component also asked for the table. + const routingTable = await getRoutingTable(env); + const ids = new Set(); + if (routingTable) { + for (const route of Object.values(routingTable.routes)) { + for (const candidate of route) { + ids.add(candidate.model); + } + } + } + if (options.codingPlanModelId) { + ids.add(options.codingPlanModelId); + } + const idList = Array.from(ids); + if (idList.length === 0) { + return new Map(); + } + + const result = new Map(); + const all = await cache.get(env); + mergeInto(result, all); + // The cache stores the union of all ids ever requested; fill the + // remainder from the DB. We don't write the partial-fill back to KV — + // a true cache miss above already wrote the full union, and a partial + // hit is rare enough that the extra round-trip is acceptable. + const missing = idList.filter(id => !result.has(id)); + if (missing.length > 0) { + const fromDb = await queryModelCapabilities(env, missing); + mergeInto(result, fromDb); + } + return result; + }; + + try { + return await raceWithBudget(load(), MODEL_CAPABILITIES_LOOKUP_BUDGET_MS); + } catch (error) { + console.warn( + JSON.stringify({ + event: 'auto_routing_capabilities_lookup_failed', + ...formatError(error), + }) + ); + return new Map(); + } +} + +// Race a promise against a millisecond budget without leaking the slow +// promise. The eventual rejection of the loser is intentionally swallowed +// so it never surfaces as an unhandled rejection after the budget has +// already fired. +function raceWithBudget(promise: Promise, budgetMs: number): Promise { + let timer: ReturnType | undefined; + const timeout = new Promise((_, reject) => { + timer = setTimeout(() => reject(new Error('capability lookup budget exceeded')), budgetMs); + }); + return Promise.race([promise, timeout]).finally(() => { + if (timer !== undefined) clearTimeout(timer); + // Attach a no-op catch so the losing promise does not surface as an + // unhandled rejection after the budget has already fired. + promise.catch(() => {}); + }); +} From a8d4d0d5d300a1c29c7dda395e538407b662f459 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Fri, 17 Jul 2026 11:50:53 +0200 Subject: [PATCH 4/4] fix(auto-routing): repair modality docs, Responses tool-output token count, and capability cache poisoning Kilobot review repair: - Remove the false doc-comment claim that MODALITY_CONTAINER_KEYS supports Gemini-style contents[].parts[] bodies; no code path can produce that shape here, so the comment was corrected instead of adding a dead-code container key. - estimateRoutingTokens now reads the output field (not content) for OpenAI Responses API function_call_output/tool_call_output items, which previously undercounted (or zero-counted) their token contribution; Anthropic/chat-style tool_result still reads content, unchanged. - model-capabilities.ts's queryAllIds now returns null (skipping the KV write, per kvReadThrough's own contract) when the routing table is transiently unavailable, instead of caching an empty {} capability map in KV for a full hour as if it were valid data for every request across the deployment. A genuinely empty-but-available routing table still caches {} correctly. --- .../src/normalize.test.ts | 49 +++++++++++++++++++ .../auto-routing-contracts/src/normalize.ts | 18 ++++--- .../src/model-capabilities.test.ts | 45 +++++++++++++++++ .../auto-routing/src/model-capabilities.ts | 13 ++--- 4 files changed, 112 insertions(+), 13 deletions(-) diff --git a/packages/auto-routing-contracts/src/normalize.test.ts b/packages/auto-routing-contracts/src/normalize.test.ts index 0df613a867..4cfc21d2e1 100644 --- a/packages/auto-routing-contracts/src/normalize.test.ts +++ b/packages/auto-routing-contracts/src/normalize.test.ts @@ -254,6 +254,25 @@ describe('detectRequiredInputModalities', () => { }) ).toEqual(['image']); }); + + it('does not claim support for Gemini native contents[].parts[] bodies', () => { + // Native Gemini request bodies never reach this helper: the gateway only + // invokes auto-routing for chat_completions / responses / messages shapes, + // and normalizeClassifierInput rejects any other top-level structure. The + // doc comment therefore does not advertise Gemini-style support. + expect( + detectRequiredInputModalities({ + contents: [ + { + parts: [ + { inline_data: { mime_type: 'image/png', data: 'iVBORw0KGgo=' } }, + { text: 'What is this?' }, + ], + }, + ], + }) + ).toEqual([]); + }); }); describe('estimateRoutingTokens', () => { @@ -395,6 +414,36 @@ describe('estimateRoutingTokens', () => { expect(estimateRoutingTokens({})).toBe(0); }); + it('counts Responses API function_call_output output field, not content', () => { + const output = 'x'.repeat(1000); + const estimate = estimateRoutingTokens({ + model: 'gpt-4o', + input: [{ type: 'function_call_output', call_id: 'call-1', output }], + }); + expect(estimate).toBe(Math.round(output.length / 4)); + }); + + it('does not zero-count a large Responses function_call_output output', () => { + const output = 'x'.repeat(100_000); + const estimate = estimateRoutingTokens({ + model: 'gpt-4o', + input: [{ type: 'function_call_output', call_id: 'call-1', output }], + }); + expect(estimate).toBe(Math.round(output.length / 4)); + }); + + it('counts Anthropic tool_result content as before', () => { + const content = 'x'.repeat(1000); + const estimate = estimateRoutingTokens({ + model: 'claude-sonnet-4-20250514', + messages: [ + { role: 'user', content: 'Read file' }, + { role: 'user', content: [{ type: 'tool_result', tool_use_id: 'tool-1', content }] }, + ], + }); + expect(estimate).toBe(Math.round((9 + content.length) / 4)); + }); + it('returns a positive integer (never fractional, never 0 when text exists)', () => { // 1 char / 4 = 0.25 => would round to 0, must floor to 1 const estimate = estimateRoutingTokens({ diff --git a/packages/auto-routing-contracts/src/normalize.ts b/packages/auto-routing-contracts/src/normalize.ts index 10a0e72fb0..e668c85d87 100644 --- a/packages/auto-routing-contracts/src/normalize.ts +++ b/packages/auto-routing-contracts/src/normalize.ts @@ -370,9 +370,8 @@ function collectModalities(value: unknown, out: Set): void { // visited. Bodies use different keys per provider: // * OpenAI chat completions / Anthropic: `messages[]` // * OpenAI Responses: `input` (string or array) - // * Gemini-style: `contents[].parts[]` // Content parts live under `content` (chat / Anthropic) or `parts` - // (Gemini). We do NOT recurse into every key — that would over-walk + // (Responses). We do NOT recurse into every key — that would over-walk // tools, metadata, and provider hints — only into the known structural // containers. for (const key of MODALITY_CONTAINER_KEYS) { @@ -522,12 +521,17 @@ function sumContentPartTextChars(part: unknown): number { return 0; } - // Tool/function result content strings + // Tool/function result content strings. + // Anthropic/Chat Completions tool results use `content`; OpenAI Responses + // `function_call_output` and `tool_call_output` items use `output`. if (type === 'tool_result' || type === 'function_call_output' || type === 'tool_call_output') { - const content = part.content; - if (typeof content === 'string') return content.length; - if (Array.isArray(content)) { - return content.reduce((sum: number, p: unknown) => sum + sumContentPartTextChars(p), 0); + const text = + type === 'function_call_output' || type === 'tool_call_output' + ? (part.output ?? part.content) + : part.content; + if (typeof text === 'string') return text.length; + if (Array.isArray(text)) { + return text.reduce((sum: number, p: unknown) => sum + sumContentPartTextChars(p), 0); } return 0; } diff --git a/services/auto-routing/src/model-capabilities.test.ts b/services/auto-routing/src/model-capabilities.test.ts index c21193110a..9933e1ecfd 100644 --- a/services/auto-routing/src/model-capabilities.test.ts +++ b/services/auto-routing/src/model-capabilities.test.ts @@ -333,6 +333,51 @@ describe('getModelCapabilities', () => { expect(result.get('coding-plan/chat')?.contextLength).toBe(200000); }); + it('distinguishes an unavailable routing table from a genuinely empty one when caching capabilities', async () => { + const put = vi.fn(async () => undefined); + const get = vi.fn(async () => null); + const env = makeEnv(null); + env.AUTO_ROUTING_CONFIG = { get, put } as unknown as KVNamespace; + + // (a) Routing table is unavailable: queryAllIds returns null, so the origin + // value for kvReadThrough is null and the model_capabilities_v1 key is NOT + // written. A later in-memory-miss must still re-check KV and re-fetch origin. + mockGetRoutingTable.mockResolvedValue(null); + const first = await getModelCapabilities(env, { codingPlanModelId: 'coding-plan/chat' }); + expect(first.size).toBe(0); + const capabilityPutsBefore = put.mock.calls.filter( + (call: unknown[]) => call[0] === 'model_capabilities_v1' + ); + expect(capabilityPutsBefore).toEqual([]); + + clearModelCapabilitiesCache(); + const second = await getModelCapabilities(env, { codingPlanModelId: 'coding-plan/chat' }); + expect(second.size).toBe(0); + expect(get).toHaveBeenCalledTimes(2); + const capabilityPutsAfter = put.mock.calls.filter( + (call: unknown[]) => call[0] === 'model_capabilities_v1' + ); + expect(capabilityPutsAfter).toEqual([]); + + // (b) Routing table resolves successfully but has zero candidates: this is + // real data, not a failure, so the empty map IS written to KV. + put.mockClear(); + get.mockClear(); + clearModelCapabilitiesCache(); + clearRoutingTableCache(); + mockGetRoutingTable.mockResolvedValue({ + ...SAMPLE_ROUTING_TABLE, + routes: {}, + }); + const third = await getModelCapabilities(env, { codingPlanModelId: 'coding-plan/chat' }); + expect(third.size).toBe(0); + const capabilityPutsEmpty = (put.mock.calls as unknown[][]).filter( + call => call[0] === 'model_capabilities_v1' + ); + expect(capabilityPutsEmpty).toHaveLength(1); + expect(JSON.parse(capabilityPutsEmpty[0][1] as unknown as string)).toEqual({}); + }); + it('returns an empty map when the routing table is missing entirely', async () => { mockGetRoutingTable.mockResolvedValue(null); const env = { diff --git a/services/auto-routing/src/model-capabilities.ts b/services/auto-routing/src/model-capabilities.ts index 7c9a112e07..d04dcef7a3 100644 --- a/services/auto-routing/src/model-capabilities.ts +++ b/services/auto-routing/src/model-capabilities.ts @@ -165,14 +165,15 @@ async function loadAll(env: ModelCapabilitiesEnv): Promise { +async function queryAllIds(env: ModelCapabilitiesEnv): Promise { const routingTable = await getRoutingTable(env); + if (!routingTable) { + return null; + } const ids = new Set(); - if (routingTable) { - for (const route of Object.values(routingTable.routes)) { - for (const candidate of route) { - ids.add(candidate.model); - } + for (const route of Object.values(routingTable.routes)) { + for (const candidate of route) { + ids.add(candidate.model); } } return queryModelCapabilities(env, Array.from(ids));