feat(auto-routing): modality and context-size aware efficient routing

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 },

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<typeof globalThis.fetch>;
@@ -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);
+  });
 });

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,

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<string, unknown>).futureFlag).toBeUndefined();
+    expect((parsed as Record<string, unknown>).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);
+  });
+});

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<typeof AutoRoutingModeSchema>;
 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<typeof RoutingConstraintsSchema>;
+
 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<typeof MirrorPayloadSchema>;
 
@@ -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';