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, 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..4cfc21d2e1 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,293 @@ 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']); + }); + + 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', () => { + 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('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({ + 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..e668c85d87 100644 --- a/packages/auto-routing-contracts/src/normalize.ts +++ b/packages/auto-routing-contracts/src/normalize.ts @@ -333,3 +333,262 @@ 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) + // Content parts live under `content` (chat / Anthropic) or `parts` + // (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) { + 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. + // 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 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; + } + } + + // 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; +} 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..9933e1ecfd --- /dev/null +++ b/services/auto-routing/src/model-capabilities.test.ts @@ -0,0 +1,402 @@ +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('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 = { + 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..d04dcef7a3 --- /dev/null +++ b/services/auto-routing/src/model-capabilities.ts @@ -0,0 +1,259 @@ +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); + if (!routingTable) { + return null; + } + const ids = new Set(); + 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(() => {}); + }); +}