diff --git a/.kilo_workflow/learnings/auto-routing-benchmark-d1-fixtures.md b/.kilo_workflow/learnings/auto-routing-benchmark-d1-fixtures.md new file mode 100644 index 0000000000..01bdc15ee8 --- /dev/null +++ b/.kilo_workflow/learnings/auto-routing-benchmark-d1-fixtures.md @@ -0,0 +1,20 @@ +# Benchmark profile D1 fixtures for owner-pool E2E (auto-routing) + +Deterministic fixtures for `services/auto-routing-benchmark` local D1 (verified 2026-07-28): + +- Engine identity is deterministic: `computeEngineIdentity('decider')` → `v1:8f2eee90` + (run `pnpm exec tsx -e "import {computeEngineIdentity} from './src/run.ts'; console.log(computeEngineIdentity('decider'))"` + from the service dir). Seed `benchmark_profiles.engine_identity` with exactly this value and + `repetitions` matching `benchmark_config.decider_repetitions`, or rows are stale and auto re-admit. +- `benchmark_config.switch_cost_factor` must be >= 1 (`BenchmarkConfigSchema` zod min 1). A seeded 0.1 + makes `/admin/custom-routing-table` 500 (ZodError at buildCustomRoutingTable), which + `loadCustomRoutingTable` converts to a SILENT null table → decide returns null → gateway serves + the balanced fallback (qwen/qwen3.7-plus) with no product error. Symptom looks like "pool ignored". +- `benchmark_config` alone is not enough: `getBenchmarkConfig` returns null (→ 400 on all profile + endpoints) unless `config_classifier_models` and `config_decider_models` each have >= 1 row. +- Custom table assembly needs BOTH ready `benchmark_profiles` rows (with `run_id`) AND + `model_summaries` rows for that exact run_id + exact (model, variant) pair (provenance binding). +- Profile status lookup auto re-admits stale rows as pending WITHOUT charging + `profile_request_events` — verify ledger count stays flat across engine-drift. +- Generate multi-statement SQL with node, not a shell for-loop: zsh does NOT word-split an unquoted + `$VAR` (`for R in $ROUTES` iterates once with the whole string), producing one garbage row key. diff --git a/.kilo_workflow/learnings/web-gateway-org-context-and-fake-login.md b/.kilo_workflow/learnings/web-gateway-org-context-and-fake-login.md new file mode 100644 index 0000000000..6b5ee5868e --- /dev/null +++ b/.kilo_workflow/learnings/web-gateway-org-context-and-fake-login.md @@ -0,0 +1,22 @@ +# Gateway/auth facts for local efficient-routing E2E (apps/web) + +Verified 2026-07-28 on the efficient-custom-pool-e6e3 worktree: + +- Session-cookie requests NEVER thread organizationId through `getUserFromAuth` (server.ts: org + header is read only in the `Authorization:` bearer path, "only used for extension-originated + requests"). Org-context routing assertions must use the user's own API JWT (rendered in the + profile page API-key textbox) as `Authorization: Bearer` plus `x-kilocode-organizationid`. + Never print the token; extract it in-page via `page.evaluate` and use it there. +- Org-context billing charges the ORGANIZATION balance; personal-context charges the user. Top up + the right entity before real calls. +- Balance top-up needs the denormalized counter: `UPDATE kilocode_users SET + total_microdollars_acquired=...` (same for `organizations`). Inserting only a + `credit_transactions` row does NOT change the balance (402 persists). +- Fake login in a browser: `/users/sign_in?fakeUser=...` 307-redirects to NEXTAUTH_URL, which in an + offset worktree points at a dead port (e.g. 192.168.1.10:3000). The session cookie is set by the + 307 response before the failed redirect — navigate to the target page directly afterwards. + Robust alternative: `GET /api/auth/csrf` then POST `/api/auth/callback/fake-login` + (form-encoded csrfToken + email + json:true) via page.evaluate from any loaded page. +- Org model deny list for fixtures: `organizations.settings.model_deny_list` (jsonb array of model + ids) — immediately makes saved pool entries render `Unavailable` through the web + annotateConfiguredPool path (enterprise orgs). diff --git a/CONTEXT.md b/CONTEXT.md index 8c1545ca03..d204b5cc2c 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -2,7 +2,7 @@ ## Scope -Kilo Code Cloud hosts Kilo Code agents, integrations, and automation. This contract defines Code Reviewer, Security Agent, and Cost Insights language plus ownership boundaries used across review execution, analytics, sync, web, email, remediation, billing alerts, tests, and product documentation. +Kilo Code Cloud hosts Kilo Code agents, integrations, and automation. This contract defines Code Reviewer, Security Agent, Cost Insights, and Auto Routing language plus ownership boundaries used across review execution, analytics, sync, web, email, remediation, billing alerts, model selection, tests, and product documentation. Covered domains must use this contract's canonical terms in code, docs, task descriptions, tests, and agent outputs. Do not introduce a synonym for a governed concept without updating this contract. Scoped `AGENTS.md` files link to this contract instead of copying it. @@ -16,6 +16,7 @@ Covered domains must use this contract's canonical terms in code, docs, task des | **Security Agent Email Delivery** | Dispatch-time revalidation, email rendering, owner-aware links, and Mailgun delivery | `apps/web/src/app/api/internal/security-agent/`, `apps/web/src/lib/email.ts`, `apps/web/src/emails/` | Accepts notification identity only and loads current data before sending | | **Shared Security Notification Policy** | Canonical config parsing, defaults, severity thresholds, and pure event eligibility rules | `packages/worker-utils/src/security-notification-policy.ts` | Web and Worker must use same policy contract | | **Cost Insights** | Spend evidence dashboard, Spend Alerts policy, alert history, and owner-scoped spend alerting | Billing, usage ingestion, usage analytics, and subscription-management surfaces | Applies to both personal users and organizations | +| **Auto Routing** | Efficient model-pool settings, shared benchmark profiles, and per-request model selection | `packages/auto-routing-contracts/`, `services/auto-routing/`, `services/auto-routing-benchmark/`, `apps/web/src/components/auto-routing/` | Personal and organization settings constrain the existing `kilo-auto/efficient` model | ## Canonical Terms @@ -44,6 +45,9 @@ Covered domains must use this contract's canonical terms in code, docs, task des | **Cost Suggestion** | Optional owner-scoped recommendation based on observed Credit spend that may improve cost efficiency through an eligible Coding Plan or Kilo Pass | Referring to recommendation evaluation, dashboard cards, emails, CTA destinations, dismissal, or settings | Alert, warning, guaranteed savings, automatic optimization | | **Suggestion dismissal** | Authorized owner action that hides one Cost Suggestion without changing billing or future suggestion eligibility | Referring to dismissing a recommendation | Alert acknowledgment, unsubscribe, disable suggestions | | **Spend owner** | Personal user or organization whose credit balance is charged for Credit spend | Referring to the Spend Alerts policy and evaluation boundary | Account when personal/org ambiguity matters | +| **Efficient model pool** | The exact model and thinking-variant pairs that `kilo-auto/efficient` may consider for an owner | Referring to either the platform default or an owner-configured pool | Custom mode, custom efficient model | +| **Pool entry** | One concrete model paired with its canonical thinking-variant key, or the model's default only when it exposes no variants | Referring to benchmark and routing identity inside an Efficient model pool | Model alone, normalized effort | +| **Benchmark profile** | Globally shared, current benchmark results for one Pool entry | Referring to readiness and measured routing evidence reused across owners | User benchmark, pool benchmark | | **Spend Anomaly Alert** | Spend Alert triggered when short-window owner Credit spend exceeds that owner's normal usage pattern | Referring to hourly burst-detection Spend Alerts | Low-balance alert, threshold alert | | **Variable Credit spend** | Credit spend created by request-metered product usage such as token usage or metered tool/API usage | Referring to spend that can burst unexpectedly during active usage | Scheduled Credit spend | | **Scheduled Credit spend** | Predictable Credit spend created by subscription-like purchases, renewals, or hosting deductions | Referring to expected recurring or explicitly purchased credit deductions | Variable Credit spend | @@ -67,6 +71,9 @@ Covered domains must use this contract's canonical terms in code, docs, task des - A **Security Remediation** belongs to one **Security Finding** and can have one or more **Security Remediation Attempts**. - A **Security Finding Activity Event** belongs to one Security Agent owner and one Security Finding, including after that finding is deleted. - **Spend Alerts** and **Cost Suggestions** belong to exactly one **Spend owner**: one personal user or one organization. +- An **Efficient model pool** belongs to one personal user or organization; an organization pool overrides members' personal pools, while no configured pool inherits the next available setting and then the platform default. +- A **Pool entry** has at most one current **Benchmark profile** for a benchmark-engine version, shared by every owner that selects that entry. +- `kilo-auto/efficient` considers only ready, request-compatible Pool entries from the effective Efficient model pool; its balanced failure fallback remains outside pool membership. - All Credit spend charged to a **Spend owner** counts toward that owner's Spend Alerts and Cost Suggestion evaluation. - **Cost Suggestions** are enabled by default, independent from Spend Alerts, and recommend an eligible Coding Plan or Kilo Pass when observed Credit spend indicates potential cost-efficiency benefit. - Cost Suggestions are advisory. They do not guarantee savings, automatically purchase or change subscriptions, or alter spend behavior. @@ -189,6 +196,8 @@ Covered domains must use this contract's canonical terms in code, docs, task des - Do not describe a Cost Suggestion as an alert, warning, guaranteed savings, automatic optimization, or required action. - Use **Spend Threshold Alert** and **Spend threshold** for the independent rolling 24-hour, rolling 7-day, and rolling 30-day threshold windows. Do not introduce warning or critical threshold tiers. - Keep Spend Alerts alert-only. Do not describe them as spend blocks, hard limits, pauses, throttles, or request admission controls. +- Use **Efficient model pool** for the candidate set and **Pool entry** for an exact model and canonical thinking variant. Do not introduce a second efficient mode or model ID. +- Treat **Benchmark profiles** as global evidence, never owner-specific benchmark results. ## Ambiguities @@ -215,6 +224,7 @@ Covered domains must use this contract's canonical terms in code, docs, task des - **Shared Security Notification Policy** defines common parsing and pure eligibility behavior; it does not perform persistence or recipient lookup. - Cross-context dispatch sends only stable notification ID from **Security Sync** to authenticated **Security Agent Email Delivery** boundary. - **Spend Alerts** evaluate personal and organization Credit spend at the owner boundary, not per product by default. +- **Auto Routing** owns effective pool resolution and per-request selection. The benchmark worker owns Benchmark profile measurement and publication; web owns catalog validation, permissions, and settings UI. - Do not assume Spend Alerts apply to owners who have not opted in. - Do not make Spend Alerts block, pause, throttle, suppress auto-top-up, reject paid requests, or emit Spend Alerts-specific HTTP 402 responses. - Do not hide v1 Cost Insights behind a release-toggle gate unless a later product decision supersedes public opt-in. diff --git a/apps/web/src/app/api/auto-routing/settings/route.test.ts b/apps/web/src/app/api/auto-routing/settings/route.test.ts new file mode 100644 index 0000000000..e3516e8e97 --- /dev/null +++ b/apps/web/src/app/api/auto-routing/settings/route.test.ts @@ -0,0 +1,706 @@ +import { beforeEach, describe, expect, test } from '@jest/globals'; +import { TRPCError } from '@trpc/server'; +import type { AutoRoutingSettingsResponse } from '@kilocode/auto-routing-contracts'; +import { NextRequest } from 'next/server'; +import { + getAutoRoutingMode, + getAutoRoutingSettings, + updateAutoRoutingSettings, +} from '@/lib/ai-gateway/auto-routing-admin-client'; +import { poolValidationMessage } from '@/lib/ai-gateway/auto-routing-pool-validation'; +import { requireActiveSubscriptionOrTrial } from '@/lib/organizations/trial-middleware'; +import { getUserFromAuth } from '@/lib/user/server'; +import { ensureOrganizationAccess } from '@/routers/organizations/utils'; +import { GET, PUT } from './route'; + +jest.mock('@/lib/ai-gateway/auto-routing-admin-client'); +jest.mock('@/lib/organizations/trial-middleware'); +jest.mock('@/lib/user/server'); +jest.mock('@/routers/organizations/utils'); +jest.mock('@/lib/ai-gateway/providers/openrouter', () => ({ + getEnhancedOpenRouterModels: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/experiments/list-available-experiment-models', () => ({ + listAvailableExperimentModels: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/providers/direct-byok', () => ({ + getDirectByokModelsForUser: jest.fn(), + getDirectByokModelsForOrganization: jest.fn(), +})); +jest.mock('@/lib/organizations/organization-models', () => ({ + getAvailableModelsForOrganization: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/models', () => ({ + kiloExclusiveModels: [ + { public_id: 'kilo/hidden-model', status: 'hidden' }, + { public_id: 'kilo/public-model', status: 'public' }, + ], +})); + +const { getEnhancedOpenRouterModels } = jest.requireMock('@/lib/ai-gateway/providers/openrouter'); +const { listAvailableExperimentModels } = jest.requireMock( + '@/lib/ai-gateway/experiments/list-available-experiment-models' +); +const { getDirectByokModelsForUser, getDirectByokModelsForOrganization } = jest.requireMock( + '@/lib/ai-gateway/providers/direct-byok' +); +const { getAvailableModelsForOrganization } = jest.requireMock( + '@/lib/organizations/organization-models' +); + +const mockedGetAutoRoutingSettings = jest.mocked(getAutoRoutingSettings); +const mockedUpdateAutoRoutingSettings = jest.mocked(updateAutoRoutingSettings); +const mockedGetAutoRoutingMode = jest.mocked(getAutoRoutingMode); +const mockedRequireActiveSubscriptionOrTrial = jest.mocked(requireActiveSubscriptionOrTrial); +const mockedGetUserFromAuth = jest.mocked(getUserFromAuth); +const mockedEnsureOrganizationAccess = jest.mocked(ensureOrganizationAccess); +const mockedGetEnhanced = jest.mocked(getEnhancedOpenRouterModels); +const mockedListExperiments = jest.mocked(listAvailableExperimentModels); +const mockedGetByokUser = jest.mocked(getDirectByokModelsForUser); +const mockedGetByokOrg = jest.mocked(getDirectByokModelsForOrganization); +const mockedGetOrgModels = jest.mocked(getAvailableModelsForOrganization); + +function modeBody( + overrides: Partial<{ + ownerType: 'user' | 'org'; + ownerId: string; + mode: 'cost_per_accuracy' | 'best_accuracy'; + configuredMode: 'cost_per_accuracy' | 'best_accuracy' | null; + defaultMode: 'cost_per_accuracy' | 'best_accuracy'; + }> = {} +) { + return { + ownerType: 'user' as const, + ownerId: USER_ID, + mode: 'best_accuracy' as const, + configuredMode: 'best_accuracy' as const, + defaultMode: 'cost_per_accuracy' as const, + ...overrides, + }; +} + +const USER_ID = 'user-1'; +const ORGANIZATION_ID = 'org-1'; + +function model(id: string, variants?: Record>) { + return { + id, + name: id, + created: 0, + description: id, + architecture: { + input_modalities: ['text'], + output_modalities: ['text'], + tokenizer: 'Other', + }, + top_provider: { is_moderated: false, context_length: 100_000 }, + pricing: { prompt: '0', completion: '0' }, + context_length: 100_000, + ...(variants ? { opencode: { variants } } : {}), + }; +} + +function makeRequest(path: string, body?: unknown) { + return new NextRequest(`http://localhost:3000${path}`, { + method: body === undefined ? 'GET' : 'PUT', + headers: body === undefined ? undefined : { 'content-type': 'application/json' }, + body: body === undefined ? undefined : JSON.stringify(body), + }); +} + +function settingsBody( + overrides: Partial = {} +): AutoRoutingSettingsResponse { + return { + ownerType: 'user', + ownerId: USER_ID, + mode: 'cost_per_accuracy', + configuredMode: null, + defaultMode: 'cost_per_accuracy', + configuredPool: null, + poolStatuses: [], + ...overrides, + }; +} + +describe('/api/auto-routing/settings', () => { + beforeEach(() => { + jest.resetAllMocks(); + mockedGetUserFromAuth.mockResolvedValue({ + user: { id: USER_ID, is_admin: false }, + authFailedResponse: null, + } as never); + mockedListExperiments.mockResolvedValue([]); + mockedGetByokUser.mockResolvedValue([]); + mockedGetByokOrg.mockResolvedValue([]); + mockedGetEnhanced.mockResolvedValue({ + data: [ + model('anthropic/claude-sonnet-4', { high: {}, low: {} }), + model('google/gemini-2.5-flash'), + model('kilo/public-model'), + ], + }); + }); + + test('returns 401 when unauthenticated', async () => { + mockedGetUserFromAuth.mockResolvedValue({ + user: null, + authFailedResponse: new Response(), + } as never); + + const response = await GET(makeRequest('/api/auto-routing/settings')); + + expect(response.status).toBe(401); + await expect(response.json()).resolves.toEqual({ error: 'Authentication required' }); + }); + + test('GET annotates departed pool entries as unavailable without dropping them', async () => { + mockedGetAutoRoutingSettings.mockResolvedValue({ + status: 200, + body: settingsBody({ + configuredPool: [ + { model: 'anthropic/claude-sonnet-4', variant: 'high' }, + { model: 'removed/model', variant: null }, + ], + poolStatuses: [ + { + entry: { model: 'anthropic/claude-sonnet-4', variant: 'high' }, + status: 'ready', + }, + { + entry: { model: 'removed/model', variant: null }, + status: 'ready', + }, + ], + }), + }); + + const response = await GET(makeRequest('/api/auto-routing/settings')); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual( + expect.objectContaining({ + poolSupported: true, + configuredPool: [ + { model: 'anthropic/claude-sonnet-4', variant: 'high', unavailable: false }, + { model: 'removed/model', variant: null, unavailable: true }, + ], + }) + ); + expect(mockedGetAutoRoutingMode).not.toHaveBeenCalled(); + }); + + test('GET falls back to legacy mode when settings endpoint is 404', async () => { + mockedGetAutoRoutingSettings.mockResolvedValue({ + status: 404, + body: { error: 'Not found' }, + }); + mockedGetAutoRoutingMode.mockResolvedValue({ + status: 200, + body: modeBody({ + mode: 'best_accuracy', + configuredMode: 'best_accuracy', + }), + }); + + const response = await GET(makeRequest('/api/auto-routing/settings')); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual({ + ownerType: 'user', + ownerId: USER_ID, + mode: 'best_accuracy', + configuredMode: 'best_accuracy', + defaultMode: 'cost_per_accuracy', + configuredPool: null, + poolStatuses: [], + poolSupported: false, + }); + expect(mockedGetAutoRoutingMode).toHaveBeenCalledWith({ + ownerType: 'user', + ownerId: USER_ID, + }); + // Legacy path must not annotate pool / fetch catalog. + expect(mockedGetEnhanced).not.toHaveBeenCalled(); + }); + + test('GET passes through legacy mode failure status after settings 404', async () => { + mockedGetAutoRoutingSettings.mockResolvedValue({ + status: 404, + body: { error: 'Not found' }, + }); + mockedGetAutoRoutingMode.mockResolvedValue({ + status: 503, + body: { error: 'Worker unavailable' }, + }); + + const response = await GET(makeRequest('/api/auto-routing/settings')); + + expect(response.status).toBe(503); + await expect(response.json()).resolves.toEqual({ error: 'Worker unavailable' }); + }); + + test('rejects org member PUT without owner/billing_manager role with 401', async () => { + // Production ensureOrganizationAccess throws UNAUTHORIZED for insufficient + // role (same as missing membership); trpcErrorResponse maps that to 401. + mockedEnsureOrganizationAccess.mockRejectedValue( + new TRPCError({ + code: 'UNAUTHORIZED', + message: 'You do not have permission to manage this organization', + }) + ); + + const response = await PUT( + makeRequest(`/api/auto-routing/settings?organizationId=${ORGANIZATION_ID}`, { + mode: null, + pool: null, + }) + ); + + expect(response.status).toBe(401); + expect(mockedUpdateAutoRoutingSettings).not.toHaveBeenCalled(); + expect(mockedEnsureOrganizationAccess).toHaveBeenCalledWith( + expect.anything(), + ORGANIZATION_ID, + ['owner', 'billing_manager'] + ); + }); + + test('allows billing_manager organization PUT when entitled', async () => { + mockedEnsureOrganizationAccess.mockResolvedValue(undefined as never); + mockedRequireActiveSubscriptionOrTrial.mockResolvedValue(undefined as never); + mockedGetOrgModels.mockResolvedValue({ + data: [model('google/gemini-2.5-flash')], + }); + mockedUpdateAutoRoutingSettings.mockResolvedValue({ + status: 200, + body: settingsBody({ + ownerType: 'org', + ownerId: ORGANIZATION_ID, + configuredMode: null, + configuredPool: [{ model: 'google/gemini-2.5-flash', variant: null }], + poolStatuses: [ + { + entry: { model: 'google/gemini-2.5-flash', variant: null }, + status: 'pending', + }, + ], + }), + }); + + const response = await PUT( + makeRequest(`/api/auto-routing/settings?organizationId=${ORGANIZATION_ID}`, { + mode: null, + pool: [{ model: 'google/gemini-2.5-flash', variant: null }], + }) + ); + + expect(response.status).toBe(200); + expect(mockedUpdateAutoRoutingSettings).toHaveBeenCalledWith({ + ownerType: 'org', + ownerId: ORGANIZATION_ID, + mode: null, + pool: [{ model: 'google/gemini-2.5-flash', variant: null }], + }); + }); + + test('blocks organization PUT without active subscription or trial', async () => { + mockedEnsureOrganizationAccess.mockResolvedValue(undefined as never); + mockedRequireActiveSubscriptionOrTrial.mockRejectedValue( + new TRPCError({ + code: 'NOT_FOUND', + message: 'Organization subscription not found', + }) + ); + + const response = await PUT( + makeRequest(`/api/auto-routing/settings?organizationId=${ORGANIZATION_ID}`, { + mode: null, + pool: null, + }) + ); + + expect(response.status).toBe(404); + expect(mockedUpdateAutoRoutingSettings).not.toHaveBeenCalled(); + }); + + test('allows personal PUT for the authenticated user', async () => { + mockedUpdateAutoRoutingSettings.mockResolvedValue({ + status: 200, + body: settingsBody({ + configuredMode: 'best_accuracy', + mode: 'best_accuracy', + configuredPool: [{ model: 'google/gemini-2.5-flash', variant: null }], + poolStatuses: [ + { + entry: { model: 'google/gemini-2.5-flash', variant: null }, + status: 'pending', + }, + ], + }), + }); + + const response = await PUT( + makeRequest('/api/auto-routing/settings', { + mode: 'best_accuracy', + pool: [{ model: 'google/gemini-2.5-flash', variant: null }], + }) + ); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual( + expect.objectContaining({ poolSupported: true }) + ); + expect(mockedUpdateAutoRoutingSettings).toHaveBeenCalledWith({ + ownerType: 'user', + ownerId: USER_ID, + mode: 'best_accuracy', + pool: [{ model: 'google/gemini-2.5-flash', variant: null }], + }); + expect(mockedRequireActiveSubscriptionOrTrial).not.toHaveBeenCalled(); + }); + + test.each([ + { + name: 'mode-only clear (pool: null)', + body: { mode: 'best_accuracy' as const, pool: null }, + }, + { + name: 'non-null pool', + body: { + mode: null, + pool: [{ model: 'google/gemini-2.5-flash', variant: null }], + }, + }, + { + name: 'retryEntries', + body: { + mode: null, + pool: [{ model: 'google/gemini-2.5-flash', variant: null }], + retryEntries: [{ model: 'google/gemini-2.5-flash', variant: null }], + }, + }, + ])('PUT with $name returns 503 when settings endpoint is 404', async ({ body }) => { + mockedUpdateAutoRoutingSettings.mockResolvedValue({ + status: 404, + body: { error: 'Not found' }, + }); + + const response = await PUT(makeRequest('/api/auto-routing/settings', body)); + + expect(response.status).toBe(503); + await expect(response.json()).resolves.toEqual({ + error: + 'Custom pools are temporarily unavailable while routing backends update. Try again shortly.', + reason: 'pool_temporarily_unavailable', + }); + expect(mockedUpdateAutoRoutingSettings).toHaveBeenCalled(); + }); + + test('returns 400 for invalid body', async () => { + const response = await PUT( + makeRequest('/api/auto-routing/settings', { + mode: 'not-a-mode', + pool: null, + }) + ); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toEqual({ error: 'Invalid routing settings' }); + expect(mockedUpdateAutoRoutingSettings).not.toHaveBeenCalled(); + }); + + test('maps worker 429 quota errors including retryAt', async () => { + const retryAt = '2026-07-29T12:00:00.000Z'; + mockedUpdateAutoRoutingSettings.mockResolvedValue({ + status: 429, + body: { + error: 'Benchmark profile request limit exceeded', + retryAt, + }, + }); + + const response = await PUT( + makeRequest('/api/auto-routing/settings', { + mode: null, + pool: [{ model: 'google/gemini-2.5-flash', variant: null }], + }) + ); + + expect(response.status).toBe(429); + await expect(response.json()).resolves.toEqual({ + error: 'Benchmark profile request limit exceeded', + retryAt, + }); + }); + + test.each([ + { + name: 'virtual model', + entry: { model: 'kilo-auto/efficient', variant: null }, + reason: 'virtual_model' as const, + setup: () => { + mockedGetEnhanced.mockResolvedValue({ + data: [model('kilo-auto/efficient')], + }); + }, + }, + { + name: 'experiment model', + entry: { model: 'experiment/active', variant: null }, + reason: 'experiment_model' as const, + setup: () => { + mockedListExperiments.mockResolvedValue([model('experiment/active')]); + }, + }, + { + name: 'hidden model', + entry: { model: 'kilo/hidden-model', variant: null }, + reason: 'hidden_model' as const, + setup: () => { + mockedGetEnhanced.mockResolvedValue({ + data: [model('kilo/hidden-model')], + }); + }, + }, + { + name: 'BYOK-only model', + entry: { model: 'openai-codex/gpt-5', variant: null }, + reason: 'byok_only_model' as const, + setup: () => { + mockedGetByokUser.mockResolvedValue([model('openai-codex/gpt-5')]); + }, + }, + ])('rejects $name with a specific 400 message', async ({ entry, reason, setup }) => { + setup(); + + const response = await PUT( + makeRequest('/api/auto-routing/settings', { + mode: null, + pool: [entry], + }) + ); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toEqual( + expect.objectContaining({ + error: poolValidationMessage(reason), + reason, + }) + ); + expect(mockedUpdateAutoRoutingSettings).not.toHaveBeenCalled(); + }); + + test('rejects organization-denied models with a specific 400 message', async () => { + mockedEnsureOrganizationAccess.mockResolvedValue(undefined as never); + mockedRequireActiveSubscriptionOrTrial.mockResolvedValue(undefined as never); + mockedGetEnhanced.mockResolvedValue({ + data: [model('openai/gpt-4o'), model('google/gemini-2.5-flash')], + }); + mockedGetOrgModels.mockResolvedValue({ + data: [model('google/gemini-2.5-flash')], + }); + + const response = await PUT( + makeRequest(`/api/auto-routing/settings?organizationId=${ORGANIZATION_ID}`, { + mode: null, + pool: [{ model: 'openai/gpt-4o', variant: null }], + }) + ); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toEqual( + expect.objectContaining({ + error: poolValidationMessage('organization_denied_model'), + reason: 'organization_denied_model', + }) + ); + }); + + test('rejects missing variant when model exposes variants', async () => { + const response = await PUT( + makeRequest('/api/auto-routing/settings', { + mode: null, + pool: [{ model: 'anthropic/claude-sonnet-4', variant: null }], + }) + ); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toEqual( + expect.objectContaining({ + reason: 'missing_variant', + error: poolValidationMessage('missing_variant'), + }) + ); + }); + + test('rejects unknown variant keys', async () => { + const response = await PUT( + makeRequest('/api/auto-routing/settings', { + mode: null, + pool: [{ model: 'anthropic/claude-sonnet-4', variant: 'max' }], + }) + ); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toEqual( + expect.objectContaining({ + reason: 'unknown_variant', + }) + ); + }); + + test('rejects unexpected variant on models without variants', async () => { + const response = await PUT( + makeRequest('/api/auto-routing/settings', { + mode: null, + pool: [{ model: 'google/gemini-2.5-flash', variant: 'high' }], + }) + ); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toEqual( + expect.objectContaining({ + reason: 'unexpected_variant', + }) + ); + }); + + test('rejects more than 10 entries', async () => { + const pool = Array.from({ length: 11 }, (_, i) => ({ + model: `provider/model-${i}`, + variant: null, + })); + + const response = await PUT( + makeRequest('/api/auto-routing/settings', { + mode: null, + pool, + }) + ); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toEqual( + expect.objectContaining({ + reason: 'too_many_entries', + error: poolValidationMessage('too_many_entries'), + }) + ); + }); + + test('rejects duplicate exact pairs but allows same model with different variants', async () => { + const duplicate = await PUT( + makeRequest('/api/auto-routing/settings', { + mode: null, + pool: [ + { model: 'anthropic/claude-sonnet-4', variant: 'high' }, + { model: 'anthropic/claude-sonnet-4', variant: 'high' }, + ], + }) + ); + expect(duplicate.status).toBe(400); + await expect(duplicate.json()).resolves.toEqual( + expect.objectContaining({ reason: 'duplicate_pair' }) + ); + + mockedUpdateAutoRoutingSettings.mockResolvedValue({ + status: 200, + body: settingsBody({ + configuredPool: [ + { model: 'anthropic/claude-sonnet-4', variant: 'high' }, + { model: 'anthropic/claude-sonnet-4', variant: 'low' }, + ], + poolStatuses: [], + }), + }); + + const allowed = await PUT( + makeRequest('/api/auto-routing/settings', { + mode: null, + pool: [ + { model: 'anthropic/claude-sonnet-4', variant: 'high' }, + { model: 'anthropic/claude-sonnet-4', variant: 'low' }, + ], + }) + ); + expect(allowed.status).toBe(200); + expect(mockedUpdateAutoRoutingSettings).toHaveBeenCalledWith( + expect.objectContaining({ + pool: [ + { model: 'anthropic/claude-sonnet-4', variant: 'high' }, + { model: 'anthropic/claude-sonnet-4', variant: 'low' }, + ], + }) + ); + }); + + test('forwards retryEntries to the worker when present', async () => { + mockedUpdateAutoRoutingSettings.mockResolvedValue({ + status: 200, + body: settingsBody({ + configuredPool: [{ model: 'google/gemini-2.5-flash', variant: null }], + poolStatuses: [ + { + entry: { model: 'google/gemini-2.5-flash', variant: null }, + status: 'pending', + }, + ], + }), + }); + + const response = await PUT( + makeRequest('/api/auto-routing/settings', { + mode: null, + pool: [{ model: 'google/gemini-2.5-flash', variant: null }], + retryEntries: [{ model: 'google/gemini-2.5-flash', variant: null }], + }) + ); + + expect(response.status).toBe(200); + expect(mockedUpdateAutoRoutingSettings).toHaveBeenCalledWith({ + ownerType: 'user', + ownerId: USER_ID, + mode: null, + pool: [{ model: 'google/gemini-2.5-flash', variant: null }], + retryEntries: [{ model: 'google/gemini-2.5-flash', variant: null }], + }); + }); + + test('rejects retryEntries when pool is null before calling the worker', async () => { + const response = await PUT( + makeRequest('/api/auto-routing/settings', { + mode: null, + pool: null, + retryEntries: [{ model: 'google/gemini-2.5-flash', variant: null }], + }) + ); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toEqual( + expect.objectContaining({ + reason: 'invalid_retry_entries', + error: 'retryEntries require a non-null pool', + }) + ); + expect(mockedUpdateAutoRoutingSettings).not.toHaveBeenCalled(); + }); + + test('rejects retryEntries that are not in pool before calling the worker', async () => { + const response = await PUT( + makeRequest('/api/auto-routing/settings', { + mode: null, + pool: [{ model: 'google/gemini-2.5-flash', variant: null }], + retryEntries: [{ model: 'anthropic/claude-sonnet-4', variant: 'high' }], + }) + ); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toEqual( + expect.objectContaining({ + reason: 'invalid_retry_entries', + error: expect.stringContaining('retryEntry must appear in pool'), + }) + ); + expect(mockedUpdateAutoRoutingSettings).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/web/src/app/api/auto-routing/settings/route.ts b/apps/web/src/app/api/auto-routing/settings/route.ts new file mode 100644 index 0000000000..318524b605 --- /dev/null +++ b/apps/web/src/app/api/auto-routing/settings/route.ts @@ -0,0 +1,270 @@ +import { + AutoRoutingModeResponseSchema, + AutoRoutingSettingsResponseSchema, + BenchmarkProfileQuotaErrorSchema, + UpdateAutoRoutingSettingsRequestSchema, + type AutoRoutingModeOwnerType, + type AutoRoutingSettingsResponse, +} from '@kilocode/auto-routing-contracts'; +import { TRPCError } from '@trpc/server'; +import { NextResponse, type NextRequest } from 'next/server'; +import { + getAutoRoutingMode, + getAutoRoutingSettings, + updateAutoRoutingSettings, +} from '@/lib/ai-gateway/auto-routing-admin-client'; +import { + annotateConfiguredPool, + poolValidationMessage, + toApiSettingsResponse, + toLegacyModeApiSettingsResponse, + validatePoolEntries, + type AutoRoutingSettingsApiResponse, +} from '@/lib/ai-gateway/auto-routing-pool-validation'; +import { getUserFromAuth } from '@/lib/user/server'; +import { ensureOrganizationAccess } from '@/routers/organizations/utils'; +import { requireActiveSubscriptionOrTrial } from '@/lib/organizations/trial-middleware'; + +export type { AutoRoutingSettingsApiResponse }; + +const POOL_TEMPORARILY_UNAVAILABLE_MESSAGE = + 'Custom pools are temporarily unavailable while routing backends update. Try again shortly.'; + +function trpcErrorResponse(error: unknown): NextResponse<{ error: string }> | null { + if (!(error instanceof TRPCError)) return null; + const status = + error.code === 'UNAUTHORIZED' + ? 401 + : error.code === 'FORBIDDEN' + ? 403 + : error.code === 'NOT_FOUND' + ? 404 + : 500; + return NextResponse.json({ error: error.message }, { status }); +} + +async function resolveOwner( + request: NextRequest, + roles?: Parameters[2] +): Promise< + | { + ownerType: AutoRoutingModeOwnerType; + ownerId: string; + userId: string; + organizationId: string | null; + } + | { response: NextResponse<{ error: string }> } +> { + const { user, authFailedResponse } = await getUserFromAuth({ adminOnly: false }); + if (!user || authFailedResponse) { + return { response: NextResponse.json({ error: 'Authentication required' }, { status: 401 }) }; + } + + const organizationId = request.nextUrl.searchParams.get('organizationId'); + if (!organizationId) { + return { + ownerType: 'user', + ownerId: user.id, + userId: user.id, + organizationId: null, + }; + } + + try { + await ensureOrganizationAccess({ user }, organizationId, roles); + } catch (error) { + const response = trpcErrorResponse(error); + if (response) return { response }; + throw error; + } + return { + ownerType: 'org', + ownerId: organizationId, + userId: user.id, + organizationId, + }; +} + +async function annotateWorkerSuccess( + status: number, + body: AutoRoutingSettingsResponse, + owner: { userId: string; organizationId: string | null } +): Promise { + const configuredPool = await annotateConfiguredPool({ + userId: owner.userId, + organizationId: owner.organizationId, + configuredPool: body.configuredPool, + }); + const apiBody: AutoRoutingSettingsApiResponse = toApiSettingsResponse(body, configuredPool); + return NextResponse.json(apiBody, { status }); +} + +function workerErrorResponse(result: { status: number; body: unknown }): NextResponse { + if (result.status === 429) { + const quota = BenchmarkProfileQuotaErrorSchema.safeParse(result.body); + if (quota.success) { + return NextResponse.json(quota.data, { status: 429 }); + } + } + return NextResponse.json(result.body, { status: result.status }); +} + +function workerResultResponse( + result: { status: number; body: unknown }, + owner: { userId: string; organizationId: string | null } +): Promise | NextResponse { + if (result.status >= 400) { + return workerErrorResponse(result); + } + + const parsed = AutoRoutingSettingsResponseSchema.safeParse(result.body); + if (!parsed.success) { + return NextResponse.json({ error: 'Invalid worker settings response' }, { status: 502 }); + } + + return annotateWorkerSuccess(result.status, parsed.data, owner); +} + +function legacyModeSuccessResponse(result: { status: number; body: unknown }): NextResponse { + if (result.status >= 400) { + return workerErrorResponse(result); + } + const parsed = AutoRoutingModeResponseSchema.safeParse(result.body); + if (!parsed.success) { + return NextResponse.json({ error: 'Invalid worker mode response' }, { status: 502 }); + } + // No annotateConfiguredPool — pool is null; no catalog fetch. + return NextResponse.json(toLegacyModeApiSettingsResponse(parsed.data), { status: 200 }); +} + +export async function GET(request: NextRequest): Promise { + const owner = await resolveOwner(request); + if ('response' in owner) return owner.response; + + const result = await getAutoRoutingSettings({ + ownerType: owner.ownerType, + ownerId: owner.ownerId, + }); + if (result.status === 404) { + const legacy = await getAutoRoutingMode({ + ownerType: owner.ownerType, + ownerId: owner.ownerId, + }); + return legacyModeSuccessResponse(legacy); + } + return workerResultResponse(result, owner); +} + +export async function PUT(request: NextRequest): Promise { + const owner = await resolveOwner(request, ['owner', 'billing_manager']); + if ('response' in owner) return owner.response; + if (owner.ownerType === 'org') { + try { + await requireActiveSubscriptionOrTrial(owner.ownerId); + } catch (error) { + const response = trpcErrorResponse(error); + if (response) return response; + throw error; + } + } + + let rawBody: unknown; + try { + rawBody = await request.json(); + } catch { + return NextResponse.json({ error: 'Invalid JSON body' }, { status: 400 }); + } + + // Reuse the shared worker contract (includes retryEntries refinements). + // Owner identity is resolved from auth/query, not the client body: spread + // the body first so auth-resolved ownerType/ownerId always win in parsed.data. + const parsed = UpdateAutoRoutingSettingsRequestSchema.safeParse({ + ...(rawBody !== null && typeof rawBody === 'object' ? rawBody : {}), + ownerType: owner.ownerType, + ownerId: owner.ownerId, + }); + if (!parsed.success) { + // Map common pool/retry shape failures to specific messages when possible. + const issue = parsed.error.issues[0]; + if (issue) { + const path = issue.path.join('.'); + if (path.startsWith('retryEntries')) { + return NextResponse.json( + { + error: issue.message, + reason: 'invalid_retry_entries', + }, + { status: 400 } + ); + } + if (path.startsWith('pool') && issue.message.toLowerCase().includes('duplicate')) { + return NextResponse.json( + { error: poolValidationMessage('duplicate_pair'), reason: 'duplicate_pair' }, + { status: 400 } + ); + } + if ( + path === 'pool' && + (issue.code === 'too_big' || issue.message.toLowerCase().includes('at most')) + ) { + return NextResponse.json( + { error: poolValidationMessage('too_many_entries'), reason: 'too_many_entries' }, + { status: 400 } + ); + } + if ( + path === 'pool' && + (issue.code === 'too_small' || issue.message.toLowerCase().includes('at least')) + ) { + return NextResponse.json( + { error: poolValidationMessage('empty_pool'), reason: 'empty_pool' }, + { status: 400 } + ); + } + } + return NextResponse.json({ error: 'Invalid routing settings' }, { status: 400 }); + } + + let pool = parsed.data.pool; + if (pool !== null) { + const validation = await validatePoolEntries({ + user: { id: owner.userId }, + organizationId: owner.organizationId, + entries: pool, + }); + if (!validation.ok) { + return NextResponse.json( + { + error: validation.error.message, + reason: validation.error.reason, + ...(validation.error.index !== undefined ? { index: validation.error.index } : {}), + }, + { status: 400 } + ); + } + pool = validation.entries; + } + + const result = await updateAutoRoutingSettings({ + ownerType: owner.ownerType, + ownerId: owner.ownerId, + mode: parsed.data.mode, + pool, + ...(parsed.data.retryEntries !== undefined ? { retryEntries: parsed.data.retryEntries } : {}), + }); + + // Settings PUT callers always intend pool-aware writes. An old worker that + // 404s cannot honor `pool: null` clear intent (or any pool mutation); return + // a retryable 503 rather than silently falling back to mode-only PUT. + if (result.status === 404) { + return NextResponse.json( + { + error: POOL_TEMPORARILY_UNAVAILABLE_MESSAGE, + reason: 'pool_temporarily_unavailable', + }, + { status: 503 } + ); + } + + return workerResultResponse(result, owner); +} diff --git a/apps/web/src/components/auto-routing/AutoRoutingModeCard.test.ts b/apps/web/src/components/auto-routing/AutoRoutingModeCard.test.ts new file mode 100644 index 0000000000..59cf8a44a6 --- /dev/null +++ b/apps/web/src/components/auto-routing/AutoRoutingModeCard.test.ts @@ -0,0 +1,1189 @@ +import { afterEach, beforeEach, describe, expect, it, jest } from '@jest/globals'; +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; +import { createRequire } from 'node:module'; +import { act, createElement } from 'react'; +import { createRoot, type Root } from 'react-dom/client'; +import { renderToStaticMarkup } from 'react-dom/server'; +import type { BenchmarkProfileEntryStatus, PoolEntry } from '@kilocode/auto-routing-contracts'; + +jest.mock('@/app/api/openrouter/hooks', () => ({ + useModelSelectorList: () => ({ + data: { data: [] }, + isLoading: false, + }), +})); + +jest.mock('sonner', () => ({ + toast: { + success: jest.fn(), + error: jest.fn(), + }, +})); + +import { + applyRetryMutationSuccess, + applySaveMutationError, + applySaveMutationSuccess, + AutoRoutingModeCard, + buildModeSaveBody, + buildRetryBody, + buildSaveBody, + emptyPoolCopy, + formatLoadErrorMessage, + formatSaveErrorMessage, + hasBenchmarkingEntries, + isDirectByokOnlyModel, + isDraftDirty, + isEligiblePoolModel, + isExperimentSelectorModel, + mapPoolEntryDisplayStatus, + modeEndpoint, + NOT_SAVED_ENTRY_LABEL, + ORGANIZATION_EMPTY_POOL_COPY, + PERSONAL_EMPTY_POOL_COPY, + POOL_ROLLOUT_NOTE, + removePoolEntry, + resolveEditableChrome, + resolveEffectiveDraft, + settingsEndpoint, + settingsQueryKey, + settingsRefetchInterval, + SETTINGS_POLL_INTERVAL_MS, + shouldShowClearPoolControl, + toEligibleModelOptions, + tryAddPoolEntry, + UNAVAILABLE_ENTRY_EXPLANATION, + variantLabel, + type AutoRoutingSettingsApiResponse, + type DraftPool, + type PoolEntryWithAvailability, +} from './AutoRoutingModeCard'; + +type LinkedomParseHtml = (html: string) => { + window: typeof globalThis & { + HTMLElement: typeof HTMLElement; + Element: typeof Element; + Node: typeof Node; + Text: typeof Text; + Comment: typeof Comment; + DocumentFragment: typeof DocumentFragment; + Document: typeof Document; + SVGElement: typeof SVGElement; + Event: typeof Event; + CustomEvent: typeof CustomEvent; + navigator: Navigator; + }; + document: Document; +}; + +/** linkedom lives in the monorepo pnpm store (transitive); resolve from this file. */ +function installLinkedomDom(): { cleanup: () => void; container: HTMLElement } { + const requireFromHere = createRequire(__filename); + const loadLinkedom = (): { parseHTML: LinkedomParseHtml } => { + try { + return requireFromHere('linkedom') as { parseHTML: LinkedomParseHtml }; + } catch { + // Transitive monorepo path (not a web package dep): test file → repo root. + return requireFromHere( + '../../../../node_modules/.pnpm/linkedom@0.18.12/node_modules/linkedom' + ) as { parseHTML: LinkedomParseHtml }; + } + }; + const { parseHTML } = loadLinkedom(); + + const { window, document } = parseHTML( + '
' + ); + + const previous = { + window: globalThis.window, + document: globalThis.document, + HTMLElement: globalThis.HTMLElement, + Element: globalThis.Element, + Node: globalThis.Node, + Text: globalThis.Text, + Comment: globalThis.Comment, + DocumentFragment: globalThis.DocumentFragment, + Document: globalThis.Document, + SVGElement: globalThis.SVGElement, + Event: globalThis.Event, + CustomEvent: globalThis.CustomEvent, + navigator: globalThis.navigator, + requestAnimationFrame: globalThis.requestAnimationFrame, + cancelAnimationFrame: globalThis.cancelAnimationFrame, + IS_REACT_ACT_ENVIRONMENT: (globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }) + .IS_REACT_ACT_ENVIRONMENT, + }; + + Object.assign(globalThis, { + window, + document, + HTMLElement: window.HTMLElement, + Element: window.Element, + Node: window.Node, + Text: window.Text, + Comment: window.Comment, + DocumentFragment: window.DocumentFragment, + Document: window.Document, + SVGElement: window.SVGElement, + Event: window.Event, + CustomEvent: window.CustomEvent, + navigator: window.navigator, + requestAnimationFrame: (cb: FrameRequestCallback) => setTimeout(() => cb(Date.now()), 0), + cancelAnimationFrame: (id: number) => clearTimeout(id), + IS_REACT_ACT_ENVIRONMENT: true, + }); + + const container = document.getElementById('root'); + if (!container) throw new Error('linkedom root missing'); + + return { + container: container as unknown as HTMLElement, + cleanup: () => { + Object.assign(globalThis, previous); + }, + }; +} + +// --------------------------------------------------------------------------- +// Fixtures +// --------------------------------------------------------------------------- + +const entryReady: PoolEntry = { model: 'anthropic/claude-sonnet-4', variant: 'high' }; +const entryFailed: PoolEntry = { model: 'google/gemini-2.5-flash', variant: null }; +const entryPending: PoolEntry = { model: 'openai/gpt-5', variant: 'xhigh' }; +const entryUnavailable: PoolEntry = { model: 'removed/model', variant: null }; +const entryDraftOnly: PoolEntry = { model: 'meta-llama/llama-4', variant: null }; + +const statuses: BenchmarkProfileEntryStatus[] = [ + { entry: entryReady, status: 'ready' }, + { + entry: entryFailed, + status: 'failed', + failureReason: 'Benchmark container exited with code 1', + }, + { entry: entryPending, status: 'pending' }, + { entry: entryUnavailable, status: 'ready' }, +]; + +function settings( + overrides: Partial = {} +): AutoRoutingSettingsApiResponse { + return { + ownerType: 'user', + ownerId: 'user-1', + mode: 'cost_per_accuracy', + configuredMode: null, + defaultMode: 'cost_per_accuracy', + configuredPool: null, + poolStatuses: [], + poolSupported: true, + ...overrides, + }; +} +function createTestQueryClient(): QueryClient { + return new QueryClient({ + defaultOptions: { + queries: { retry: false }, + mutations: { retry: false }, + }, + }); +} + +function mountCardHtml(data: AutoRoutingSettingsApiResponse): string { + const queryClient = createTestQueryClient(); + // Seed cache so the first render is the loaded settings state (not skeleton). + queryClient.setQueryData(settingsQueryKey(undefined), data); + + return renderToStaticMarkup( + createElement( + QueryClientProvider, + { client: queryClient }, + createElement(AutoRoutingModeCard, {}) + ) + ); +} + +type MountedCard = { + container: HTMLElement; + root: Root; + cleanup: () => void; +}; + +function mountCardDom(data: AutoRoutingSettingsApiResponse, organizationId?: string): MountedCard { + const dom = installLinkedomDom(); + const queryClient = createTestQueryClient(); + queryClient.setQueryData(settingsQueryKey(organizationId), data); + + let root!: Root; + act(() => { + root = createRoot(dom.container); + root.render( + createElement( + QueryClientProvider, + { client: queryClient }, + createElement(AutoRoutingModeCard, organizationId ? { organizationId } : {}) + ) + ); + }); + + return { + container: dom.container, + root, + cleanup: () => { + act(() => { + root.unmount(); + }); + dom.cleanup(); + }, + }; +} + +function findSaveButton(container: HTMLElement): HTMLButtonElement { + const buttons = Array.from(container.querySelectorAll('button')); + const save = buttons.find(button => button.textContent?.includes('Save auto routing')); + if (!save) throw new Error('Save auto routing button not found'); + return save as HTMLButtonElement; +} + +// --------------------------------------------------------------------------- +// Endpoint / query key +// --------------------------------------------------------------------------- + +describe('settings endpoint and query key', () => { + it('uses the settings route for pool-aware reads/writes', () => { + expect(settingsEndpoint(undefined)).toBe('/api/auto-routing/settings'); + expect(settingsEndpoint('org-1')).toBe('/api/auto-routing/settings?organizationId=org-1'); + expect(settingsEndpoint(undefined)).not.toContain('/mode'); + expect(settingsQueryKey(undefined)).toEqual(['auto-routing-settings', 'personal']); + expect(settingsQueryKey('org-9')).toEqual(['auto-routing-settings', 'org-9']); + }); + + it('exposes the legacy mode endpoint for unsupported-pool saves', () => { + expect(modeEndpoint(undefined)).toBe('/api/auto-routing/mode'); + expect(modeEndpoint('org-1')).toBe('/api/auto-routing/mode?organizationId=org-1'); + }); +}); + +// --------------------------------------------------------------------------- +// Empty state copy (exact strings from the card module) +// --------------------------------------------------------------------------- + +describe('empty / inherited pool copy', () => { + it('uses the exact personal empty string', () => { + expect(PERSONAL_EMPTY_POOL_COPY).toBe( + 'No custom pool. Efficient uses the platform model pool.' + ); + expect(emptyPoolCopy(undefined)).toBe(PERSONAL_EMPTY_POOL_COPY); + }); + + it('uses the exact organization empty string and never implies a merged member pool', () => { + expect(ORGANIZATION_EMPTY_POOL_COPY).toBe( + 'No organization override. Members use their personal pool, or the platform model pool if they have none.' + ); + expect(emptyPoolCopy('org-1')).toBe(ORGANIZATION_EMPTY_POOL_COPY); + expect(ORGANIZATION_EMPTY_POOL_COPY.toLowerCase()).not.toContain('organization owns'); + expect(ORGANIZATION_EMPTY_POOL_COPY.toLowerCase()).not.toContain('merged'); + }); +}); + +// --------------------------------------------------------------------------- +// Status mapping (real helper used by the card rows) +// --------------------------------------------------------------------------- + +describe('mapPoolEntryDisplayStatus', () => { + it('maps ready, pending/running, failed, and unavailable exactly', () => { + expect(mapPoolEntryDisplayStatus({ entry: entryReady, poolStatuses: statuses })).toEqual({ + kind: 'ready', + label: 'Ready', + }); + + expect(mapPoolEntryDisplayStatus({ entry: entryPending, poolStatuses: statuses })).toEqual({ + kind: 'benchmarking', + label: 'Benchmarking', + }); + + expect( + mapPoolEntryDisplayStatus({ + entry: entryPending, + poolStatuses: [{ entry: entryPending, status: 'running' }], + }) + ).toEqual({ kind: 'benchmarking', label: 'Benchmarking' }); + + expect(mapPoolEntryDisplayStatus({ entry: entryFailed, poolStatuses: statuses })).toEqual({ + kind: 'failed', + label: 'Failed', + failureReason: 'Benchmark container exited with code 1', + }); + + expect( + mapPoolEntryDisplayStatus({ + entry: entryUnavailable, + unavailable: true, + poolStatuses: statuses, + }) + ).toEqual({ + kind: 'unavailable', + label: 'Unavailable', + explanation: UNAVAILABLE_ENTRY_EXPLANATION, + }); + }); + + it('treats missing status on a saved entry as Benchmarking', () => { + expect( + mapPoolEntryDisplayStatus({ entry: entryReady, poolStatuses: [], isSaved: true }) + ).toEqual({ + kind: 'benchmarking', + label: 'Benchmarking', + }); + }); + + it('labels unsaved draft-only rows Not saved, never Benchmarking', () => { + const display = mapPoolEntryDisplayStatus({ + entry: entryDraftOnly, + poolStatuses: statuses, + isSaved: false, + }); + expect(display).toEqual({ kind: 'not_saved', label: NOT_SAVED_ENTRY_LABEL }); + expect(display.label).toBe('Not saved'); + expect(display.kind).not.toBe('benchmarking'); + }); + + it('does not map unsaved rows through unavailable/failed even if keys collide in statuses', () => { + const display = mapPoolEntryDisplayStatus({ + entry: entryFailed, + unavailable: true, + poolStatuses: statuses, + isSaved: false, + }); + expect(display.kind).toBe('not_saved'); + expect(display.label).toBe('Not saved'); + }); +}); + +// --------------------------------------------------------------------------- +// Polling stop condition +// --------------------------------------------------------------------------- + +describe('settingsRefetchInterval / hasBenchmarkingEntries', () => { + const configured: PoolEntryWithAvailability[] = [ + { ...entryReady, unavailable: false }, + { ...entryPending, unavailable: false }, + ]; + + it('never polls when poolSupported is false', () => { + expect( + settingsRefetchInterval( + settings({ + poolSupported: false, + configuredPool: configured, + poolStatuses: statuses, + }) + ) + ).toBe(false); + }); + + it('polls while any saved entry is pending/running', () => { + expect(hasBenchmarkingEntries(configured, statuses)).toBe(true); + expect( + settingsRefetchInterval( + settings({ + configuredPool: configured, + poolStatuses: statuses, + }) + ) + ).toBe(SETTINGS_POLL_INTERVAL_MS); + }); + + it('stops polling once every entry is terminal (ready/failed) or unavailable', () => { + const terminalPool: PoolEntryWithAvailability[] = [ + { ...entryReady, unavailable: false }, + { ...entryFailed, unavailable: false }, + { ...entryUnavailable, unavailable: true }, + ]; + const terminalStatuses: BenchmarkProfileEntryStatus[] = [ + { entry: entryReady, status: 'ready' }, + { entry: entryFailed, status: 'failed', failureReason: 'x' }, + { entry: entryUnavailable, status: 'ready' }, + ]; + expect(hasBenchmarkingEntries(terminalPool, terminalStatuses)).toBe(false); + expect( + settingsRefetchInterval( + settings({ configuredPool: terminalPool, poolStatuses: terminalStatuses }) + ) + ).toBe(false); + }); + + it('does not poll for null/empty pools', () => { + expect(hasBenchmarkingEntries(null, [])).toBe(false); + expect(settingsRefetchInterval(settings())).toBe(false); + expect(settingsRefetchInterval(undefined)).toBe(false); + }); +}); + +// --------------------------------------------------------------------------- +// Draft resolution the card executes on every render (poll-safe overrides) +// --------------------------------------------------------------------------- + +describe('resolveEffectiveDraft (card-executed override-over-saved)', () => { + it('keeps overrides when a poll updates the saved snapshot (dirty draft preserved)', () => { + const previousSaved: DraftPool = [entryReady]; + const polledSaved: DraftPool = [entryReady, entryFailed]; + const dirtyPool: DraftPool = [entryReady, entryDraftOnly]; + + // Simulate a refetch that changes savedMode/savedPool while the user holds overrides. + const afterPoll = resolveEffectiveDraft({ + savedMode: 'cost_per_accuracy', + savedPool: polledSaved, + modeOverride: 'best_accuracy', + poolOverride: dirtyPool, + }); + + expect(afterPoll.selectedMode).toBe('best_accuracy'); + expect(afterPoll.draftPool).toEqual(dirtyPool); + expect(afterPoll.draftPool).not.toEqual(polledSaved); + expect(previousSaved).not.toEqual(polledSaved); + + // Dirty relative to the polled server snapshot — Save stays enabled / error panel keeps draft. + expect( + isDraftDirty({ + selectedMode: afterPoll.selectedMode, + draftPool: afterPoll.draftPool, + savedMode: 'cost_per_accuracy', + savedPool: polledSaved, + }) + ).toBe(true); + }); + + it('follows the server when overrides are undefined (clean poll path)', () => { + const previousSaved: DraftPool = [entryReady]; + const polledSaved: DraftPool = [entryReady, entryFailed]; + + const afterPoll = resolveEffectiveDraft({ + savedMode: 'best_accuracy', + savedPool: polledSaved, + modeOverride: undefined, + poolOverride: undefined, + }); + + expect(afterPoll.selectedMode).toBe('best_accuracy'); + expect(afterPoll.draftPool).toEqual(polledSaved); + expect(afterPoll.draftPool).not.toEqual(previousSaved); + expect( + isDraftDirty({ + selectedMode: afterPoll.selectedMode, + draftPool: afterPoll.draftPool, + savedMode: 'best_accuracy', + savedPool: polledSaved, + }) + ).toBe(false); + }); + + it('treats poolOverride null as clear-to-inherit, distinct from undefined (follow server)', () => { + const savedPool: DraftPool = [entryReady]; + + expect( + resolveEffectiveDraft({ + savedMode: 'inherit', + savedPool, + modeOverride: undefined, + poolOverride: null, + }).draftPool + ).toBeNull(); + + expect( + resolveEffectiveDraft({ + savedMode: 'inherit', + savedPool, + modeOverride: undefined, + poolOverride: undefined, + }).draftPool + ).toEqual(savedPool); + }); + + it('save-failure Try again rebuilds the body from the still-held dirty draft', () => { + const savedPool: DraftPool = [entryReady]; + const draftPool: DraftPool = [entryReady, entryFailed]; + const held = resolveEffectiveDraft({ + savedMode: 'cost_per_accuracy', + savedPool, + modeOverride: 'best_accuracy', + poolOverride: draftPool, + }); + + expect( + isDraftDirty({ + selectedMode: held.selectedMode, + draftPool: held.draftPool, + savedMode: 'cost_per_accuracy', + savedPool, + }) + ).toBe(true); + + // Retry uses the same buildSaveBody draft the user still holds. + expect( + buildSaveBody({ + mode: held.selectedMode, + pool: held.draftPool, + }) + ).toEqual({ + mode: 'best_accuracy', + pool: draftPool, + }); + }); +}); + +// --------------------------------------------------------------------------- +// Save / retry mutation callbacks the card wires into useMutation +// --------------------------------------------------------------------------- + +describe('applySaveMutationSuccess / applySaveMutationError', () => { + it('on save success updates the query cache and clears the save-error panel', () => { + const setQueryData = jest.fn(); + const setSaveError = jest.fn(); + const setRetryingKey = jest.fn(); + const markClearOverridesAfterSave = jest.fn(); + const setModeOverride = jest.fn(); + const setPoolOverride = jest.fn(); + const toastSuccess = jest.fn(); + const queryKey = settingsQueryKey(undefined); + const data = settings({ + configuredMode: 'best_accuracy', + configuredPool: [{ ...entryReady, unavailable: false }], + poolStatuses: [{ entry: entryReady, status: 'ready' }], + }); + + applySaveMutationSuccess({ + queryClient: { setQueryData }, + queryKey, + data, + setSaveError, + setRetryingKey, + markClearOverridesAfterSave, + setModeOverride, + setPoolOverride, + toastSuccess, + }); + + expect(setQueryData).toHaveBeenCalledWith(queryKey, data); + expect(setSaveError).toHaveBeenCalledWith(null); + expect(setRetryingKey).toHaveBeenCalledWith(null); + expect(markClearOverridesAfterSave).toHaveBeenCalled(); + expect(setModeOverride).toHaveBeenCalledWith('best_accuracy'); + expect(setPoolOverride).toHaveBeenCalledWith([entryReady]); + expect(toastSuccess).toHaveBeenCalledWith('Auto routing settings saved'); + }); + + it('on save error sets the inline Try again message and does not touch draft overrides', () => { + const setRetryingKey = jest.fn(); + const setSaveError = jest.fn(); + const toastError = jest.fn(); + + applySaveMutationError({ + error: new Error('Invalid routing settings'), + setRetryingKey, + setSaveError, + toastError, + }); + + expect(setRetryingKey).toHaveBeenCalledWith(null); + expect(setSaveError).toHaveBeenCalledWith('Invalid routing settings'); + expect(toastError).toHaveBeenCalledWith('Invalid routing settings'); + }); +}); + +describe('applyRetryMutationSuccess', () => { + it('invalidates the settings query key immediately after a successful retry', () => { + const setQueryData = jest.fn(); + const invalidateQueries = jest.fn(); + const setRetryingKey = jest.fn(); + const toastSuccess = jest.fn(); + const queryKey = settingsQueryKey('org-1'); + const data = settings({ + ownerType: 'org', + ownerId: 'org-1', + configuredPool: [{ ...entryFailed, unavailable: false }], + poolStatuses: [{ entry: entryFailed, status: 'pending' }], + }); + + applyRetryMutationSuccess({ + queryClient: { setQueryData, invalidateQueries }, + queryKey, + data, + setRetryingKey, + toastSuccess, + }); + + expect(setQueryData).toHaveBeenCalledWith(queryKey, data); + expect(invalidateQueries).toHaveBeenCalledWith({ queryKey }); + expect(invalidateQueries).toHaveBeenCalledWith({ + queryKey: ['auto-routing-settings', 'org-1'], + }); + expect(setRetryingKey).toHaveBeenCalledWith(null); + expect(toastSuccess).toHaveBeenCalledWith('Benchmark retry requested'); + }); +}); + +// --------------------------------------------------------------------------- +// Add flow +// --------------------------------------------------------------------------- + +describe('tryAddPoolEntry', () => { + it('requires a variant when the model exposes variants', () => { + const result = tryAddPoolEntry({ + draftPool: null, + modelId: 'anthropic/claude-sonnet-4', + variant: undefined, + modelVariants: ['low', 'high'], + }); + expect(result).toEqual({ + ok: false, + reason: 'missing_variant', + message: 'Choose a variant for this model.', + }); + }); + + it('adds a unique pair and rejects duplicates', () => { + const first = tryAddPoolEntry({ + draftPool: null, + modelId: 'anthropic/claude-sonnet-4', + variant: 'high', + modelVariants: ['low', 'high'], + }); + expect(first.ok).toBe(true); + if (!first.ok) return; + + const duplicate = tryAddPoolEntry({ + draftPool: first.pool, + modelId: 'anthropic/claude-sonnet-4', + variant: 'high', + modelVariants: ['low', 'high'], + }); + expect(duplicate).toEqual({ + ok: false, + reason: 'duplicate', + message: 'That model and variant pair is already in the pool.', + }); + + const otherVariant = tryAddPoolEntry({ + draftPool: first.pool, + modelId: 'anthropic/claude-sonnet-4', + variant: 'low', + modelVariants: ['low', 'high'], + }); + expect(otherVariant.ok).toBe(true); + }); + + it('allows null variant when the model has no variants', () => { + const result = tryAddPoolEntry({ + draftPool: null, + modelId: 'google/gemini-2.5-flash', + variant: undefined, + modelVariants: undefined, + }); + expect(result).toEqual({ + ok: true, + pool: [{ model: 'google/gemini-2.5-flash', variant: null }], + }); + }); + + it('enforces the 10-entry cap', () => { + const pool: PoolEntry[] = Array.from({ length: 10 }, (_, i) => ({ + model: `provider/model-${i}`, + variant: null, + })); + const result = tryAddPoolEntry({ + draftPool: pool, + modelId: 'provider/extra', + variant: null, + modelVariants: [], + }); + expect(result.ok).toBe(false); + if (result.ok) return; + expect(result.reason).toBe('pool_full'); + }); +}); + +// --------------------------------------------------------------------------- +// Eligibility +// --------------------------------------------------------------------------- + +describe('isDirectByokOnlyModel / isExperimentSelectorModel / isEligiblePoolModel', () => { + it('excludes direct-BYOK-only entries via hasUserByokAvailable + provider prefix', () => { + expect( + isDirectByokOnlyModel({ + id: 'chutes-byok/some-model', + hasUserByokAvailable: true, + }) + ).toBe(true); + expect( + isEligiblePoolModel({ + id: 'chutes-byok/some-model', + hasUserByokAvailable: true, + }) + ).toBe(false); + + // Managed OpenRouter model with a user BYOK key remains eligible. + expect( + isDirectByokOnlyModel({ + id: 'anthropic/claude-sonnet-4', + hasUserByokAvailable: true, + }) + ).toBe(false); + expect( + isEligiblePoolModel({ + id: 'anthropic/claude-sonnet-4', + hasUserByokAvailable: true, + isFree: false, + pricing: { prompt: '0.000003' }, + }) + ).toBe(true); + }); + + it('excludes experiment selector entries via zero pricing without isFree', () => { + const experiment = { + id: 'partner/preview-model', + isFree: undefined, + pricing: { prompt: '0.0000000' }, + }; + expect(isExperimentSelectorModel(experiment)).toBe(true); + expect(isEligiblePoolModel(experiment)).toBe(false); + + // Managed free models set isFree: true and must stay eligible. + expect( + isExperimentSelectorModel({ + id: 'openrouter/free-ish', + isFree: true, + pricing: { prompt: '0' }, + }) + ).toBe(false); + expect( + isEligiblePoolModel({ + id: 'openrouter/free-ish', + isFree: true, + pricing: { prompt: '0' }, + }) + ).toBe(true); + }); + + it('still excludes virtual and custom ids', () => { + expect(isEligiblePoolModel({ id: 'anthropic/claude' })).toBe(true); + expect(isEligiblePoolModel({ id: 'kilo-auto/efficient' })).toBe(false); + expect(isEligiblePoolModel({ id: 'kilo-internal/x' })).toBe(false); + }); +}); + +describe('toEligibleModelOptions', () => { + it('excludes kilo-auto, custom LLM, BYOK-only, experiments, and pairs already in the draft', () => { + const options = toEligibleModelOptions( + [ + { id: 'kilo-auto/efficient', name: 'Efficient' }, + { id: 'kilo-internal/custom', name: 'Custom' }, + { + id: 'chutes-byok/direct-only', + name: 'Direct BYOK', + hasUserByokAvailable: true, + }, + { + id: 'partner/preview-model', + name: 'Experiment', + pricing: { prompt: '0.0000000' }, + }, + { + id: 'anthropic/claude-sonnet-4', + name: 'Sonnet', + hasUserByokAvailable: true, + pricing: { prompt: '0.000003' }, + opencode: { variants: { high: {}, low: {} } }, + }, + { + id: 'google/gemini-2.5-flash', + name: 'Flash', + pricing: { prompt: '0.0000001' }, + }, + ], + [{ model: 'google/gemini-2.5-flash', variant: null }] + ); + expect(options.map(o => o.id)).toEqual(['anthropic/claude-sonnet-4']); + expect(options[0]?.variants).toEqual(['high', 'low']); + }); + + it('keeps a multi-variant model when only one variant is drafted', () => { + const options = toEligibleModelOptions( + [ + { + id: 'anthropic/claude-sonnet-4', + name: 'Sonnet', + pricing: { prompt: '0.000003' }, + opencode: { variants: { high: {}, low: {} } }, + }, + ], + [{ model: 'anthropic/claude-sonnet-4', variant: 'high' }] + ); + expect(options).toHaveLength(1); + }); +}); + +// --------------------------------------------------------------------------- +// Save / retry bodies +// --------------------------------------------------------------------------- + +describe('buildSaveBody / buildRetryBody / buildModeSaveBody', () => { + it('PUTs { mode, pool } with nulls for inherit (happy save payload)', () => { + expect(buildSaveBody({ mode: 'inherit', pool: null })).toEqual({ + mode: null, + pool: null, + }); + expect( + buildSaveBody({ + mode: 'best_accuracy', + pool: [entryReady], + }) + ).toEqual({ + mode: 'best_accuracy', + pool: [entryReady], + }); + }); + + it('buildModeSaveBody is mode-only (no pool fields)', () => { + expect(buildModeSaveBody('best_accuracy')).toEqual({ mode: 'best_accuracy' }); + expect(buildModeSaveBody('inherit')).toEqual({ mode: null }); + expect(buildModeSaveBody('best_accuracy')).not.toHaveProperty('pool'); + }); + + it('retry PUTs the current saved pool with retryEntries: [entry]', () => { + const body = buildRetryBody({ + mode: 'cost_per_accuracy', + savedPool: [entryReady, entryFailed], + retryEntry: entryFailed, + }); + expect(body).toEqual({ + mode: 'cost_per_accuracy', + pool: [entryReady, entryFailed], + retryEntries: [entryFailed], + }); + expect(body.retryEntries).toHaveLength(1); + expect(body.retryEntries[0]).toEqual(entryFailed); + }); +}); + +describe('removePoolEntry / clear', () => { + it('removes one entry and clears to inherit when the last entry is removed', () => { + const withTwo = [entryReady, entryFailed]; + expect(removePoolEntry(withTwo, entryReady)).toEqual([entryFailed]); + expect(removePoolEntry([entryFailed], entryFailed)).toBeNull(); + }); +}); + +// --------------------------------------------------------------------------- +// Clear pool + editable chrome (card-executed visibility) +// --------------------------------------------------------------------------- + +describe('shouldShowClearPoolControl / resolveEditableChrome', () => { + it('shows Clear pool only when a saved configured pool exists', () => { + expect(shouldShowClearPoolControl(true)).toBe(true); + expect(shouldShowClearPoolControl(false)).toBe(false); + + expect( + resolveEditableChrome({ + readonly: false, + hasConfiguredPool: true, + hasSaveError: false, + }).showClearPool + ).toBe(true); + expect( + resolveEditableChrome({ + readonly: false, + hasConfiguredPool: false, + hasSaveError: false, + }).showClearPool + ).toBe(false); + }); + + it('save-error panel exposes Try again only when editable', () => { + expect( + resolveEditableChrome({ + readonly: false, + hasConfiguredPool: false, + hasSaveError: true, + }).showSaveErrorRetry + ).toBe(true); + expect( + resolveEditableChrome({ + readonly: true, + hasConfiguredPool: true, + hasSaveError: true, + }).showSaveErrorRetry + ).toBe(false); + }); + + it('org member readonly hides Save / Add / Remove / Retry / Clear', () => { + const chrome = resolveEditableChrome({ + readonly: true, + hasConfiguredPool: true, + hasSaveError: true, + }); + expect(chrome).toEqual({ + showClearPool: false, + showAddModel: false, + showSave: false, + showSaveErrorRetry: false, + showRemove: false, + showRetryBenchmarkForFailed: false, + }); + }); +}); + +// --------------------------------------------------------------------------- +// Errors +// --------------------------------------------------------------------------- + +describe('error messages', () => { + it('load failure message surfaces the API error string', () => { + expect(formatLoadErrorMessage({ error: 'Authentication required' })).toBe( + 'Authentication required' + ); + expect(formatLoadErrorMessage(null)).toBe('Failed to load auto routing settings'); + }); + + it('save failure surfaces a specific message for the inline Try again path', () => { + const message = formatSaveErrorMessage({ error: 'Invalid routing settings' }, 400); + expect(message).toBe('Invalid routing settings'); + }); + + it('429 message includes retry timing from retryAt', () => { + const retryAt = '2026-07-29T12:00:00.000Z'; + const message = formatSaveErrorMessage( + { + error: 'Benchmark profile request limit exceeded', + retryAt, + }, + 429 + ); + expect(message).toContain('Benchmark profile request limit exceeded'); + expect(message).toContain('New benchmarks can be requested'); + expect(message.toLowerCase()).toMatch(/after /); + }); +}); + +// --------------------------------------------------------------------------- +// Readonly / permission isolation (org member) — call-site derivation +// --------------------------------------------------------------------------- + +describe('readonly org member controls', () => { + it('billing-manager edit permission is isolated from page-level canEdit', () => { + // Mirrors OrganizationProvidersAndModelsPage derivation. + const derive = (role: string, isKiloAdmin = false) => { + const canEdit = isKiloAdmin || role === 'owner'; + const canEditAutoRouting = isKiloAdmin || role === 'owner' || role === 'billing_manager'; + return { canEdit, canEditAutoRouting }; + }; + + expect(derive('billing_manager')).toEqual({ + canEdit: false, + canEditAutoRouting: true, + }); + expect(derive('member')).toEqual({ + canEdit: false, + canEditAutoRouting: false, + }); + expect(derive('owner')).toEqual({ + canEdit: true, + canEditAutoRouting: true, + }); + expect(derive('member', true)).toEqual({ + canEdit: true, + canEditAutoRouting: true, + }); + + // Card consumes readonly={!canEditAutoRouting}; member → no edit chrome. + expect( + resolveEditableChrome({ + readonly: !derive('member').canEditAutoRouting, + hasConfiguredPool: true, + hasSaveError: false, + }).showSave + ).toBe(false); + }); +}); + +// --------------------------------------------------------------------------- +// Variant labels used by real card rows +// --------------------------------------------------------------------------- + +describe('variantLabel', () => { + it('labels null as Default and known effort keys for display', () => { + expect(variantLabel(null)).toBe('Default'); + expect(variantLabel('high')).toBe('High'); + }); +}); + +// --------------------------------------------------------------------------- +// Mounted card: poolSupported === false (legacy worker fallback) +// --------------------------------------------------------------------------- + +describe('AutoRoutingModeCard poolSupported=false', () => { + const originalFetch = global.fetch; + + beforeEach(() => { + jest.clearAllMocks(); + }); + + afterEach(() => { + global.fetch = originalFetch; + }); + + it('renders the rollout note and no pool controls; disables polling', () => { + const data = settings({ + poolSupported: false, + configuredMode: 'cost_per_accuracy', + mode: 'cost_per_accuracy', + }); + const html = mountCardHtml(data); + + expect(html).toContain(POOL_ROLLOUT_NOTE); + expect(html).not.toContain('Add model'); + expect(html).not.toContain('Clear pool'); + expect(html).not.toContain(PERSONAL_EMPTY_POOL_COPY); + expect(html).not.toContain('Retry benchmark'); + expect(html).toContain('Routing mode'); + expect(html).toContain('Save auto routing'); + expect(html).toContain('Efficient model pool'); + + // Card wires settingsRefetchInterval into useQuery — must not poll. + expect( + settingsRefetchInterval( + settings({ + poolSupported: false, + configuredPool: [{ ...entryPending, unavailable: false }], + poolStatuses: [{ entry: entryPending, status: 'pending' }], + }) + ) + ).toBe(false); + }); + + it('Save control PUTs { mode } to the legacy mode endpoint, not settings', async () => { + const fetchMock = jest.fn(async (input: RequestInfo | URL, init?: RequestInit) => { + const url = typeof input === 'string' ? input : input.toString(); + if (url.includes('/api/auto-routing/mode') && init?.method === 'PUT') { + return { + ok: true, + status: 200, + json: async () => ({ + ownerType: 'user', + ownerId: 'user-1', + mode: 'best_accuracy', + configuredMode: 'best_accuracy', + defaultMode: 'cost_per_accuracy', + }), + } as Response; + } + // invalidate/refetch after legacy save + if (url.includes('/api/auto-routing/settings')) { + return { + ok: true, + status: 200, + json: async () => + settings({ + poolSupported: false, + configuredMode: 'best_accuracy', + mode: 'best_accuracy', + }), + } as Response; + } + throw new Error(`Unexpected fetch: ${url} ${init?.method ?? 'GET'}`); + }); + global.fetch = fetchMock as typeof fetch; + + // Saved mode is inherit (configuredMode null). Dirty by selecting best_accuracy + // so the real Save control is enabled, then click it. + const data = settings({ + poolSupported: false, + configuredMode: null, + mode: 'cost_per_accuracy', + }); + const mounted = mountCardDom(data); + + try { + await act(async () => { + // Radix Select stores onValueChange on the provider; linkedom cannot open + // the portal menu, so drive the same change handler the Select wires up + // by finding the combobox trigger's React props is unreliable. Instead + // use the Select root via React fiber is fragile — set the mode through + // the same public path as Select: dispatch on the trigger after finding + // the hidden select value change. Practically: click is enough if we + // force-enable by changing draft via a second approach — + // query the SelectItem buttons after opening is hard without pointer + // events. Use React internals on the Select: the Select root is a + // context provider with value/onValueChange. Walk fibers: + type Fiber = { + memoizedProps?: { onValueChange?: (v: string) => void; value?: string }; + child?: Fiber | null; + sibling?: Fiber | null; + }; + const reactKey = Object.keys(mounted.container).find(key => + key.startsWith('__reactContainer') + ); + const host = reactKey + ? (mounted.container as unknown as Record)[ + reactKey + ] + : undefined; + const walk = (fiber: Fiber | null | undefined, visit: (f: Fiber) => void) => { + if (!fiber) return; + visit(fiber); + walk(fiber.child, visit); + walk(fiber.sibling, visit); + }; + let onValueChange: ((v: string) => void) | undefined; + const rootFiber = host?.stateNode?.current; + walk(rootFiber, fiber => { + const props = fiber.memoizedProps; + if ( + props && + typeof props.onValueChange === 'function' && + (props.value === 'inherit' || props.value === undefined) + ) { + onValueChange = props.onValueChange; + } + }); + if (!onValueChange) { + throw new Error('Select onValueChange not found on fiber tree'); + } + onValueChange('best_accuracy'); + }); + + const saveButton = findSaveButton(mounted.container); + expect(saveButton.disabled).toBe(false); + + await act(async () => { + saveButton.click(); + }); + + await act(async () => { + await Promise.resolve(); + await Promise.resolve(); + await Promise.resolve(); + }); + + const putCalls = fetchMock.mock.calls.filter( + ([, init]) => init && typeof init === 'object' && init.method === 'PUT' + ); + expect(putCalls).toHaveLength(1); + const [putUrl, putInit] = putCalls[0]!; + expect(String(putUrl)).toBe(modeEndpoint(undefined)); + expect(String(putUrl)).not.toContain('/settings'); + expect(putInit?.method).toBe('PUT'); + const putBody = JSON.parse(String(putInit?.body)); + expect(putBody).toEqual({ mode: 'best_accuracy' }); + expect(putBody).not.toHaveProperty('pool'); + expect(putBody).not.toHaveProperty('retryEntries'); + } finally { + mounted.cleanup(); + } + }); +}); + +beforeEach(() => { + jest.clearAllMocks?.(); +}); diff --git a/apps/web/src/components/auto-routing/AutoRoutingModeCard.tsx b/apps/web/src/components/auto-routing/AutoRoutingModeCard.tsx index 4c6c7b2bfb..dfeb154241 100644 --- a/apps/web/src/components/auto-routing/AutoRoutingModeCard.tsx +++ b/apps/web/src/components/auto-routing/AutoRoutingModeCard.tsx @@ -2,13 +2,25 @@ import { AutoRoutingModeSchema, - AutoRoutingModeResponseSchema, + isVirtualAutoModelId, + MAX_POOL_ENTRIES, + poolEntryKey, type AutoRoutingMode, + type BenchmarkProfileEntryStatus, + type BenchmarkProfileQuotaError, + type PoolEntry, } from '@kilocode/auto-routing-contracts'; import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'; -import { Route } from 'lucide-react'; -import { useEffect, useState } from 'react'; +import { Plus, Route, Trash2, X } from 'lucide-react'; +// Default `React` is required by Jest's SWC classic JSX transform (see apps/web/jest.config.ts). +// tsconfig uses "jsx": "react-jsx" for production; changing the Jest transform is out of this +// slice's owned paths. Keep the default until Jest is switched to the automatic runtime. +import React, { useEffect, useId, useMemo, useRef, useState } from 'react'; import { toast } from 'sonner'; +import { useModelSelectorList } from '@/app/api/openrouter/hooks'; +import { ModelCombobox, type ModelOption } from '@/components/shared/ModelCombobox'; +import { VariantCombobox } from '@/components/shared/VariantCombobox'; +import { Alert, AlertDescription, AlertTitle } from '@/components/ui/alert'; import { Button } from '@/components/ui/button'; import { Card, CardContent, CardDescription, CardHeader, CardTitle } from '@/components/ui/card'; import { Label } from '@/components/ui/label'; @@ -19,13 +31,112 @@ import { SelectTrigger, SelectValue, } from '@/components/ui/select'; +import { Skeleton } from '@/components/ui/skeleton'; +import { thinkingEffortLabel } from '@/lib/code-reviews/core/model-variants'; +import { CUSTOM_LLM_PREFIX } from '@/lib/ai-gateway/model-utils'; +import { DIRECT_BYOK_PROVIDERS_META } from '@/lib/ai-gateway/providers/direct-byok/direct-byok-meta'; +import { cn } from '@/lib/utils'; + +// --------------------------------------------------------------------------- +// Types (mirror the web settings API response; client-safe, no server-only import) +// --------------------------------------------------------------------------- + +export type PoolEntryWithAvailability = PoolEntry & { + unavailable: boolean; +}; + +export type AutoRoutingSettingsApiResponse = { + ownerType: 'user' | 'org'; + ownerId: string; + mode: AutoRoutingMode; + configuredMode: AutoRoutingMode | null; + defaultMode: AutoRoutingMode; + configuredPool: PoolEntryWithAvailability[] | null; + poolStatuses: BenchmarkProfileEntryStatus[]; + /** + * False when the BFF fell back to legacy mode-only workers (deploy window / + * rollback). Pool editor is replaced by an informational note; saves go to + * the legacy mode endpoint (`{ mode }` only) so a server-side pool is never + * cleared or silently preserved via a mismatched settings PUT. + */ + poolSupported: boolean; +}; + +export type ModeSelection = AutoRoutingMode | 'inherit'; + +export type DraftPool = PoolEntry[] | null; + +export type PoolEntryDisplayStatus = + | { kind: 'ready'; label: 'Ready' } + | { kind: 'benchmarking'; label: 'Benchmarking' } + | { kind: 'failed'; label: 'Failed'; failureReason: string | null } + | { kind: 'unavailable'; label: 'Unavailable'; explanation: string } + | { kind: 'not_saved'; label: 'Not saved' }; + +export const NOT_SAVED_ENTRY_LABEL = 'Not saved' as const; type Props = { organizationId?: string; readonly?: boolean; }; -type ModeSelection = AutoRoutingMode | 'inherit'; +// --------------------------------------------------------------------------- +// Copy (exact strings from plan task 5.3) +// --------------------------------------------------------------------------- + +export const PERSONAL_EMPTY_POOL_COPY = 'No custom pool. Efficient uses the platform model pool.'; + +export const ORGANIZATION_EMPTY_POOL_COPY = + 'No organization override. Members use their personal pool, or the platform model pool if they have none.'; + +export const UNAVAILABLE_ENTRY_EXPLANATION = + 'This model or variant is no longer available in your catalog and cannot be used for routing.'; + +export const SETTINGS_POLL_INTERVAL_MS = 15_000; + +export const POOL_ROLLOUT_NOTE = 'Custom pools are being rolled out and will be available shortly.'; + +// --------------------------------------------------------------------------- +// Pure helpers (exported for focused component tests) +// --------------------------------------------------------------------------- + +export function settingsEndpoint(organizationId: string | undefined): string { + if (!organizationId) return '/api/auto-routing/settings'; + const params = new URLSearchParams({ organizationId }); + return `/api/auto-routing/settings?${params}`; +} + +/** + * Legacy mode-only BFF used when `poolSupported === false`. Mode PUT never + * touches pool keys at any worker version. + */ +export function modeEndpoint(organizationId: string | undefined): string { + if (!organizationId) return '/api/auto-routing/mode'; + const params = new URLSearchParams({ organizationId }); + return `/api/auto-routing/mode?${params}`; +} + +export function settingsQueryKey(organizationId: string | undefined) { + return ['auto-routing-settings', organizationId ?? 'personal'] as const; +} + +export function emptyPoolCopy(organizationId: string | undefined): string { + return organizationId ? ORGANIZATION_EMPTY_POOL_COPY : PERSONAL_EMPTY_POOL_COPY; +} + +export function unsetModeOption(organizationId: string | undefined) { + return organizationId + ? { + value: 'inherit' as const, + label: 'No organization override', + description: "Uses the member's personal setting, then the default.", + } + : { + value: 'inherit' as const, + label: 'Use default setting', + description: 'Uses best accuracy per dollar.', + }; +} const modeOptions: Array<{ value: AutoRoutingMode; label: string; description: string }> = [ { @@ -41,88 +152,841 @@ const modeOptions: Array<{ value: AutoRoutingMode; label: string; description: s }, ]; -function unsetModeOption(organizationId: string | undefined) { - return organizationId - ? { - value: 'inherit' as const, - label: 'No organization override', - description: "Uses the member's personal setting, then the default.", - } - : { - value: 'inherit' as const, - label: 'Use default setting', - description: 'Uses best accuracy per dollar.', - }; +export function findStatusForEntry( + entry: PoolEntry, + poolStatuses: BenchmarkProfileEntryStatus[] +): BenchmarkProfileEntryStatus | undefined { + const key = poolEntryKey(entry); + return poolStatuses.find(status => poolEntryKey(status.entry) === key); } -function endpoint(organizationId: string | undefined): string { - if (!organizationId) return '/api/auto-routing/mode'; - const params = new URLSearchParams({ organizationId }); - return `/api/auto-routing/mode?${params}`; +export function mapPoolEntryDisplayStatus(params: { + entry: PoolEntry; + unavailable?: boolean; + poolStatuses: BenchmarkProfileEntryStatus[]; + /** When false, the pair exists only in the local draft and has no saved status. */ + isSaved?: boolean; +}): PoolEntryDisplayStatus { + if (params.isSaved === false) { + return { kind: 'not_saved', label: NOT_SAVED_ENTRY_LABEL }; + } + + if (params.unavailable) { + return { + kind: 'unavailable', + label: 'Unavailable', + explanation: UNAVAILABLE_ENTRY_EXPLANATION, + }; + } + + const status = findStatusForEntry(params.entry, params.poolStatuses); + if (!status || status.status === 'pending' || status.status === 'running') { + return { kind: 'benchmarking', label: 'Benchmarking' }; + } + if (status.status === 'failed') { + return { + kind: 'failed', + label: 'Failed', + failureReason: status.failureReason ?? null, + }; + } + return { kind: 'ready', label: 'Ready' }; } -async function fetchMode(organizationId: string | undefined) { - const response = await fetch(endpoint(organizationId)); - const body: unknown = await response.json(); - if (!response.ok) { - throw new Error( +/** + * True when the draft differs from the last server snapshot (mode and/or pool). + * Used so polls/refetches do not clobber unsaved edits or an active save error. + */ +export function isDraftDirty(params: { + selectedMode: ModeSelection; + draftPool: DraftPool; + savedMode: ModeSelection; + savedPool: DraftPool; +}): boolean { + return ( + params.selectedMode !== params.savedMode || !draftsEqual(params.draftPool, params.savedPool) + ); +} + +/** + * Override-over-saved draft resolution the card uses on every render. + * `undefined` overrides mean "follow the latest server snapshot", so polls + * update statuses/saved pool without clobbering unsaved edits. + */ +export function resolveEffectiveDraft(params: { + savedMode: ModeSelection; + savedPool: DraftPool; + modeOverride: ModeSelection | undefined; + poolOverride: DraftPool | undefined; +}): { selectedMode: ModeSelection; draftPool: DraftPool } { + return { + selectedMode: params.modeOverride ?? params.savedMode, + draftPool: params.poolOverride !== undefined ? params.poolOverride : params.savedPool, + }; +} + +/** Clear pool is offered only when a saved configured pool exists (not draft-only). */ +export function shouldShowClearPoolControl(hasConfiguredPool: boolean): boolean { + return hasConfiguredPool; +} + +/** + * Edit-action visibility for the card chrome and pool rows. + * Members (`readonly`) never see Save / Add / Remove / Retry / Clear. + */ +export function resolveEditableChrome(params: { + readonly: boolean; + hasConfiguredPool: boolean; + hasSaveError: boolean; +}): { + showClearPool: boolean; + showAddModel: boolean; + showSave: boolean; + showSaveErrorRetry: boolean; + showRemove: boolean; + showRetryBenchmarkForFailed: boolean; +} { + const editable = !params.readonly; + return { + showClearPool: editable && shouldShowClearPoolControl(params.hasConfiguredPool), + showAddModel: editable, + showSave: editable, + showSaveErrorRetry: editable && params.hasSaveError, + showRemove: editable, + showRetryBenchmarkForFailed: editable, + }; +} + +/** True when any saved (non-unavailable) entry is still pending/running. */ +export function hasBenchmarkingEntries( + configuredPool: PoolEntryWithAvailability[] | null, + poolStatuses: BenchmarkProfileEntryStatus[] +): boolean { + if (!configuredPool || configuredPool.length === 0) return false; + return configuredPool.some(entry => { + if (entry.unavailable) return false; + const display = mapPoolEntryDisplayStatus({ + entry, + unavailable: entry.unavailable, + poolStatuses, + }); + return display.kind === 'benchmarking'; + }); +} + +/** + * React Query refetchInterval: poll while any saved entry is Benchmarking; + * stop once every entry is terminal (ready/failed) or unavailable. + * Never poll when the worker does not support pools (legacy fallback). + */ +export function settingsRefetchInterval( + data: AutoRoutingSettingsApiResponse | undefined +): number | false { + if (!data) return false; + if (data.poolSupported === false) return false; + return hasBenchmarkingEntries(data.configuredPool, data.poolStatuses) + ? SETTINGS_POLL_INTERVAL_MS + : false; +} + +export function draftsEqual(a: DraftPool, b: DraftPool): boolean { + if (a === null && b === null) return true; + if (a === null || b === null) return false; + if (a.length !== b.length) return false; + return a.every((entry, index) => { + const other = b[index]; + return other !== undefined && poolEntryKey(entry) === poolEntryKey(other); + }); +} + +export function buildSaveBody(params: { mode: ModeSelection; pool: DraftPool }): { + mode: AutoRoutingMode | null; + pool: PoolEntry[] | null; +} { + return { + mode: params.mode === 'inherit' ? null : params.mode, + pool: params.pool, + }; +} + +/** Mode-only body for the legacy `/api/auto-routing/mode` PUT. */ +export function buildModeSaveBody(mode: ModeSelection): { mode: AutoRoutingMode | null } { + return { mode: mode === 'inherit' ? null : mode }; +} + +export function buildRetryBody(params: { + mode: ModeSelection; + savedPool: PoolEntry[]; + retryEntry: PoolEntry; +}): { mode: AutoRoutingMode | null; pool: PoolEntry[]; retryEntries: PoolEntry[] } { + return { + mode: params.mode === 'inherit' ? null : params.mode, + pool: params.savedPool, + retryEntries: [params.retryEntry], + }; +} + +export type AddPoolEntryResult = + | { ok: true; pool: PoolEntry[] } + | { + ok: false; + reason: 'duplicate' | 'missing_variant' | 'pool_full' | 'no_model'; + message: string; + }; + +export function tryAddPoolEntry(params: { + draftPool: DraftPool; + modelId: string | undefined; + variant: string | null | undefined; + modelVariants: string[] | undefined; +}): AddPoolEntryResult { + const modelId = params.modelId?.trim(); + if (!modelId) { + return { ok: false, reason: 'no_model', message: 'Choose a model to add.' }; + } + + const hasVariants = (params.modelVariants?.length ?? 0) > 0; + const variant = + params.variant === undefined || params.variant === null || params.variant === '' + ? null + : params.variant; + + if (hasVariants && variant === null) { + return { + ok: false, + reason: 'missing_variant', + message: 'Choose a variant for this model.', + }; + } + if (!hasVariants && variant !== null) { + // Models without variants always store null. + } + + const entry: PoolEntry = { + model: modelId, + variant: hasVariants ? variant : null, + }; + + const current = params.draftPool ?? []; + if (current.length >= MAX_POOL_ENTRIES) { + return { + ok: false, + reason: 'pool_full', + message: `An Efficient pool can include at most ${MAX_POOL_ENTRIES} models.`, + }; + } + + const key = poolEntryKey(entry); + if (current.some(existing => poolEntryKey(existing) === key)) { + return { + ok: false, + reason: 'duplicate', + message: 'That model and variant pair is already in the pool.', + }; + } + + return { ok: true, pool: [...current, entry] }; +} + +export function removePoolEntry(draftPool: DraftPool, entry: PoolEntry): DraftPool { + if (!draftPool) return draftPool; + const key = poolEntryKey(entry); + const next = draftPool.filter(existing => poolEntryKey(existing) !== key); + return next.length === 0 ? null : next; +} + +const DIRECT_BYOK_PROVIDER_PREFIXES = new Set( + Object.keys(DIRECT_BYOK_PROVIDERS_META).map(id => `${id}/`) +); + +/** + * Direct-BYOK-only catalog entries use a direct-BYOK provider id prefix + * (`chutes-byok/…`, `kimi-coding/…`, …) and are always marked + * `hasUserByokAvailable: true` in the selector list. Managed OpenRouter models + * may also carry `hasUserByokAvailable` when the user has a matching key — those + * stay eligible. Server revalidates regardless. + */ +export function isDirectByokOnlyModel(model: { + id: string; + hasUserByokAvailable?: boolean; +}): boolean { + if (model.hasUserByokAvailable !== true) return false; + const slash = model.id.indexOf('/'); + if (slash <= 0) return false; + const prefix = model.id.slice(0, slash + 1); + return DIRECT_BYOK_PROVIDER_PREFIXES.has(prefix); +} + +/** + * Experiment public ids appear in the selector list with ordinary partner ids + * (not `experiment/…`). `listAvailableExperimentModels` always sets zero + * pricing and omits `isFree`; managed free models set `isFree: true`. Combined + * with zero prompt pricing this is the client-visible experiment signal. + * Server remains authoritative. + */ +export function isExperimentSelectorModel(model: { + id: string; + isFree?: boolean; + pricing?: { prompt?: string } | null; +}): boolean { + const id = model.id; + if (!id) return false; + if (id.includes('/experiment') || id.startsWith('experiment/')) return true; + if (model.isFree === true) return false; + const prompt = model.pricing?.prompt; + if (typeof prompt !== 'string') return false; + const amount = Number.parseFloat(prompt); + return Number.isFinite(amount) && amount === 0; +} + +/** Client-side usability filter; server revalidates on save. */ +export function isEligiblePoolModel(model: { + id: string; + name?: string; + isFree?: boolean; + hasUserByokAvailable?: boolean; + pricing?: { prompt?: string } | null; +}): boolean { + const id = model.id; + if (!id) return false; + if (isVirtualAutoModelId(id)) return false; + if (id.startsWith(CUSTOM_LLM_PREFIX)) return false; + if (isDirectByokOnlyModel(model)) return false; + if (isExperimentSelectorModel(model)) return false; + return true; +} + +export type SelectorListModel = { + id: string; + name: string; + isFree?: boolean; + mayTrainOnYourPrompts?: boolean; + hasUserByokAvailable?: boolean; + pricing?: { prompt?: string } | null; + opencode?: { variants?: Record } | null; +}; + +export function toEligibleModelOptions( + models: SelectorListModel[], + draftPool: DraftPool +): ModelOption[] { + const draftKeys = new Set((draftPool ?? []).map(poolEntryKey)); + return models + .filter(model => isEligiblePoolModel(model)) + .map(model => { + const variantKeys = model.opencode?.variants + ? Object.keys(model.opencode.variants).filter(key => key.trim().length > 0) + : []; + return { + id: model.id, + name: model.name, + isFree: model.isFree, + mayTrainOnYourPrompts: model.mayTrainOnYourPrompts, + hasUserByokAvailable: model.hasUserByokAvailable, + variants: variantKeys.length > 0 ? variantKeys : undefined, + }; + }) + .filter(model => { + // Keep models that still have at least one (model, variant) pair not in the draft. + // Models with no variants are excluded only when the null-variant pair is already drafted. + if (!model.variants || model.variants.length === 0) { + return !draftKeys.has(poolEntryKey({ model: model.id, variant: null })); + } + return model.variants.some( + variant => !draftKeys.has(poolEntryKey({ model: model.id, variant })) + ); + }); +} + +export function formatSaveErrorMessage(body: unknown, status: number): string { + if (status === 429) { + const retryAt = + body && typeof body === 'object' && 'retryAt' in body && typeof body.retryAt === 'string' + ? body.retryAt + : null; + const base = body && typeof body === 'object' && 'error' in body && typeof body.error === 'string' ? body.error - : 'Failed to load auto routing mode' - ); + : 'Benchmark request limit reached.'; + if (retryAt) { + const when = formatRetryAt(retryAt); + return `${base} New benchmarks can be requested ${when}.`; + } + return base; + } + + if (body && typeof body === 'object' && 'error' in body && typeof body.error === 'string') { + return body.error; + } + return 'Failed to save auto routing settings'; +} + +export function formatRetryAt(iso: string): string { + const date = new Date(iso); + if (Number.isNaN(date.getTime())) return `after ${iso}`; + return `after ${date.toLocaleString(undefined, { + dateStyle: 'medium', + timeStyle: 'short', + })}`; +} + +export function formatLoadErrorMessage(body: unknown): string { + if (body && typeof body === 'object' && 'error' in body && typeof body.error === 'string') { + return body.error; + } + return 'Failed to load auto routing settings'; +} + +export function variantLabel(variant: string | null): string { + if (variant === null) return 'Default'; + return thinkingEffortLabel(variant); +} + +/** + * Save mutation `onSuccess` side effects the card wires into `useMutation`. + * Keeps the committed response in the query cache and clears the save-error panel. + */ +export function applySaveMutationSuccess(params: { + queryClient: { setQueryData: (key: readonly unknown[], data: unknown) => void }; + queryKey: readonly unknown[]; + data: AutoRoutingSettingsApiResponse; + setSaveError: (message: string | null) => void; + setRetryingKey: (key: string | null) => void; + markClearOverridesAfterSave: () => void; + setModeOverride: (mode: ModeSelection) => void; + setPoolOverride: (pool: DraftPool) => void; + toastSuccess: (message: string) => void; +}): void { + params.queryClient.setQueryData(params.queryKey, params.data); + params.setSaveError(null); + params.setRetryingKey(null); + params.markClearOverridesAfterSave(); + params.setModeOverride(params.data.configuredMode ?? 'inherit'); + params.setPoolOverride(savedPoolSnapshot(params.data.configuredPool)); + params.toastSuccess('Auto routing settings saved'); +} + +/** + * Save mutation `onError` side effects: surface the message for the inline + * Try again panel and toast (draft overrides are intentionally left alone). + */ +export function applySaveMutationError(params: { + error: unknown; + setRetryingKey: (key: string | null) => void; + setSaveError: (message: string) => void; + toastError: (message: string) => void; +}): void { + params.setRetryingKey(null); + const message = + params.error instanceof Error ? params.error.message : 'Failed to save auto routing settings'; + params.setSaveError(message); + params.toastError(message); +} + +/** + * Retry-benchmark mutation `onSuccess`: cache the response and invalidate so + * polling/status refresh picks up the re-queued profile immediately. + */ +export function applyRetryMutationSuccess(params: { + queryClient: { + setQueryData: (key: readonly unknown[], data: unknown) => void; + invalidateQueries: (filters: { queryKey: readonly unknown[] }) => unknown; + }; + queryKey: readonly unknown[]; + data: AutoRoutingSettingsApiResponse; + setRetryingKey: (key: string | null) => void; + toastSuccess: (message: string) => void; +}): void { + params.queryClient.setQueryData(params.queryKey, params.data); + void params.queryClient.invalidateQueries({ queryKey: params.queryKey }); + params.setRetryingKey(null); + params.toastSuccess('Benchmark retry requested'); +} + +function savedPoolSnapshot(configuredPool: PoolEntryWithAvailability[] | null): DraftPool { + if (!configuredPool) return null; + return configuredPool.map(({ model, variant }) => ({ model, variant })); +} + +function availabilityByKey( + configuredPool: PoolEntryWithAvailability[] | null +): Map { + const map = new Map(); + if (!configuredPool) return map; + for (const entry of configuredPool) { + map.set(poolEntryKey(entry), entry.unavailable); } - return AutoRoutingModeResponseSchema.parse(body); + return map; } -async function saveMode(organizationId: string | undefined, mode: AutoRoutingMode | null) { - const response = await fetch(endpoint(organizationId), { +// --------------------------------------------------------------------------- +// Fetch / mutate +// --------------------------------------------------------------------------- + +async function fetchSettings( + organizationId: string | undefined +): Promise { + const response = await fetch(settingsEndpoint(organizationId)); + const body: unknown = await response.json().catch(() => null); + if (!response.ok) { + throw new Error(formatLoadErrorMessage(body)); + } + return body as AutoRoutingSettingsApiResponse; +} + +function throwSaveError(body: unknown, status: number): never { + const message = formatSaveErrorMessage(body, status); + const error = new Error(message) as Error & { + status?: number; + retryAt?: string; + quota?: BenchmarkProfileQuotaError; + }; + error.status = status; + if (body && typeof body === 'object' && 'retryAt' in body && typeof body.retryAt === 'string') { + error.retryAt = body.retryAt; + } + throw error; +} + +async function putSettings( + organizationId: string | undefined, + payload: { + mode: AutoRoutingMode | null; + pool: PoolEntry[] | null; + retryEntries?: PoolEntry[]; + } +): Promise { + const response = await fetch(settingsEndpoint(organizationId), { method: 'PUT', headers: { 'content-type': 'application/json' }, - body: JSON.stringify({ mode }), + body: JSON.stringify(payload), }); - const body: unknown = await response.json(); + const body: unknown = await response.json().catch(() => null); if (!response.ok) { - throw new Error( - body && typeof body === 'object' && 'error' in body && typeof body.error === 'string' - ? body.error - : 'Failed to save auto routing mode' - ); + throwSaveError(body, response.status); + } + return body as AutoRoutingSettingsApiResponse; +} + +/** + * Legacy mode-only save. On success the caller invalidates the settings query + * so GET re-synthesizes `{ poolSupported: false, … }` from the mode endpoint. + */ +async function putMode( + organizationId: string | undefined, + payload: { mode: AutoRoutingMode | null } +): Promise { + const response = await fetch(modeEndpoint(organizationId), { + method: 'PUT', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify(payload), + }); + const body: unknown = await response.json().catch(() => null); + if (!response.ok) { + throwSaveError(body, response.status); } - return AutoRoutingModeResponseSchema.parse(body); } +// --------------------------------------------------------------------------- +// Component +// --------------------------------------------------------------------------- + export function AutoRoutingModeCard({ organizationId, readonly = false }: Props) { const queryClient = useQueryClient(); - const queryKey = ['auto-routing-mode', organizationId ?? 'personal']; + const queryKey = settingsQueryKey(organizationId); + const idPrefix = useId(); + const modeFieldId = `${idPrefix}-mode`; + const modeHelpId = `${idPrefix}-mode-help`; + const poolHelpId = `${idPrefix}-pool-help`; + const addErrorId = `${idPrefix}-add-error`; + const loadErrorId = `${idPrefix}-load-error`; + const saveErrorId = `${idPrefix}-save-error`; + const addModelSectionRef = useRef(null); + const query = useQuery({ queryKey, - queryFn: () => fetchMode(organizationId), + queryFn: () => fetchSettings(organizationId), + refetchInterval: query => settingsRefetchInterval(query.state.data), + refetchOnWindowFocus: true, + }); + + const modelsQuery = useModelSelectorList(organizationId); + + // Local draft overrides. `undefined` means "follow the latest server snapshot". + // Polls/refetches update statuses via query.data without touching these. + const [modeOverride, setModeOverride] = useState(undefined); + const [poolOverride, setPoolOverride] = useState(undefined); + const [addModelId, setAddModelId] = useState(undefined); + const [addVariant, setAddVariant] = useState(undefined); + const [addError, setAddError] = useState(null); + const [showAddFlow, setShowAddFlow] = useState(false); + const [saveError, setSaveError] = useState(null); + const [retryingKey, setRetryingKey] = useState(null); + // After a successful save, drop overrides once the query reflects the new snapshot. + const clearOverridesAfterSaveRef = useRef(false); + + const savedMode: ModeSelection = query.data?.configuredMode ?? 'inherit'; + const savedPool = useMemo( + () => savedPoolSnapshot(query.data?.configuredPool ?? null), + [query.data?.configuredPool] + ); + const availabilityMap = useMemo( + () => availabilityByKey(query.data?.configuredPool ?? null), + [query.data?.configuredPool] + ); + const poolStatuses = query.data?.poolStatuses ?? []; + + const { selectedMode, draftPool } = resolveEffectiveDraft({ + savedMode, + savedPool, + modeOverride, + poolOverride, }); - const [selectedMode, setSelectedMode] = useState('inherit'); - const currentSelection: ModeSelection = query.data?.configuredMode ?? 'inherit'; useEffect(() => { - setSelectedMode(currentSelection); - }, [currentSelection]); + if (!clearOverridesAfterSaveRef.current) return; + if (!query.data) return; + // Drop local overrides so the draft tracks the post-save server snapshot + // (and subsequent clean polls). Do not clear while a later edit is dirty. + const nextSavedMode: ModeSelection = query.data.configuredMode ?? 'inherit'; + const nextSavedPool = savedPoolSnapshot(query.data.configuredPool); + const stillDirty = + (modeOverride !== undefined && modeOverride !== nextSavedMode) || + (poolOverride !== undefined && !draftsEqual(poolOverride, nextSavedPool)); + if (stillDirty) return; + clearOverridesAfterSaveRef.current = false; + setModeOverride(undefined); + setPoolOverride(undefined); + }, [query.data, modeOverride, poolOverride]); - const mutation = useMutation({ - mutationFn: (mode: ModeSelection) => saveMode(organizationId, mode === 'inherit' ? null : mode), + const saveMutation = useMutation({ + mutationFn: async (payload: { + /** When true, write mode via the legacy mode endpoint (no pool keys). */ + useLegacyModeEndpoint: boolean; + mode: AutoRoutingMode | null; + pool: PoolEntry[] | null; + retryEntries?: PoolEntry[]; + }): Promise => { + if (payload.useLegacyModeEndpoint) { + await putMode(organizationId, { mode: payload.mode }); + return 'legacy-mode-saved'; + } + return putSettings(organizationId, { + mode: payload.mode, + pool: payload.pool, + ...(payload.retryEntries !== undefined ? { retryEntries: payload.retryEntries } : {}), + }); + }, + onSuccess: data => { + if (data === 'legacy-mode-saved') { + // Refetch settings so GET fallback re-synthesizes poolSupported: false. + setSaveError(null); + setRetryingKey(null); + clearOverridesAfterSaveRef.current = true; + void queryClient.invalidateQueries({ queryKey }); + toast.success('Auto routing settings saved'); + return; + } + // Align overrides with the saved response immediately so the UI shows + // the committed pool; the effect above clears them once query.data matches. + applySaveMutationSuccess({ + queryClient, + queryKey, + data, + setSaveError, + setRetryingKey, + markClearOverridesAfterSave: () => { + clearOverridesAfterSaveRef.current = true; + }, + setModeOverride, + setPoolOverride, + toastSuccess: message => toast.success(message), + }); + }, + onError: error => { + applySaveMutationError({ + error, + setRetryingKey, + setSaveError, + toastError: message => toast.error(message), + }); + }, + }); + + const retryMutation = useMutation({ + mutationFn: (payload: { + mode: AutoRoutingMode | null; + pool: PoolEntry[]; + retryEntries: PoolEntry[]; + }) => putSettings(organizationId, payload), onSuccess: data => { - queryClient.setQueryData(queryKey, data); - toast.success('Auto routing mode saved'); + applyRetryMutationSuccess({ + queryClient, + queryKey, + data, + setRetryingKey, + toastSuccess: message => toast.success(message), + }); }, onError: error => { - toast.error(error instanceof Error ? error.message : 'Failed to save auto routing mode'); + setRetryingKey(null); + const message = error instanceof Error ? error.message : 'Failed to retry benchmark'; + toast.error(message); }, }); + const eligibleModels = useMemo( + () => toEligibleModelOptions(modelsQuery.data?.data ?? [], draftPool), + [modelsQuery.data?.data, draftPool] + ); + + const selectedAddModel = eligibleModels.find(model => model.id === addModelId); + const addModelVariants = selectedAddModel?.variants ?? []; + const resetOption = unsetModeOption(organizationId); const selectedOption = selectedMode === 'inherit' ? resetOption : (modeOptions.find(option => option.value === selectedMode) ?? modeOptions[0]); - const disabled = readonly || query.isLoading || mutation.isPending; - const hasChanges = selectedMode !== currentSelection; + + const controlsDisabled = readonly || query.isLoading || saveMutation.isPending; + const hasChanges = isDraftDirty({ + selectedMode, + draftPool, + savedMode, + savedPool, + }); + const hasConfiguredPool = savedPool !== null && savedPool.length > 0; + const draftHasEntries = draftPool !== null && draftPool.length > 0; + const editableChrome = resolveEditableChrome({ + readonly, + hasConfiguredPool, + hasSaveError: saveError !== null, + }); + const focusAddFlow = () => { + setShowAddFlow(true); + setAddError(null); + // Defer focus until the add controls mount. + requestAnimationFrame(() => { + addModelSectionRef.current + ?.querySelector('button[role="combobox"], button') + ?.focus(); + }); + }; + + const handleAdd = () => { + const result = tryAddPoolEntry({ + draftPool, + modelId: addModelId, + variant: addVariant, + modelVariants: addModelVariants, + }); + if (!result.ok) { + setAddError(result.message); + return; + } + setPoolOverride(result.pool); + setAddModelId(undefined); + setAddVariant(undefined); + setAddError(null); + setShowAddFlow(false); + }; + + const handleRemove = (entry: PoolEntry) => { + setPoolOverride(removePoolEntry(draftPool, entry)); + setAddError(null); + }; + + const handleClearPool = () => { + setPoolOverride(null); + setAddError(null); + }; + + const poolSupported = query.data?.poolSupported !== false; + + const handleSave = () => { + setSaveError(null); + if (!poolSupported) { + // Mode-only legacy endpoint never touches pool keys (any worker version). + const { mode } = buildModeSaveBody(selectedMode); + saveMutation.mutate({ + useLegacyModeEndpoint: true, + mode, + pool: null, + }); + return; + } + const body = buildSaveBody({ mode: selectedMode, pool: draftPool }); + saveMutation.mutate({ + useLegacyModeEndpoint: false, + mode: body.mode, + pool: body.pool, + }); + }; + + const handleRetryBenchmark = (entry: PoolEntry) => { + if (!savedPool) return; + setRetryingKey(poolEntryKey(entry)); + retryMutation.mutate( + buildRetryBody({ + mode: savedMode, + savedPool, + retryEntry: entry, + }) + ); + }; + + const poolSectionDescription = emptyPoolCopy(organizationId); + + if (query.isLoading && !query.data) { + return ( + + + + + Auto routing + + + Choose how Kilo ranks models for kilo-auto/efficient and which models belong in the + Efficient pool. + + + + + + + + + ); + } + + if (query.isError && !query.data) { + const message = + query.error instanceof Error ? query.error.message : formatLoadErrorMessage(null); + return ( + + + + + Auto routing + + + Choose how Kilo ranks models for kilo-auto/efficient and which models belong in the + Efficient pool. + + + + + Could not load auto routing settings + {message} + + + + + ); + } return ( @@ -131,21 +995,22 @@ export function AutoRoutingModeCard({ organizationId, readonly = false }: Props) Auto routing - Choose how Kilo ranks models for kilo-auto/efficient. + + Choose how Kilo ranks models for kilo-auto/efficient and which models belong in the + Efficient pool. + - +
- + -

{selectedOption.description}

+

+ {selectedOption.description} +

+
+ +
+
+
+

Efficient model pool

+ {poolSupported ? ( +

+ Up to {MAX_POOL_ENTRIES} exact model and variant pairs. Leave empty to inherit. +

+ ) : null} +
+ {poolSupported && editableChrome.showAddModel && ( +
+ {editableChrome.showClearPool ? ( + + ) : null} + +
+ )} +
+ + {!poolSupported ? ( +
+

{POOL_ROLLOUT_NOTE}

+
+ ) : !draftHasEntries ? ( +
+

{poolSectionDescription}

+ {editableChrome.showAddModel && ( + + )} +
+ ) : ( +
    + {(draftPool ?? []).map(entry => { + const unavailable = availabilityMap.get(poolEntryKey(entry)) ?? false; + const isSaved = + savedPool?.some(saved => poolEntryKey(saved) === poolEntryKey(entry)) ?? false; + const display = mapPoolEntryDisplayStatus({ + entry, + unavailable, + poolStatuses, + isSaved, + }); + + return ( +
  • +
    +

    {entry.model}

    +

    + Variant: {variantLabel(entry.variant)} +

    +

    + + {display.label} + + {display.kind === 'failed' && display.failureReason ? ( + + {display.failureReason} + + ) : null} + {display.kind === 'unavailable' ? ( + + {display.explanation} + + ) : null} +

    +
    +
    + {display.kind === 'failed' && + editableChrome.showRetryBenchmarkForFailed && + isSaved ? ( + + ) : null} + {editableChrome.showRemove ? ( + + ) : null} +
    +
    • + ); + })} +
+ )} + + {poolSupported && showAddFlow && editableChrome.showAddModel ? ( +
+
+

Add model

+ +
+
+
+ { + setAddModelId(value); + setAddVariant(undefined); + setAddError(null); + }} + isLoading={modelsQuery.isLoading} + disabled={controlsDisabled} + required + error={addError && !addModelId ? addError : undefined} + /> +
+ {addModelVariants.length > 0 ? ( +
+ + { + setAddVariant(value); + setAddError(null); + }} + disabled={controlsDisabled || !addModelId} + className="w-full" + /> +
+ ) : null} + +
+ {addError ? ( + + ) : null} +
+ ) : null}
- {!readonly && ( + + {saveError ? ( + + Could not save auto routing settings + + {saveError} + {editableChrome.showSaveErrorRetry ? ( + + ) : null} + + + ) : null} + + {editableChrome.showSave && ( )} diff --git a/apps/web/src/components/organizations/providers-and-models/OrganizationProvidersAndModelsPage.tsx b/apps/web/src/components/organizations/providers-and-models/OrganizationProvidersAndModelsPage.tsx index 21ecdfcc57..21e80e1edd 100644 --- a/apps/web/src/components/organizations/providers-and-models/OrganizationProvidersAndModelsPage.tsx +++ b/apps/web/src/components/organizations/providers-and-models/OrganizationProvidersAndModelsPage.tsx @@ -96,6 +96,10 @@ export function OrganizationProvidersAndModelsPage({ organizationId, role }: Pro const isKiloAdmin = assumedRole === 'KILO ADMIN'; const currentRole = (isKiloAdmin ? 'owner' : assumedRole) ?? role; const canEdit = isKiloAdmin || currentRole === 'owner'; + // Auto routing card: owners, billing managers, and platform admins may edit. + // Do not widen page-level canEdit (provider allow-list stays owner/admin only). + const canEditAutoRouting = + isKiloAdmin || currentRole === 'owner' || currentRole === 'billing_manager'; const updateOrganizationSettings = useUpdateOrganizationSettings(); @@ -462,7 +466,7 @@ export function OrganizationProvidersAndModelsPage({ organizationId, role }: Pro showBackButton={false} /> - + diff --git a/apps/web/src/lib/ai-gateway/auto-model/resolution.test.ts b/apps/web/src/lib/ai-gateway/auto-model/resolution.test.ts index 3ad34cdebf..96062319e8 100644 --- a/apps/web/src/lib/ai-gateway/auto-model/resolution.test.ts +++ b/apps/web/src/lib/ai-gateway/auto-model/resolution.test.ts @@ -162,6 +162,228 @@ describe('resolveAutoModel — kilo-auto/efficient branch', () => { expect(thunk).toHaveBeenCalledTimes(1); }); + + it('applies complete catalog settings for variant xhigh (distinct from max)', async () => { + // Claude catalog: xhigh → effort xhigh + verbosity xhigh; max → effort xhigh + verbosity max + const claudeModel = 'anthropic/claude-sonnet-5'; + const xhighResult = await resolveAutoModel( + { + ...baseParams, + apiKind: 'chat_completions', + efficientDecision: async () => ({ + ...sampleDecision, + model: claudeModel, + variant: 'xhigh', + }), + }, + nullUserPromise, + zeroBalancePromise + ); + const maxResult = await resolveAutoModel( + { + ...baseParams, + apiKind: 'chat_completions', + efficientDecision: async () => ({ + ...sampleDecision, + model: claudeModel, + variant: 'max', + }), + }, + nullUserPromise, + zeroBalancePromise + ); + + expect(xhighResult).toEqual({ + kind: 'ok', + resolved: { + model: claudeModel, + reasoning: { enabled: true, effort: 'xhigh' }, + verbosity: 'xhigh', + }, + }); + expect(maxResult).toEqual({ + kind: 'ok', + resolved: { + model: claudeModel, + reasoning: { enabled: true, effort: 'xhigh' }, + verbosity: 'max', + }, + }); + expect(xhighResult).not.toEqual(maxResult); + }); + + it('applies Claude max variant with both reasoning and verbosity', async () => { + const claudeModel = 'anthropic/claude-haiku-4.5'; + const result = await resolveAutoModel( + { + ...baseParams, + apiKind: 'chat_completions', + efficientDecision: async () => ({ + ...sampleDecision, + model: claudeModel, + variant: 'max', + }), + }, + nullUserPromise, + zeroBalancePromise + ); + + expect(result).toEqual({ + kind: 'ok', + resolved: { + model: claudeModel, + reasoning: { enabled: true, effort: 'xhigh' }, + verbosity: 'max', + }, + }); + }); + + it('falls back to BALANCED_QWEN_MODEL when variant is absent from the model catalog', async () => { + // Claude has no "thinking" key — only none/low/medium/high/xhigh/max + const result = await resolveAutoModel( + { + ...baseParams, + apiKind: 'chat_completions', + efficientDecision: async () => ({ + ...sampleDecision, + model: 'anthropic/claude-sonnet-5', + variant: 'thinking', + }), + }, + nullUserPromise, + zeroBalancePromise + ); + + expect(result).toEqual({ kind: 'ok', resolved: BALANCED_QWEN_MODEL }); + }); + + it('falls back to BALANCED_QWEN_MODEL when the model exposes no variants but decision has a variant', async () => { + const result = await resolveAutoModel( + { + ...baseParams, + apiKind: 'chat_completions', + efficientDecision: async () => ({ + ...sampleDecision, + model: 'some-provider/unknown-model-without-variants', + variant: 'high', + }), + }, + nullUserPromise, + zeroBalancePromise + ); + + expect(result).toEqual({ kind: 'ok', resolved: BALANCED_QWEN_MODEL }); + }); + + it('applies exact thinking and instant variant settings', async () => { + // kimi-k2 uses REASONING_VARIANTS_BINARY: instant + thinking + const kimiModel = 'moonshotai/kimi-k2.5'; + const thinkingResult = await resolveAutoModel( + { + ...baseParams, + apiKind: 'chat_completions', + efficientDecision: async () => ({ + ...sampleDecision, + model: kimiModel, + variant: 'thinking', + }), + }, + nullUserPromise, + zeroBalancePromise + ); + const instantResult = await resolveAutoModel( + { + ...baseParams, + apiKind: 'chat_completions', + efficientDecision: async () => ({ + ...sampleDecision, + model: kimiModel, + variant: 'instant', + }), + }, + nullUserPromise, + zeroBalancePromise + ); + + expect(thinkingResult).toEqual({ + kind: 'ok', + resolved: { + model: kimiModel, + reasoning: { enabled: true, effort: 'high' }, + }, + }); + expect(instantResult).toEqual({ + kind: 'ok', + resolved: { + model: kimiModel, + reasoning: { enabled: false, effort: 'none' }, + }, + }); + }); + + it('preserves legacy effort-only behavior when variant is absent', async () => { + const withEffort = await resolveAutoModel( + { + ...baseParams, + apiKind: 'chat_completions', + efficientDecision: async () => ({ ...sampleDecision, reasoningEffort: 'high' }), + }, + nullUserPromise, + zeroBalancePromise + ); + const withoutEffort = await resolveAutoModel( + { + ...baseParams, + apiKind: 'chat_completions', + efficientDecision: async () => sampleDecision, + }, + nullUserPromise, + zeroBalancePromise + ); + const nullEffort = await resolveAutoModel( + { + ...baseParams, + apiKind: 'chat_completions', + efficientDecision: async () => ({ ...sampleDecision, reasoningEffort: null }), + }, + nullUserPromise, + zeroBalancePromise + ); + + expect(withEffort).toEqual({ + kind: 'ok', + resolved: { + model: 'anthropic/claude-haiku-4', + reasoning: { enabled: true, effort: 'high' }, + }, + }); + expect(withoutEffort).toEqual({ + kind: 'ok', + resolved: { model: 'anthropic/claude-haiku-4' }, + }); + expect(nullEffort).toEqual({ + kind: 'ok', + resolved: { model: 'anthropic/claude-haiku-4' }, + }); + }); + + it('still falls back when the decision model is a virtual auto model even with a variant', async () => { + const result = await resolveAutoModel( + { + ...baseParams, + apiKind: 'chat_completions', + efficientDecision: async () => ({ + ...sampleDecision, + model: KILO_AUTO_EFFICIENT_MODEL.id, + variant: 'high', + }), + }, + nullUserPromise, + zeroBalancePromise + ); + + expect(result).toEqual({ kind: 'ok', resolved: BALANCED_QWEN_MODEL }); + }); }); describe('resolveAutoModel — Organization Auto branch', () => { diff --git a/apps/web/src/lib/ai-gateway/auto-model/resolution.ts b/apps/web/src/lib/ai-gateway/auto-model/resolution.ts index 3ff53046a0..c82a69ccb7 100644 --- a/apps/web/src/lib/ai-gateway/auto-model/resolution.ts +++ b/apps/web/src/lib/ai-gateway/auto-model/resolution.ts @@ -43,6 +43,9 @@ import { isOrganizationAutoTargetModel, validateOrganizationAutoTarget, } from '@/lib/organizations/organization-auto-model'; +import { getModelVariants } from '@/lib/ai-gateway/providers/model-settings'; +import type { OpenCodeVariant } from '@kilocode/db/schema-types'; +import type { OpenRouterReasoningConfig } from '@/lib/ai-gateway/providers/openrouter/types'; type ResolveAutoModelParams = { model: string; @@ -231,6 +234,44 @@ async function resolveOrganizationAutoModel( }; } +/** + * Map an efficient routing decision onto a concrete model + catalog settings. + * + * Prefer canonical `variant` (complete OpenCode variant settings). When the + * variant key is absent from the model's catalog, return null so the caller + * falls back to balanced rather than serving implicit defaults. When `variant` + * is absent, preserve legacy effort-only behavior for rolling deploys. + */ +function resolveEfficientDecisionModel(decision: AutoRoutingDecision): ResolvedAutoModel | null { + // `variant` is only on the benchmark decision branch of the discriminated + // union; coding-plan defaults never carry it. + if ('variant' in decision && decision.variant != null) { + const variants = getModelVariants(decision.model); + const variantSettings: OpenCodeVariant | undefined = variants?.[decision.variant]; + if (!variantSettings) { + return null; + } + // Catalog variants are the source of truth; cast into ResolvedAutoModel's + // OpenRouter-shaped fields (catalog effort may include values like `max` + // beyond ChatCompletionReasoningEffort). + return { + model: decision.model, + ...(variantSettings.reasoning + ? { reasoning: { ...variantSettings.reasoning } as OpenRouterReasoningConfig } + : {}), + ...(variantSettings.verbosity ? { verbosity: variantSettings.verbosity } : {}), + }; + } + + // Legacy effort-only decisions (old workers during rolling deploy). + return { + model: decision.model, + ...('reasoningEffort' in decision && decision.reasoningEffort + ? { reasoning: { enabled: true, effort: decision.reasoningEffort } } + : {}), + }; +} + export async function resolveAutoModel( params: ResolveAutoModelParams, userPromise: Promise, @@ -265,17 +306,13 @@ export async function resolveAutoModel( if (model === KILO_AUTO_EFFICIENT_MODEL.id) { const decision = params.efficientDecision ? await params.efficientDecision() : null; if (decision && !isVirtualAutoModelId(decision.model)) { - // Apply the candidate's pinned reasoning effort so the model runs under - // the same conditions the benchmark measured it at. - return { - kind: 'ok', - resolved: { - model: decision.model, - ...(decision.reasoningEffort - ? { reasoning: { enabled: true, effort: decision.reasoningEffort } } - : {}), - }, - }; + const resolvedFromDecision = resolveEfficientDecisionModel(decision); + if (resolvedFromDecision) { + return { kind: 'ok', resolved: resolvedFromDecision }; + } + // Exact catalog variant missing or removed: never serve the chosen model + // with implicit defaults — same balanced fallback as the no-decision path. + return { kind: 'ok', resolved: BALANCED_QWEN_MODEL }; } // Static fallback when the worker is slow/unavailable: same model as // balanced so an efficient request never degrades below balanced. diff --git a/apps/web/src/lib/ai-gateway/auto-routing-admin-client.test.ts b/apps/web/src/lib/ai-gateway/auto-routing-admin-client.test.ts index adc58001b3..d45f28ac05 100644 --- a/apps/web/src/lib/ai-gateway/auto-routing-admin-client.test.ts +++ b/apps/web/src/lib/ai-gateway/auto-routing-admin-client.test.ts @@ -1,7 +1,9 @@ import { getAutoRoutingClassifierAnalytics, getAutoRoutingClassifierModel, + getAutoRoutingSettings, updateAutoRoutingClassifierModel, + updateAutoRoutingSettings, } from './auto-routing-admin-client'; jest.mock('@/lib/config.server', () => ({ @@ -38,6 +40,21 @@ const classifierAnalyticsResponse = { classifierModelBreakdown: [], }; +const settingsResponse = { + ownerType: 'user', + ownerId: 'user-1', + mode: 'cost_per_accuracy', + configuredMode: null, + defaultMode: 'cost_per_accuracy', + configuredPool: [{ model: 'google/gemini-2.5-flash', variant: null }], + poolStatuses: [ + { + entry: { model: 'google/gemini-2.5-flash', variant: null }, + status: 'ready', + }, + ], +}; + describe('auto routing admin client', () => { beforeEach(() => { mockFetch.mockReset(); @@ -129,4 +146,137 @@ describe('auto routing admin client', () => { } ); }); + + it('gets routing settings using worker bearer auth', async () => { + mockFetch.mockResolvedValue({ + status: 200, + ok: true, + json: () => Promise.resolve(settingsResponse), + }); + + await expect(getAutoRoutingSettings({ ownerType: 'user', ownerId: 'user-1' })).resolves.toEqual( + { + status: 200, + body: settingsResponse, + } + ); + + expect(mockFetch).toHaveBeenCalledWith( + 'https://auto-routing.example.com/admin/routing-settings?ownerType=user&ownerId=user-1', + { + method: 'GET', + headers: { + authorization: 'Bearer test-internal-secret', + }, + } + ); + }); + + it('updates routing settings and forwards optional retryEntries', async () => { + mockFetch.mockResolvedValue({ + status: 200, + ok: true, + json: () => Promise.resolve(settingsResponse), + }); + + await updateAutoRoutingSettings({ + ownerType: 'org', + ownerId: 'org-1', + mode: 'best_accuracy', + pool: [{ model: 'google/gemini-2.5-flash', variant: null }], + retryEntries: [{ model: 'google/gemini-2.5-flash', variant: null }], + }); + + expect(mockFetch).toHaveBeenCalledWith( + 'https://auto-routing.example.com/admin/routing-settings', + { + method: 'PUT', + headers: { + authorization: 'Bearer test-internal-secret', + 'content-type': 'application/json', + }, + body: JSON.stringify({ + ownerType: 'org', + ownerId: 'org-1', + mode: 'best_accuracy', + pool: [{ model: 'google/gemini-2.5-flash', variant: null }], + retryEntries: [{ model: 'google/gemini-2.5-flash', variant: null }], + }), + } + ); + }); + + it('preserves benchmark quota 429 bodies including retryAt', async () => { + const quotaBody = { + error: 'Benchmark profile request limit exceeded', + retryAt: '2026-07-29T12:00:00.000Z', + }; + mockFetch.mockResolvedValue({ + status: 429, + ok: false, + json: () => Promise.resolve(quotaBody), + }); + + await expect( + updateAutoRoutingSettings({ + ownerType: 'user', + ownerId: 'user-1', + mode: null, + pool: [{ model: 'google/gemini-2.5-flash', variant: null }], + }) + ).resolves.toEqual({ + status: 429, + body: quotaBody, + }); + }); + + it('returns 502 for 2xx bodies that fail settings schema validation', async () => { + mockFetch.mockResolvedValue({ + status: 200, + ok: true, + json: () => Promise.resolve({ unexpected: true }), + }); + + await expect(getAutoRoutingSettings({ ownerType: 'user', ownerId: 'user-1' })).resolves.toEqual( + { + status: 502, + body: { error: 'Invalid worker settings response' }, + } + ); + }); + + it('returns 502 for 2xx non-JSON bodies without throwing', async () => { + mockFetch.mockResolvedValue({ + status: 200, + ok: true, + json: () => Promise.reject(new SyntaxError('Unexpected token < in JSON')), + }); + + await expect( + updateAutoRoutingSettings({ + ownerType: 'user', + ownerId: 'user-1', + mode: null, + pool: null, + }) + ).resolves.toEqual({ + status: 502, + body: { error: 'Invalid worker settings response' }, + }); + }); + + it('passes through non-2xx worker error bodies', async () => { + mockFetch.mockResolvedValue({ + status: 404, + ok: false, + json: () => Promise.resolve({ error: 'Settings not found' }), + }); + + await expect(getAutoRoutingSettings({ ownerType: 'user', ownerId: 'user-1' })).resolves.toEqual( + { + status: 404, + body: { error: 'Settings not found' }, + } + ); + }); }); diff --git a/apps/web/src/lib/ai-gateway/auto-routing-admin-client.ts b/apps/web/src/lib/ai-gateway/auto-routing-admin-client.ts index d7a4b1bb69..8737e188fc 100644 --- a/apps/web/src/lib/ai-gateway/auto-routing-admin-client.ts +++ b/apps/web/src/lib/ai-gateway/auto-routing-admin-client.ts @@ -2,18 +2,33 @@ import { AutoRoutingClassifierAnalyticsResponseSchema, AutoRoutingClassifierModelResponseSchema, AutoRoutingModeResponseSchema, + AutoRoutingSettingsResponseSchema, + BenchmarkProfileQuotaErrorSchema, type AutoRoutingMode, type AutoRoutingModeOwnerType, type AutoRoutingAnalyticsPeriod, + type AutoRoutingSettingsResponse, + type EfficientModelPool, + type PoolEntry, } from '@kilocode/auto-routing-contracts'; -import { AUTO_ROUTING_WORKER_URL } from '@/lib/config.server'; -import { createWorkerAdminFetch } from './worker-admin-fetch'; +import { AUTO_ROUTING_WORKER_URL, INTERNAL_API_SECRET } from '@/lib/config.server'; +import { + createWorkerAdminFetch, + ErrorBodySchema, + type ErrorBody, + type WorkerAdminResult, +} from './worker-admin-fetch'; +import type { BenchmarkProfileQuotaError } from '@kilocode/auto-routing-contracts'; const fetchAutoRoutingAdmin = createWorkerAdminFetch({ workerUrl: AUTO_ROUTING_WORKER_URL, unconfiguredError: 'Auto routing worker is not configured', }); +export type AutoRoutingSettingsWorkerResult = WorkerAdminResult< + AutoRoutingSettingsResponse | BenchmarkProfileQuotaError | ErrorBody +>; + export function getAutoRoutingClassifierModel() { return fetchAutoRoutingAdmin( '/admin/classifier-model', @@ -75,3 +90,94 @@ export function updateAutoRoutingMode(owner: { AutoRoutingModeResponseSchema ); } + +/** + * Settings fetch preserves 429 quota bodies (`error` + `retryAt`) that the + * generic worker admin helper would strip to `{ error }` only. + */ +async function fetchAutoRoutingSettingsAdmin( + path: string, + init: Omit & { headers?: Record } +): Promise { + if (!AUTO_ROUTING_WORKER_URL || !INTERNAL_API_SECRET) { + return { + status: 500, + body: { error: 'Auto routing worker is not configured' }, + }; + } + + const response = await fetch(`${AUTO_ROUTING_WORKER_URL}${path}`, { + ...init, + headers: { + authorization: `Bearer ${INTERNAL_API_SECRET}`, + ...init.headers, + }, + }); + + let body: unknown; + try { + body = await response.json(); + } catch { + return { + status: 502, + body: { error: 'Invalid worker settings response' }, + }; + } + + if (!response.ok) { + if (response.status === 429) { + const quota = BenchmarkProfileQuotaErrorSchema.safeParse(body); + if (quota.success) { + return { status: 429, body: quota.data }; + } + } + const parsedError = ErrorBodySchema.safeParse(body); + return { + status: response.status, + body: parsedError.success + ? parsedError.data + : { error: `Request failed: ${response.status}` }, + }; + } + + const parsed = AutoRoutingSettingsResponseSchema.safeParse(body); + if (!parsed.success) { + return { + status: 502, + body: { error: 'Invalid worker settings response' }, + }; + } + + return { + status: response.status, + body: parsed.data, + }; +} + +export function getAutoRoutingSettings(owner: { + ownerType: AutoRoutingModeOwnerType; + ownerId: string; +}): Promise { + const searchParams = new URLSearchParams(owner); + return fetchAutoRoutingSettingsAdmin(`/admin/routing-settings?${searchParams}`, { + method: 'GET', + }); +} + +export function updateAutoRoutingSettings(params: { + ownerType: AutoRoutingModeOwnerType; + ownerId: string; + mode: AutoRoutingMode | null; + pool: EfficientModelPool | null; + retryEntries?: PoolEntry[]; +}): Promise { + const { retryEntries, ...rest } = params; + return fetchAutoRoutingSettingsAdmin('/admin/routing-settings', { + method: 'PUT', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + ...rest, + ...(retryEntries !== undefined ? { retryEntries } : {}), + }), + }); +} diff --git a/apps/web/src/lib/ai-gateway/auto-routing-pool-validation.test.ts b/apps/web/src/lib/ai-gateway/auto-routing-pool-validation.test.ts new file mode 100644 index 0000000000..ff7eb432c5 --- /dev/null +++ b/apps/web/src/lib/ai-gateway/auto-routing-pool-validation.test.ts @@ -0,0 +1,488 @@ +import { beforeEach, describe, expect, test } from '@jest/globals'; +import { + annotatePoolAvailability, + poolValidationMessage, + toApiSettingsResponse, + toLegacyModeApiSettingsResponse, + validatePoolEntries, + type EligibleCatalog, +} from './auto-routing-pool-validation'; + +jest.mock('@/lib/ai-gateway/providers/openrouter', () => ({ + getEnhancedOpenRouterModels: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/experiments/list-available-experiment-models', () => ({ + listAvailableExperimentModels: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/providers/direct-byok', () => ({ + getDirectByokModelsForUser: jest.fn(), + getDirectByokModelsForOrganization: jest.fn(), +})); +jest.mock('@/lib/organizations/organization-models', () => ({ + getAvailableModelsForOrganization: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/models', () => ({ + kiloExclusiveModels: [ + { + public_id: 'kilo/hidden-model', + status: 'hidden', + }, + { + public_id: 'kilo/public-model', + status: 'public', + }, + ], +})); + +const { getEnhancedOpenRouterModels } = jest.requireMock('@/lib/ai-gateway/providers/openrouter'); +const { listAvailableExperimentModels } = jest.requireMock( + '@/lib/ai-gateway/experiments/list-available-experiment-models' +); +const { getDirectByokModelsForUser, getDirectByokModelsForOrganization } = jest.requireMock( + '@/lib/ai-gateway/providers/direct-byok' +); +const { getAvailableModelsForOrganization } = jest.requireMock( + '@/lib/organizations/organization-models' +); + +const mockedGetEnhanced = jest.mocked(getEnhancedOpenRouterModels); +const mockedListExperiments = jest.mocked(listAvailableExperimentModels); +const mockedGetByokUser = jest.mocked(getDirectByokModelsForUser); +const mockedGetByokOrg = jest.mocked(getDirectByokModelsForOrganization); +const mockedGetOrgModels = jest.mocked(getAvailableModelsForOrganization); + +function model(id: string, variants?: Record) { + return { + id, + name: id, + created: 0, + description: id, + architecture: { + input_modalities: ['text'], + output_modalities: ['text'], + tokenizer: 'Other', + }, + top_provider: { is_moderated: false, context_length: 100_000 }, + pricing: { + prompt: '0', + completion: '0', + }, + context_length: 100_000, + ...(variants ? { opencode: { variants } } : {}), + }; +} + +describe('auto-routing-pool-validation', () => { + beforeEach(() => { + jest.resetAllMocks(); + mockedListExperiments.mockResolvedValue([]); + mockedGetByokUser.mockResolvedValue([]); + mockedGetByokOrg.mockResolvedValue([]); + }); + + describe('validatePoolEntries', () => { + test('accepts eligible managed model with required catalog variant', async () => { + mockedGetEnhanced.mockResolvedValue({ + data: [ + model('anthropic/claude-sonnet-4', { + low: { reasoning: { enabled: true } }, + high: { reasoning: { enabled: true } }, + }), + ], + }); + + const result = await validatePoolEntries({ + user: { id: 'user-1' }, + organizationId: null, + entries: [{ model: 'anthropic/claude-sonnet-4', variant: ' high ' }], + }); + + expect(result).toEqual({ + ok: true, + entries: [{ model: 'anthropic/claude-sonnet-4', variant: 'high' }], + }); + }); + + test('accepts model with no variants when variant is null', async () => { + mockedGetEnhanced.mockResolvedValue({ + data: [model('google/gemini-2.5-flash')], + }); + + const result = await validatePoolEntries({ + user: { id: 'user-1' }, + organizationId: null, + entries: [{ model: 'google/gemini-2.5-flash', variant: null }], + }); + + expect(result).toEqual({ + ok: true, + entries: [{ model: 'google/gemini-2.5-flash', variant: null }], + }); + }); + + test('allows same model with different variants', async () => { + mockedGetEnhanced.mockResolvedValue({ + data: [ + model('anthropic/claude-sonnet-4', { + low: {}, + high: {}, + }), + ], + }); + + const result = await validatePoolEntries({ + user: { id: 'user-1' }, + organizationId: null, + entries: [ + { model: 'anthropic/claude-sonnet-4', variant: 'low' }, + { model: 'anthropic/claude-sonnet-4', variant: 'high' }, + ], + }); + + expect(result.ok).toBe(true); + }); + + test('rejects virtual auto model ids', async () => { + mockedGetEnhanced.mockResolvedValue({ + data: [model('kilo-auto/efficient')], + }); + + const result = await validatePoolEntries({ + user: { id: 'user-1' }, + organizationId: null, + entries: [{ model: 'kilo-auto/efficient', variant: null }], + }); + + expect(result).toEqual({ + ok: false, + error: expect.objectContaining({ + reason: 'virtual_model', + message: poolValidationMessage('virtual_model'), + }), + }); + }); + + test('rejects active experiment model ids', async () => { + mockedGetEnhanced.mockResolvedValue({ data: [model('openrouter/base')] }); + mockedListExperiments.mockResolvedValue([model('experiment/active-model')]); + + const result = await validatePoolEntries({ + user: { id: 'user-1' }, + organizationId: null, + entries: [{ model: 'experiment/active-model', variant: null }], + }); + + expect(result).toEqual({ + ok: false, + error: expect.objectContaining({ + reason: 'experiment_model', + message: poolValidationMessage('experiment_model'), + }), + }); + }); + + test('rejects hidden exclusive models', async () => { + mockedGetEnhanced.mockResolvedValue({ + data: [model('kilo/hidden-model')], + }); + + const result = await validatePoolEntries({ + user: { id: 'user-1' }, + organizationId: null, + entries: [{ model: 'kilo/hidden-model', variant: null }], + }); + + expect(result).toEqual({ + ok: false, + error: expect.objectContaining({ + reason: 'hidden_model', + message: poolValidationMessage('hidden_model'), + }), + }); + }); + + test('rejects direct BYOK-only model ids', async () => { + mockedGetEnhanced.mockResolvedValue({ data: [model('openrouter/base')] }); + mockedGetByokUser.mockResolvedValue([model('openai-codex/gpt-5')]); + + const result = await validatePoolEntries({ + user: { id: 'user-1' }, + organizationId: null, + entries: [{ model: 'openai-codex/gpt-5', variant: null }], + }); + + expect(result).toEqual({ + ok: false, + error: expect.objectContaining({ + reason: 'byok_only_model', + message: poolValidationMessage('byok_only_model'), + }), + }); + }); + + test('rejects organization-denied models with a specific reason', async () => { + mockedGetEnhanced.mockResolvedValue({ + data: [model('openai/gpt-4o'), model('anthropic/claude-sonnet-4')], + }); + mockedGetOrgModels.mockResolvedValue({ + data: [model('anthropic/claude-sonnet-4')], + }); + + const result = await validatePoolEntries({ + user: { id: 'user-1' }, + organizationId: 'org-1', + entries: [{ model: 'openai/gpt-4o', variant: null }], + }); + + expect(mockedGetByokOrg).toHaveBeenCalledWith('org-1'); + expect(result).toEqual({ + ok: false, + error: expect.objectContaining({ + reason: 'organization_denied_model', + message: poolValidationMessage('organization_denied_model'), + }), + }); + }); + + test('rejects unknown models', async () => { + mockedGetEnhanced.mockResolvedValue({ data: [model('openrouter/base')] }); + + const result = await validatePoolEntries({ + user: { id: 'user-1' }, + organizationId: null, + entries: [{ model: 'totally/unknown', variant: null }], + }); + + expect(result).toEqual({ + ok: false, + error: expect.objectContaining({ + reason: 'unknown_model', + message: poolValidationMessage('unknown_model'), + }), + }); + }); + + test('rejects missing variant when model exposes variants', async () => { + mockedGetEnhanced.mockResolvedValue({ + data: [model('anthropic/claude-sonnet-4', { high: {} })], + }); + + const result = await validatePoolEntries({ + user: { id: 'user-1' }, + organizationId: null, + entries: [{ model: 'anthropic/claude-sonnet-4', variant: null }], + }); + + expect(result).toEqual({ + ok: false, + error: expect.objectContaining({ + reason: 'missing_variant', + message: poolValidationMessage('missing_variant'), + }), + }); + }); + + test('rejects unknown variant keys', async () => { + mockedGetEnhanced.mockResolvedValue({ + data: [model('anthropic/claude-sonnet-4', { high: {} })], + }); + + const result = await validatePoolEntries({ + user: { id: 'user-1' }, + organizationId: null, + entries: [{ model: 'anthropic/claude-sonnet-4', variant: 'max' }], + }); + + expect(result).toEqual({ + ok: false, + error: expect.objectContaining({ + reason: 'unknown_variant', + message: poolValidationMessage('unknown_variant'), + }), + }); + }); + + test('rejects unexpected variant when model exposes none', async () => { + mockedGetEnhanced.mockResolvedValue({ + data: [model('google/gemini-2.5-flash')], + }); + + const result = await validatePoolEntries({ + user: { id: 'user-1' }, + organizationId: null, + entries: [{ model: 'google/gemini-2.5-flash', variant: 'high' }], + }); + + expect(result).toEqual({ + ok: false, + error: expect.objectContaining({ + reason: 'unexpected_variant', + message: poolValidationMessage('unexpected_variant'), + }), + }); + }); + + test('rejects duplicate exact pairs', async () => { + mockedGetEnhanced.mockResolvedValue({ + data: [model('anthropic/claude-sonnet-4', { high: {} })], + }); + + const result = await validatePoolEntries({ + user: { id: 'user-1' }, + organizationId: null, + entries: [ + { model: 'anthropic/claude-sonnet-4', variant: 'high' }, + { model: 'anthropic/claude-sonnet-4', variant: 'high' }, + ], + }); + + expect(result).toEqual({ + ok: false, + error: expect.objectContaining({ + reason: 'duplicate_pair', + message: poolValidationMessage('duplicate_pair'), + }), + }); + }); + + test('rejects more than 10 entries', async () => { + mockedGetEnhanced.mockResolvedValue({ + data: [model('google/gemini-2.5-flash')], + }); + + // Length is checked before catalog/duplicate logic. + const entries = Array.from({ length: 11 }, () => ({ + model: 'google/gemini-2.5-flash', + variant: null as null, + })); + + const result = await validatePoolEntries({ + user: { id: 'user-1' }, + organizationId: null, + entries, + }); + + expect(result).toEqual({ + ok: false, + error: expect.objectContaining({ + reason: 'too_many_entries', + message: poolValidationMessage('too_many_entries'), + }), + }); + }); + + test('rejects empty pool arrays', async () => { + mockedGetEnhanced.mockResolvedValue({ data: [] }); + + const result = await validatePoolEntries({ + user: { id: 'user-1' }, + organizationId: null, + entries: [], + }); + + expect(result).toEqual({ + ok: false, + error: expect.objectContaining({ + reason: 'empty_pool', + message: poolValidationMessage('empty_pool'), + }), + }); + }); + }); + + describe('annotatePoolAvailability', () => { + const catalog: EligibleCatalog = { + byId: new Map([ + [ + 'anthropic/claude-sonnet-4', + { + id: 'anthropic/claude-sonnet-4', + variantKeys: new Set(['high', 'low']), + }, + ], + [ + 'google/gemini-2.5-flash', + { + id: 'google/gemini-2.5-flash', + variantKeys: null, + }, + ], + ]), + experimentIds: new Set(), + byokOnlyIds: new Set(), + managedIds: new Set(['anthropic/claude-sonnet-4', 'google/gemini-2.5-flash']), + ownerCatalogIds: new Set(['anthropic/claude-sonnet-4', 'google/gemini-2.5-flash']), + organizationId: null, + }; + + test('marks null configured pool as null', () => { + expect(annotatePoolAvailability({ entries: null, catalog })).toBeNull(); + }); + + test('marks still-eligible entries available and departed ones unavailable', () => { + const annotated = annotatePoolAvailability({ + catalog, + entries: [ + { model: 'anthropic/claude-sonnet-4', variant: 'high' }, + { model: 'anthropic/claude-sonnet-4', variant: 'max' }, + { model: 'removed/model', variant: null }, + { model: 'google/gemini-2.5-flash', variant: null }, + ], + }); + + expect(annotated).toEqual([ + { model: 'anthropic/claude-sonnet-4', variant: 'high', unavailable: false }, + { model: 'anthropic/claude-sonnet-4', variant: 'max', unavailable: true }, + { model: 'removed/model', variant: null, unavailable: true }, + { model: 'google/gemini-2.5-flash', variant: null, unavailable: false }, + ]); + }); + }); + + describe('toApiSettingsResponse / toLegacyModeApiSettingsResponse', () => { + test('toApiSettingsResponse includes poolSupported: true', () => { + const response = toApiSettingsResponse( + { + ownerType: 'user', + ownerId: 'user-1', + mode: 'cost_per_accuracy', + configuredMode: null, + defaultMode: 'cost_per_accuracy', + configuredPool: null, + poolStatuses: [], + }, + null + ); + expect(response.poolSupported).toBe(true); + expect(response).toEqual( + expect.objectContaining({ + ownerType: 'user', + ownerId: 'user-1', + configuredPool: null, + poolSupported: true, + }) + ); + }); + + test('toLegacyModeApiSettingsResponse synthesizes poolSupported: false', () => { + expect( + toLegacyModeApiSettingsResponse({ + ownerType: 'user', + ownerId: 'user-1', + mode: 'best_accuracy', + configuredMode: 'best_accuracy', + defaultMode: 'cost_per_accuracy', + }) + ).toEqual({ + ownerType: 'user', + ownerId: 'user-1', + mode: 'best_accuracy', + configuredMode: 'best_accuracy', + defaultMode: 'cost_per_accuracy', + configuredPool: null, + poolStatuses: [], + poolSupported: false, + }); + }); + }); +}); diff --git a/apps/web/src/lib/ai-gateway/auto-routing-pool-validation.ts b/apps/web/src/lib/ai-gateway/auto-routing-pool-validation.ts new file mode 100644 index 0000000000..9e860ad80f --- /dev/null +++ b/apps/web/src/lib/ai-gateway/auto-routing-pool-validation.ts @@ -0,0 +1,414 @@ +import 'server-only'; + +import { + isVirtualAutoModelId, + MAX_POOL_ENTRIES, + poolEntryKey, + type AutoRoutingSettingsResponse, + type PoolEntry, +} from '@kilocode/auto-routing-contracts'; +import { CUSTOM_LLM_PREFIX } from '@/lib/ai-gateway/model-utils'; +import { kiloExclusiveModels } from '@/lib/ai-gateway/models'; +import { listAvailableExperimentModels } from '@/lib/ai-gateway/experiments/list-available-experiment-models'; +import { + getDirectByokModelsForOrganization, + getDirectByokModelsForUser, +} from '@/lib/ai-gateway/providers/direct-byok'; +import { getEnhancedOpenRouterModels } from '@/lib/ai-gateway/providers/openrouter'; +import { getAvailableModelsForOrganization } from '@/lib/organizations/organization-models'; +import type { OpenRouterModel } from '@/lib/organizations/organization-types'; + +export type PoolValidationReason = + | 'unknown_model' + | 'virtual_model' + | 'experiment_model' + | 'hidden_model' + | 'byok_only_model' + | 'organization_denied_model' + | 'missing_variant' + | 'unknown_variant' + | 'unexpected_variant' + | 'duplicate_pair' + | 'too_many_entries' + | 'empty_pool'; + +export type PoolValidationError = { + reason: PoolValidationReason; + message: string; + index?: number; + entry?: PoolEntry; +}; + +export type PoolEntryWithAvailability = PoolEntry & { + unavailable: boolean; +}; + +/** Web API response: worker settings plus per-entry availability for the UI. */ +export type AutoRoutingSettingsApiResponse = Omit & { + configuredPool: PoolEntryWithAvailability[] | null; + /** + * False when the worker only supports legacy mode endpoints (deploy-order + * window or worker rollback). UI must hide pool editing and never send pool mutations. + */ + poolSupported: boolean; +}; + +export type EligibleModelInfo = { + /** Canonical catalog model id. */ + id: string; + /** + * Canonical variant keys when the model exposes `opencode.variants`. + * `null` means the model exposes no variants (variant must be null). + */ + variantKeys: ReadonlySet | null; +}; + +export type EligibleCatalog = { + byId: ReadonlyMap; + experimentIds: ReadonlySet; + byokOnlyIds: ReadonlySet; + /** Managed catalog ids before org filtering (for org-denied classification). */ + managedIds: ReadonlySet; + /** Ids present in the owner-scoped catalog after org filtering. */ + ownerCatalogIds: ReadonlySet; + organizationId: string | null; +}; + +const REJECTION_MESSAGES: Record = { + unknown_model: 'Unknown model. Choose a managed model from your catalog.', + virtual_model: 'Virtual auto-routing models cannot be added to an Efficient pool.', + experiment_model: 'Model experiment IDs cannot be added to an Efficient pool.', + hidden_model: 'This model is not visible in your model picker.', + byok_only_model: 'Direct BYOK-only models cannot be added to an Efficient pool.', + organization_denied_model: 'This model is not allowed for the organization.', + missing_variant: 'A catalog variant is required for this model.', + unknown_variant: 'Unknown variant for this model.', + unexpected_variant: 'This model does not expose variants; variant must be null.', + duplicate_pair: 'Duplicate model and variant pair in the pool.', + too_many_entries: `An Efficient pool can include at most ${MAX_POOL_ENTRIES} models.`, + empty_pool: 'An Efficient pool must include at least one model.', +}; + +export function poolValidationMessage(reason: PoolValidationReason): string { + return REJECTION_MESSAGES[reason]; +} + +function variantKeysForModel(model: OpenRouterModel): ReadonlySet | null { + const variants = model.opencode?.variants; + if (!variants) return null; + const keys = Object.keys(variants).filter(key => key.trim().length > 0); + return keys.length > 0 ? new Set(keys) : null; +} + +function isHiddenExclusiveModel(modelId: string): boolean { + const exclusive = kiloExclusiveModels.find(model => model.public_id === modelId); + return exclusive !== undefined && exclusive.status !== 'public'; +} + +function isCustomLlmId(modelId: string): boolean { + return modelId.startsWith(CUSTOM_LLM_PREFIX); +} + +/** + * Build the eligible managed catalog for pool membership. + * Client-side filtering is never trusted; this is the authorization source. + */ +export async function buildEligibleCatalog(params: { + userId: string; + organizationId: string | null; +}): Promise { + const { userId, organizationId } = params; + + const [enhanced, experimentModels, byokModels] = await Promise.all([ + getEnhancedOpenRouterModels(), + listAvailableExperimentModels(), + organizationId + ? getDirectByokModelsForOrganization(organizationId) + : getDirectByokModelsForUser(userId), + ]); + + const managedModels = Array.isArray(enhanced.data) ? enhanced.data : []; + const managedIds = new Set(managedModels.map(model => model.id)); + const experimentIds = new Set(experimentModels.map(model => model.id)); + const byokOnlyIds = new Set(byokModels.map(model => model.id)); + + let ownerModels: OpenRouterModel[]; + if (organizationId) { + const orgCatalog = await getAvailableModelsForOrganization(organizationId); + ownerModels = orgCatalog?.data ?? []; + } else { + ownerModels = managedModels; + } + + const ownerCatalogIds = new Set(ownerModels.map(model => model.id)); + const byId = new Map(); + + for (const model of ownerModels) { + const id = model.id; + if (isVirtualAutoModelId(id)) continue; + if (experimentIds.has(id)) continue; + if (byokOnlyIds.has(id)) continue; + if (isCustomLlmId(id)) continue; + if (isHiddenExclusiveModel(id)) continue; + + byId.set(id, { + id, + variantKeys: variantKeysForModel(model), + }); + } + + return { + byId, + experimentIds, + byokOnlyIds, + managedIds, + ownerCatalogIds, + organizationId, + }; +} + +function classifyIneligibleModel( + modelId: string, + catalog: EligibleCatalog +): Exclude< + PoolValidationReason, + | 'missing_variant' + | 'unknown_variant' + | 'unexpected_variant' + | 'duplicate_pair' + | 'too_many_entries' + | 'empty_pool' +> { + if (isVirtualAutoModelId(modelId)) return 'virtual_model'; + if (catalog.experimentIds.has(modelId)) return 'experiment_model'; + if (catalog.byokOnlyIds.has(modelId)) return 'byok_only_model'; + if (isHiddenExclusiveModel(modelId)) return 'hidden_model'; + if ( + catalog.organizationId && + catalog.managedIds.has(modelId) && + !catalog.ownerCatalogIds.has(modelId) + ) { + return 'organization_denied_model'; + } + // Present in owner catalog but subtracted (custom LLM, etc.) → treat as unknown. + return 'unknown_model'; +} + +function validateSingleEntry( + entry: PoolEntry, + index: number, + catalog: EligibleCatalog +): { ok: true; entry: PoolEntry } | { ok: false; error: PoolValidationError } { + const modelId = entry.model.trim(); + if (!modelId) { + return { + ok: false, + error: { + reason: 'unknown_model', + message: poolValidationMessage('unknown_model'), + index, + entry, + }, + }; + } + + const eligible = catalog.byId.get(modelId); + if (!eligible) { + // Try case-sensitive exact match only (catalog ids are canonical). + const reason = classifyIneligibleModel(modelId, catalog); + return { + ok: false, + error: { + reason, + message: poolValidationMessage(reason), + index, + entry, + }, + }; + } + + const rawVariant = entry.variant; + const variant = + rawVariant === null || rawVariant === undefined ? null : rawVariant.trim() || null; + + if (eligible.variantKeys === null) { + if (variant !== null) { + return { + ok: false, + error: { + reason: 'unexpected_variant', + message: poolValidationMessage('unexpected_variant'), + index, + entry, + }, + }; + } + return { ok: true, entry: { model: eligible.id, variant: null } }; + } + + if (variant === null) { + return { + ok: false, + error: { + reason: 'missing_variant', + message: poolValidationMessage('missing_variant'), + index, + entry, + }, + }; + } + + // Match exact catalog key (trim already applied); keys are as cataloged. + if (!eligible.variantKeys.has(variant)) { + // Also accept if a catalog key equals after the same trim (keys are already trimmed filters). + const canonical = [...eligible.variantKeys].find(key => key === variant); + if (!canonical) { + return { + ok: false, + error: { + reason: 'unknown_variant', + message: poolValidationMessage('unknown_variant'), + index, + entry, + }, + }; + } + return { ok: true, entry: { model: eligible.id, variant: canonical } }; + } + + return { ok: true, entry: { model: eligible.id, variant } }; +} + +/** + * Revalidate every submitted Pool entry against the owner's effective managed catalog. + * Returns canonicalized entries on success. + */ +export async function validatePoolEntries(params: { + user: { id: string }; + organizationId: string | null; + entries: PoolEntry[]; +}): Promise<{ ok: true; entries: PoolEntry[] } | { ok: false; error: PoolValidationError }> { + const { entries } = params; + + if (entries.length === 0) { + return { + ok: false, + error: { + reason: 'empty_pool', + message: poolValidationMessage('empty_pool'), + }, + }; + } + + if (entries.length > MAX_POOL_ENTRIES) { + return { + ok: false, + error: { + reason: 'too_many_entries', + message: poolValidationMessage('too_many_entries'), + }, + }; + } + + const catalog = await buildEligibleCatalog({ + userId: params.user.id, + organizationId: params.organizationId, + }); + + const canonical: PoolEntry[] = []; + const seen = new Set(); + + for (let index = 0; index < entries.length; index++) { + const entry = entries[index]; + if (!entry) continue; + const result = validateSingleEntry(entry, index, catalog); + if (!result.ok) return result; + + const key = poolEntryKey(result.entry); + if (seen.has(key)) { + return { + ok: false, + error: { + reason: 'duplicate_pair', + message: poolValidationMessage('duplicate_pair'), + index, + entry: result.entry, + }, + }; + } + seen.add(key); + canonical.push(result.entry); + } + + return { ok: true, entries: canonical }; +} + +function isEntryAvailable(entry: PoolEntry, catalog: EligibleCatalog): boolean { + const eligible = catalog.byId.get(entry.model); + if (!eligible) return false; + + if (eligible.variantKeys === null) { + return entry.variant === null; + } + if (entry.variant === null) return false; + return eligible.variantKeys.has(entry.variant); +} + +/** + * Classify saved pool entries against today's catalog so GET can mark + * departed entries unavailable without rejecting the read. + */ +export function annotatePoolAvailability(params: { + entries: PoolEntry[] | null; + catalog: EligibleCatalog; +}): PoolEntryWithAvailability[] | null { + const { entries, catalog } = params; + if (entries === null) return null; + return entries.map(entry => ({ + ...entry, + unavailable: !isEntryAvailable(entry, catalog), + })); +} + +export async function annotateConfiguredPool(params: { + userId: string; + organizationId: string | null; + configuredPool: PoolEntry[] | null; +}): Promise { + if (params.configuredPool === null) return null; + const catalog = await buildEligibleCatalog({ + userId: params.userId, + organizationId: params.organizationId, + }); + return annotatePoolAvailability({ entries: params.configuredPool, catalog }); +} + +export function toApiSettingsResponse( + workerBody: AutoRoutingSettingsResponse, + configuredPool: PoolEntryWithAvailability[] | null +): AutoRoutingSettingsApiResponse { + return { + ...workerBody, + configuredPool, + poolSupported: true, + }; +} + +/** + * Synthesize a settings API response from the legacy mode-only worker body + * when `/admin/routing-settings` is not yet deployed (or was rolled back). + */ +export function toLegacyModeApiSettingsResponse(modeBody: { + ownerType: AutoRoutingSettingsResponse['ownerType']; + ownerId: string; + mode: AutoRoutingSettingsResponse['mode']; + configuredMode: AutoRoutingSettingsResponse['configuredMode']; + defaultMode: AutoRoutingSettingsResponse['defaultMode']; +}): AutoRoutingSettingsApiResponse { + return { + ...modeBody, + configuredPool: null, + poolStatuses: [], + poolSupported: false, + }; +} diff --git a/docs/adr/0003-efficient-model-pools.md b/docs/adr/0003-efficient-model-pools.md new file mode 100644 index 0000000000..656f2ef381 --- /dev/null +++ b/docs/adr/0003-efficient-model-pools.md @@ -0,0 +1,129 @@ +# ADR 0003: Efficient Model Pools + +## Status + +Accepted + +## Context + +[ADR 0002](./0002-auto-routing-efficient.md) introduced `kilo-auto/efficient`, a +hidden virtual model backed by a benchmark-driven decision engine that publishes +a single platform routing table. Every request is routed through the same +published table — there is no per-user or per-organization candidate set. + +Users and organizations want to constrain `kilo-auto/efficient` to a specific +subset of models they trust or have budget approval for, without forking the +routing infrastructure or introducing a new virtual model ID. They also want +benchmark measurements for those custom subsets without maintaining their own +benchmark pipeline. + +## Decision + +Introduce **Efficient model pools** as an owner-configurable setting on the +existing `kilo-auto/efficient` virtual model. A pool is a set of 1–10 **pool +entries**, each an exact `(managed model, canonical catalog variant)` pair. Pool +entries select candidates for custom routing; absent or cleared pools restore +platform-default behavior (the published platform efficient routing table). +Balanced is only the gateway fallback when `/decide` returns null (no table, +no ready/compatible selected pair, or an omitted sparse route). + +The architecture adds three concepts without a new routing mode or virtual model +id: + +1. **Owner pool settings** live in a Durable Object per owner (personal user or + organization). Mode (efficient/balanced) and pool are committed in one DO + storage transaction so a mixed save — one field set, the other cleared — can + never half-apply. Inheritance is independent for mode and pool (organization + → personal → platform default). Clearing a configured pool restores + inheritance. + +2. **Benchmark profiles** are global measurement records keyed by exact pool + entry `(model, variant)` plus engine identity. The benchmark worker + (`services/auto-routing-benchmark`) remains the sole writer of all profiles. + Profiles are generated on demand when an owner save admits missing or stale + work, rather than requiring a pre-provisioned catalogue. + +3. **Sparse assembled routing tables** are built at decision time from ready + and current selected profiles. When no selected pair is ready or compatible, + `/decide` returns null and the gateway falls back to balanced. Platform and + scheduled decider runs keep publishing the default routing table + (`ROUTING_TABLE_KV_KEY`); profile runs update profile state only and never + replace the platform artifact. + +## Invariants (what not to change without revisiting this ADR) + +1. **Pool entries are exact pairs.** An entry is `(managed model, canonical + catalog variant)`. Null variant is allowed only for models that expose no + variants. Profile identity and carried-result gating both use this exact + pair. + +2. **Global profiles, single writer.** Profiles are stored globally (not + per-owner) in the benchmark worker's D1, keyed by exact pool entry plus + engine identity. The benchmark worker is the sole writer; the decision engine + reads profiles through the existing cache chain (isolate → KV → service + binding to D1) and never writes back. + +3. **On-demand generation with admission limits.** Profiles are generated when + an owner save admits missing or stale work. Admission is limited to 10 + previously unbenchmarked or explicitly retried profiles per owner per rolling + 24 hours, with global deduplication of concurrent exact-pair work. Ready or + already-pending profiles do not consume request quota. + +4. **Single decider slot, shared platform and profile.** Platform and profile + decider work share the single active decider slot enforced by the partial + unique index (see ADR 0002 invariant 6). Profile runs fail closed when + `markProfilesRunningForRun` rejects after insert: the run is marked failed, + the slot is freed for drain, and no queue work is enqueued. Stale runs sweep + profile claims along with platform rows. + +5. **Platform vs profile publication.** Platform/admin decider runs publish the + default routing table. Profile runs transition profile state + (`running` → `ready` / `failed`) and never touch the platform artifact. A + separate sweep+drain handler ensures a failed final queue message or platform + run cannot strand pending profile work. + +6. **Fail-closed custom decide.** Sparse custom tables omit routes with no + graded candidates rather than inventing empty lists. Omitted routes yield no + decision and the gateway's balanced fallback. A configured pool with no + ready/compatible selected pair likewise returns no decision. + +7. **Deploy-order contract.** Production promotes the Vercel web deployment + *before* deploying workers, so every merge runs new web against old workers + for the worker build window. During that window, + `GET /admin/routing-settings` 404s on old workers; the web BFF synthesizes + the settings response from the legacy `/admin/routing-mode` route with + `poolSupported: false` (no pool annotation), and the Auto routing card hides + pool controls and saves mode-only through the legacy web mode route — + mode-only at every worker version, never touching pool keys. A settings PUT + that 404s (mid-session worker rollback from a supported UI) answers a + retryable 503 `pool_temporarily_unavailable` for every body shape; there is + deliberately no legacy PUT fallback, because `pool: null` from a supported UI + is clear intent an old worker cannot honor (it would silently preserve the + pool), and a stale unsupported-UI body must never reach a new worker's + settings PUT (it would silently clear the pool). + +8. **Inheritance.** Mode and pool inherit independently (organization → + personal → platform default). A configured pool does not force the owner into + efficient mode; a balanced-mode owner with a pool is inert until mode is + switched. Clearing a configured pool restores inheritance. + +## Consequences + +Profiles add a new dimension to the benchmark worker's D1 schema (profile +tables, pending-profile queue tables) and a new scheduled sweep handler to +prevent stranded work. The single decider slot shared between platform and +profile work means profile runs compete with platform runs for the slot, but +the admission limits and drain mechanism keep the slot from starving either +purpose. + +The deploy-order contract constrains when pool-related routes can land in +workers relative to the web deployment. Pool settings are additive to the +existing worker schema; rollback to pre-pool workers runs prior code against +the migrated schema (migration 0005 rebuilds `case_results` / `model_summaries` +/ `run_models` with generated rebuild SQL preserving legacy rows; 0006/0007 are +additive profile and pending-profile tables). + +## References + +- [ADR 0002: Benchmark-Driven Auto Routing](./0002-auto-routing-efficient.md) — + the base efficient routing decision that pools extend. diff --git a/packages/auto-routing-contracts/src/benchmark.ts b/packages/auto-routing-contracts/src/benchmark.ts index d7c7bdea24..a3fcac878c 100644 --- a/packages/auto-routing-contracts/src/benchmark.ts +++ b/packages/auto-routing-contracts/src/benchmark.ts @@ -1,21 +1,51 @@ import * as z from 'zod'; -import { RoutingTableSchema } from './routing-table'; +import { + CustomRoutingTableSchema, + EfficientModelPoolSchema, + PoolEntrySchema, + poolEntryKey, + RoutingTableSchema, +} from './routing-table'; import { ReasoningEffortSchema } from './reasoning'; import { TaxonomyRouteKeySchema } from './taxonomy'; export { ReasoningEffortSchema } from './reasoning'; export type { ReasoningEffort } from './reasoning'; +// Matches AutoRoutingModeOwnerTypeSchema in index.ts. Declared here (not +// imported) to avoid a circular package-root dependency. +const BenchmarkProfileOwnerTypeSchema = z.enum(['user', 'org']); + export const BenchmarkKindSchema = z.enum(['classifier', 'decider']); export type BenchmarkKind = z.infer; -export const BenchmarkDeciderModelSchema = z.object({ - id: z.string().trim().min(1), - // Passed to the kilo CLI as --variant during the benchmark and carried into - // the routing table so serving uses the same effort the model was graded - // with. Null for models without (or not using) configurable reasoning. - reasoningEffort: ReasoningEffortSchema.nullable().default(null), -}); +/** + * Decider model identity for a benchmark run. Platform/admin config still + * selects one legacy `reasoningEffort` per model. Profile runs (and exact + * Pool-entry identity) use canonical `variant`. Both non-null is malformed — + * same both-set rule as RankedCandidate / decisions. + */ +export const BenchmarkDeciderModelSchema = z + .object({ + id: z.string().trim().min(1), + // Canonical catalog variant key. Optional so platform admin config (effort + // only) still parses. Prefer this over reasoningEffort for new writers. + variant: z.string().trim().min(1).nullable().optional(), + // Passed to the kilo CLI as --variant during the benchmark and carried into + // the platform routing table so serving uses the same effort the model was + // graded with. Null for models without (or not using) configurable + // reasoning. Legacy; prefer `variant` when present. + reasoningEffort: ReasoningEffortSchema.nullable().default(null), + }) + .superRefine((model, ctx) => { + if (model.variant != null && model.reasoningEffort != null) { + ctx.addIssue({ + code: 'custom', + path: ['variant'], + message: 'Decider model must not set both variant and reasoningEffort; emit variant only', + }); + } + }); export type BenchmarkDeciderModel = z.infer; export const AUTO_DECIDER_DEFAULT_MIN_COST_USD = 15; @@ -155,6 +185,9 @@ export type BenchmarkRunStatus = z.infer; export const BenchmarkModelSummarySchema = z.object({ model: z.string(), + // Exact-pair identity for summary rows. Optional so already-persisted run + // payloads and old admin-UI responses still parse. Null = default/no variant. + variant: z.string().trim().min(1).nullable().optional(), // '*' for classifier runs, otherwise "/". routeKey: z.union([TaxonomyRouteKeySchema, z.literal('*')]), accuracy: z.number(), @@ -220,3 +253,80 @@ export const ClassifierWinnerResponseSchema = z.object({ winner: ClassifierWinnerSchema.nullable(), }); export type ClassifierWinnerResponse = z.infer; + +// --- Benchmark profile registry (global per Pool entry + engine) --- + +/** + * Wire status for a Benchmark profile. Presentation maps pending/running to + * "Benchmarking"; "Unavailable" is a web-derived state and never a wire status. + */ +export const BenchmarkProfileStatusSchema = z.enum(['pending', 'running', 'ready', 'failed']); +export type BenchmarkProfileStatus = z.infer; + +/** Bounded failure text stored on failed profile rows. */ +export const BENCHMARK_PROFILE_FAILURE_REASON_MAX_LENGTH = 500; + +export const BenchmarkProfileEntryStatusSchema = z.object({ + entry: PoolEntrySchema, + status: BenchmarkProfileStatusSchema, + failureReason: z.string().max(BENCHMARK_PROFILE_FAILURE_REASON_MAX_LENGTH).nullable().optional(), +}); +export type BenchmarkProfileEntryStatus = z.infer; + +export const RegisterBenchmarkProfilesRequestSchema = z + .object({ + ownerType: BenchmarkProfileOwnerTypeSchema, + ownerId: z.string().trim().min(1), + entries: EfficientModelPoolSchema, + // Subset of `entries` the owner explicitly retries after a terminal failure. + // Absent/empty means no explicit retries (failed current profiles are reported + // without re-admission). + retryEntries: z.array(PoolEntrySchema).optional(), + }) + .superRefine((request, ctx) => { + if (!request.retryEntries || request.retryEntries.length === 0) return; + const entryKeys = new Set(request.entries.map(poolEntryKey)); + request.retryEntries.forEach((entry, index) => { + const key = poolEntryKey(entry); + if (!entryKeys.has(key)) { + ctx.addIssue({ + code: 'custom', + path: ['retryEntries', index], + message: `retryEntry must appear in entries: ${key}`, + }); + } + }); + }); +export type RegisterBenchmarkProfilesRequest = z.infer< + typeof RegisterBenchmarkProfilesRequestSchema +>; + +export const BenchmarkProfileStatusesRequestSchema = z.object({ + entries: EfficientModelPoolSchema, +}); +export type BenchmarkProfileStatusesRequest = z.infer; + +export const BenchmarkProfileStatusesResponseSchema = z.object({ + statuses: z.array(BenchmarkProfileEntryStatusSchema), +}); +export type BenchmarkProfileStatusesResponse = z.infer< + typeof BenchmarkProfileStatusesResponseSchema +>; + +/** 429 body when an owner exceeds the rolling 24h profile admission limit. */ +export const BenchmarkProfileQuotaErrorSchema = z.object({ + error: z.string(), + // ISO timestamp when the owner may request new benchmarks again. + retryAt: z.string(), +}); +export type BenchmarkProfileQuotaError = z.infer; + +export const CustomRoutingTableRequestSchema = z.object({ + entries: EfficientModelPoolSchema, +}); +export type CustomRoutingTableRequest = z.infer; + +export const CustomRoutingTableResponseSchema = z.object({ + table: CustomRoutingTableSchema.nullable(), +}); +export type CustomRoutingTableResponse = z.infer; diff --git a/packages/auto-routing-contracts/src/contracts.test.ts b/packages/auto-routing-contracts/src/contracts.test.ts index 6dbae647d8..7d66845795 100644 --- a/packages/auto-routing-contracts/src/contracts.test.ts +++ b/packages/auto-routing-contracts/src/contracts.test.ts @@ -3,8 +3,11 @@ import { AutoRoutingClassifierAnalyticsResponseSchema, AutoRoutingClassifierModelResponseSchema, AutoRoutingDecisionResponseSchema, + AutoRoutingDecisionSchema, + AutoRoutingSettingsResponseSchema, MirrorPayloadSchema, RoutingConstraintsSchema, + UpdateAutoRoutingSettingsRequestSchema, UpdateClassifierModelRequestSchema, detectRequiredInputModalities, estimateRoutingTokens, @@ -12,8 +15,12 @@ import { import type { RoutingConstraints } from './index'; import { BenchmarkConfigSchema, + BenchmarkDeciderModelSchema, + BenchmarkModelSummarySchema, + BenchmarkProfileStatusesRequestSchema, DEFAULT_BENCHMARK_ORG_ID, DEFAULT_BENCHMARK_USER_ID, + RegisterBenchmarkProfilesRequestSchema, resolveBenchmarkIdentity, } from './benchmark'; @@ -484,3 +491,224 @@ describe('package root re-exports', () => { expect(RoutingConstraintsSchema.safeParse({}).success).toBe(true); }); }); + +describe('BenchmarkModelSummarySchema variant field', () => { + const baseSummary = { + model: 'provider/model', + routeKey: '*' as const, + accuracy: 0.9, + avgCostUsd: 0.001, + avgLatencyMs: 100, + p50LatencyMs: 90, + p95LatencyMs: 120, + cases: 10, + errors: 0, + }; + + it('parses legacy summaries without variant', () => { + const parsed = BenchmarkModelSummarySchema.parse(baseSummary); + expect(parsed.variant).toBeUndefined(); + }); + + it('parses optional nullable variant', () => { + expect(BenchmarkModelSummarySchema.parse({ ...baseSummary, variant: 'xhigh' }).variant).toBe( + 'xhigh' + ); + expect(BenchmarkModelSummarySchema.parse({ ...baseSummary, variant: null }).variant).toBeNull(); + }); +}); + +describe('BenchmarkDeciderModelSchema variant compatibility', () => { + it('parses legacy effort-only models', () => { + const parsed = BenchmarkDeciderModelSchema.parse({ id: 'm', reasoningEffort: 'high' }); + expect(parsed).toMatchObject({ id: 'm', reasoningEffort: 'high' }); + expect(parsed.variant).toBeUndefined(); + }); + + it('parses variant-only models', () => { + const parsed = BenchmarkDeciderModelSchema.parse({ id: 'm', variant: 'xhigh' }); + expect(parsed).toMatchObject({ id: 'm', variant: 'xhigh', reasoningEffort: null }); + }); + + it('rejects both non-null variant and reasoningEffort', () => { + expect( + BenchmarkDeciderModelSchema.safeParse({ + id: 'm', + variant: 'xhigh', + reasoningEffort: 'high', + }).success + ).toBe(false); + }); + + it('allows both null/absent (default effort null)', () => { + expect(BenchmarkDeciderModelSchema.parse({ id: 'm' })).toMatchObject({ + id: 'm', + reasoningEffort: null, + }); + }); +}); + +describe('AutoRoutingDecisionSchema variant compatibility', () => { + const baseBenchmark = { + model: 'provider/model', + taskType: 'implementation' as const, + subtaskType: 'feature_development' as const, + source: 'benchmark' as const, + tableVersion: 'v1', + sticky: false, + }; + + it('parses a legacy decision with reasoningEffort only', () => { + const parsed = AutoRoutingDecisionSchema.parse({ + ...baseBenchmark, + reasoningEffort: 'high', + }); + expect(parsed).toMatchObject({ reasoningEffort: 'high' }); + expect('variant' in parsed ? parsed.variant : undefined).toBeUndefined(); + }); + + it('parses a decision with variant only', () => { + const parsed = AutoRoutingDecisionSchema.parse({ + ...baseBenchmark, + variant: 'xhigh', + }); + expect(parsed).toMatchObject({ variant: 'xhigh' }); + }); + + it('rejects a decision with both non-null variant and reasoningEffort', () => { + expect( + AutoRoutingDecisionSchema.safeParse({ + ...baseBenchmark, + variant: 'xhigh', + reasoningEffort: 'high', + }).success + ).toBe(false); + }); +}); + +describe('owner pool settings wire contracts', () => { + it('accepts null pool (inherit) on update', () => { + const parsed = UpdateAutoRoutingSettingsRequestSchema.parse({ + ownerType: 'user', + ownerId: 'user-1', + mode: null, + pool: null, + }); + expect(parsed.pool).toBeNull(); + }); + + it('accepts a configured pool on settings response with statuses', () => { + const parsed = AutoRoutingSettingsResponseSchema.parse({ + ownerType: 'org', + ownerId: 'org-1', + mode: 'cost_per_accuracy', + configuredMode: null, + defaultMode: 'cost_per_accuracy', + configuredPool: [{ model: 'a/b', variant: null }], + poolStatuses: [{ entry: { model: 'a/b', variant: null }, status: 'pending' }], + }); + expect(parsed.poolStatuses).toHaveLength(1); + }); + + it('accepts optional retryEntries that are a subset of pool by poolEntryKey', () => { + const pool = [ + { model: 'a/b', variant: null as string | null }, + { model: 'c/d', variant: 'xhigh' as string | null }, + ]; + const parsed = UpdateAutoRoutingSettingsRequestSchema.parse({ + ownerType: 'user', + ownerId: 'user-1', + mode: 'cost_per_accuracy', + pool, + retryEntries: [{ model: 'c/d', variant: 'xhigh' }], + }); + expect(parsed.retryEntries).toEqual([{ model: 'c/d', variant: 'xhigh' }]); + }); + + it('rejects retryEntries that are not in pool', () => { + const result = UpdateAutoRoutingSettingsRequestSchema.safeParse({ + ownerType: 'user', + ownerId: 'user-1', + mode: null, + pool: [{ model: 'a/b', variant: null }], + retryEntries: [{ model: 'other/model', variant: 'max' }], + }); + expect(result.success).toBe(false); + }); + + it('rejects retryEntries when pool is null', () => { + const result = UpdateAutoRoutingSettingsRequestSchema.safeParse({ + ownerType: 'user', + ownerId: 'user-1', + mode: null, + pool: null, + retryEntries: [{ model: 'a/b', variant: null }], + }); + expect(result.success).toBe(false); + }); +}); + +describe('benchmark profile request entry bounds', () => { + const eleven = Array.from({ length: 11 }, (_, i) => ({ + model: `model/${i}`, + variant: null as string | null, + })); + + it('bounds RegisterBenchmarkProfilesRequestSchema entries at 10', () => { + expect( + RegisterBenchmarkProfilesRequestSchema.safeParse({ + ownerType: 'user', + ownerId: 'u1', + entries: eleven, + }).success + ).toBe(false); + expect( + RegisterBenchmarkProfilesRequestSchema.safeParse({ + ownerType: 'user', + ownerId: 'u1', + entries: eleven.slice(0, 10), + }).success + ).toBe(true); + }); + + it('bounds BenchmarkProfileStatusesRequestSchema entries at 10', () => { + expect(BenchmarkProfileStatusesRequestSchema.safeParse({ entries: eleven }).success).toBe( + false + ); + expect( + BenchmarkProfileStatusesRequestSchema.safeParse({ entries: eleven.slice(0, 1) }).success + ).toBe(true); + }); + + it('accepts optional retryEntries that are a subset of entries', () => { + const entries = [ + { model: 'a/b', variant: null as string | null }, + { model: 'c/d', variant: 'xhigh' as string | null }, + ]; + expect( + RegisterBenchmarkProfilesRequestSchema.safeParse({ + ownerType: 'user', + ownerId: 'u1', + entries, + }).success + ).toBe(true); + expect( + RegisterBenchmarkProfilesRequestSchema.safeParse({ + ownerType: 'user', + ownerId: 'u1', + entries, + retryEntries: [{ model: 'a/b', variant: null }], + }).success + ).toBe(true); + }); + + it('rejects retryEntries that are not in entries', () => { + const result = RegisterBenchmarkProfilesRequestSchema.safeParse({ + ownerType: 'org', + ownerId: 'o1', + entries: [{ model: 'a/b', variant: null }], + retryEntries: [{ model: 'other/model', variant: 'max' }], + }); + expect(result.success).toBe(false); + }); +}); diff --git a/packages/auto-routing-contracts/src/index.ts b/packages/auto-routing-contracts/src/index.ts index 608f991a95..4415425f31 100644 --- a/packages/auto-routing-contracts/src/index.ts +++ b/packages/auto-routing-contracts/src/index.ts @@ -1,6 +1,8 @@ import * as z from 'zod'; +import { BenchmarkProfileEntryStatusSchema } from './benchmark'; import { NormalizedClassifierInputSchema } from './input'; import { ReasoningEffortSchema } from './reasoning'; +import { EfficientModelPoolSchema, PoolEntrySchema, poolEntryKey } from './routing-table'; import { ClassifierSubtaskTypeSchema, ClassifierTaskTypeSchema, @@ -99,26 +101,45 @@ export const ClassifierOutputSchema = z }); export type ClassifierOutput = z.infer; -const BenchmarkAutoRoutingDecisionSchema = z.object({ - model: z.string(), - taskType: ClassifierTaskTypeSchema, - subtaskType: ClassifierSubtaskTypeSchema, - source: z.literal('benchmark'), - tableVersion: z.string(), - // Mirrors the effort the chosen model was benchmarked with, when set. - reasoningEffort: ReasoningEffortSchema.nullable().optional(), - // True when the session's incumbent model was kept over a cheaper fresh - // pick. Defaulted so responses from a not-yet-redeployed worker still - // parse. - sticky: z.boolean().default(false), - // Why the session's incumbent model was abandoned, when it was: - // 'threshold' — the incumbent is denied, off the route, or below the - // accuracy bar; 'cost' — eligible, but the mode's switch condition made - // the fresh pick worth it; 'capability' — ejected by the modality/context - // capability filters. Null when there was no incumbent or it was kept. - // Optional so responses from a not-yet-redeployed worker still parse. - switchReason: z.enum(['threshold', 'cost', 'capability']).nullable().optional(), -}); +/** + * Reader precedence for decision variant identity: prefer `variant`; when + * absent, a valid legacy `reasoningEffort` string may be interpreted as the + * legacy variant key. New writers emit `variant` only; old writers emitted + * `reasoningEffort` only. Both non-null is malformed. + */ +const BenchmarkAutoRoutingDecisionSchema = z + .object({ + model: z.string(), + taskType: ClassifierTaskTypeSchema, + subtaskType: ClassifierSubtaskTypeSchema, + source: z.literal('benchmark'), + tableVersion: z.string(), + // Canonical catalog variant key the chosen model was benchmarked with. + variant: z.string().trim().min(1).nullable().optional(), + // Legacy effort the chosen model was benchmarked with; retained for + // rolling deploy reads. Prefer `variant` when present. + reasoningEffort: ReasoningEffortSchema.nullable().optional(), + // True when the session's incumbent model was kept over a cheaper fresh + // pick. Defaulted so responses from a not-yet-redeployed worker still + // parse. + sticky: z.boolean().default(false), + // Why the session's incumbent model was abandoned, when it was: + // 'threshold' — the incumbent is denied, off the route, or below the + // accuracy bar; 'cost' — eligible, but the mode's switch condition made + // the fresh pick worth it; 'capability' — ejected by the modality/context + // capability filters. Null when there was no incumbent or it was kept. + // Optional so responses from a not-yet-redeployed worker still parse. + switchReason: z.enum(['threshold', 'cost', 'capability']).nullable().optional(), + }) + .superRefine((decision, ctx) => { + if (decision.variant != null && decision.reasoningEffort != null) { + ctx.addIssue({ + code: 'custom', + path: ['variant'], + message: 'Decision must not set both variant and reasoningEffort; emit variant only', + }); + } + }); const CodingPlanDefaultDecisionSchema = z.object({ model: z.string(), @@ -188,6 +209,50 @@ export const AutoRoutingModeResponseSchema = AutoRoutingModeOwnerQuerySchema.ext }); export type AutoRoutingModeResponse = z.infer; +// Combined mode + Efficient model pool settings. Null on each field means +// inherit independently (mode and pool do not imply each other). +export const UpdateAutoRoutingSettingsRequestSchema = AutoRoutingModeOwnerQuerySchema.extend({ + mode: AutoRoutingModeSchema.nullable(), + pool: EfficientModelPoolSchema.nullable(), + // Subset of `pool` the owner explicitly retries after a terminal failure. + // Forwarded to benchmark profile registration when present. + retryEntries: z.array(PoolEntrySchema).optional(), +}).superRefine((request, ctx) => { + if (!request.retryEntries || request.retryEntries.length === 0) return; + if (request.pool === null) { + ctx.addIssue({ + code: 'custom', + path: ['retryEntries'], + message: 'retryEntries require a non-null pool', + }); + return; + } + const entryKeys = new Set(request.pool.map(poolEntryKey)); + request.retryEntries.forEach((entry, index) => { + const key = poolEntryKey(entry); + if (!entryKeys.has(key)) { + ctx.addIssue({ + code: 'custom', + path: ['retryEntries', index], + message: `retryEntry must appear in pool: ${key}`, + }); + } + }); +}); +export type UpdateAutoRoutingSettingsRequest = z.infer< + typeof UpdateAutoRoutingSettingsRequestSchema +>; + +export const AutoRoutingSettingsResponseSchema = AutoRoutingModeOwnerQuerySchema.extend({ + mode: AutoRoutingModeSchema, + configuredMode: AutoRoutingModeSchema.nullable(), + defaultMode: AutoRoutingModeSchema, + configuredPool: EfficientModelPoolSchema.nullable(), + // Empty when configuredPool is null. + poolStatuses: z.array(BenchmarkProfileEntryStatusSchema), +}); +export type AutoRoutingSettingsResponse = z.infer; + export const AutoRoutingAnalyticsPeriodSchema = z.enum(['1h', '24h', '7d', '30d']); export type AutoRoutingAnalyticsPeriod = z.infer; diff --git a/packages/auto-routing-contracts/src/routing-table.test.ts b/packages/auto-routing-contracts/src/routing-table.test.ts index a4830ce117..70aeea66ba 100644 --- a/packages/auto-routing-contracts/src/routing-table.test.ts +++ b/packages/auto-routing-contracts/src/routing-table.test.ts @@ -1,5 +1,13 @@ import { describe, expect, it } from 'vitest'; -import { rankCandidates, RoutingTableSchema } from './routing-table'; +import { + CustomRoutingTableSchema, + EfficientModelPoolSchema, + MAX_POOL_ENTRIES, + poolEntryKey, + rankCandidates, + RankedCandidateSchema, + RoutingTableSchema, +} from './routing-table'; const candidate = (model: string, accuracy: number, avgCostUsd: number) => ({ model, @@ -64,4 +72,134 @@ describe('RoutingTableSchema', () => { expect(parsed.routes['implementation/code_generation']?.[0]?.model).toBe('impl'); }); + + it('accepts a legacy candidate with reasoningEffort only', () => { + const parsed = RankedCandidateSchema.parse({ + model: 'm', + accuracy: 0.9, + avgCostUsd: 1, + meetsThreshold: true, + reasoningEffort: 'high', + }); + expect(parsed.reasoningEffort).toBe('high'); + expect(parsed.variant).toBeUndefined(); + }); + + it('accepts a candidate with variant only', () => { + const parsed = RankedCandidateSchema.parse({ + model: 'm', + accuracy: 0.9, + avgCostUsd: 1, + meetsThreshold: true, + variant: 'xhigh', + }); + expect(parsed.variant).toBe('xhigh'); + expect(parsed.reasoningEffort).toBeUndefined(); + }); + + it('rejects a candidate with both non-null variant and reasoningEffort', () => { + expect( + RankedCandidateSchema.safeParse({ + model: 'm', + accuracy: 0.9, + avgCostUsd: 1, + meetsThreshold: true, + variant: 'xhigh', + reasoningEffort: 'high', + }).success + ).toBe(false); + }); +}); + +describe('CustomRoutingTableSchema', () => { + const base = { + version: 'v', + generatedAt: new Date(0).toISOString(), + minAccuracy: 0.7, + switchCostFactor: 3, + source: 'benchmark' as const, + }; + + it('accepts a sparse routes record (omitted taxonomy keys)', () => { + const parsed = CustomRoutingTableSchema.parse({ + ...base, + routes: { + 'implementation/code_generation': [candidate('impl', 0.9, 1)], + }, + }); + expect(Object.keys(parsed.routes)).toEqual(['implementation/code_generation']); + }); + + it('rejects unknown route keys', () => { + expect( + CustomRoutingTableSchema.safeParse({ + ...base, + routes: { + 'not-a-real/route': [candidate('m', 1, 1)], + }, + }).success + ).toBe(false); + }); +}); + +describe('EfficientModelPoolSchema', () => { + it('rejects duplicate exact pairs', () => { + expect( + EfficientModelPoolSchema.safeParse([ + { model: 'a/b', variant: 'xhigh' }, + { model: 'a/b', variant: 'xhigh' }, + ]).success + ).toBe(false); + }); + + it('accepts the same model with different variants', () => { + const parsed = EfficientModelPoolSchema.parse([ + { model: 'a/b', variant: 'xhigh' }, + { model: 'a/b', variant: 'max' }, + ]); + expect(parsed).toHaveLength(2); + }); + + it('rejects zero entries', () => { + expect(EfficientModelPoolSchema.safeParse([]).success).toBe(false); + }); + + it('rejects more than MAX_POOL_ENTRIES entries', () => { + const entries = Array.from({ length: MAX_POOL_ENTRIES + 1 }, (_, i) => ({ + model: `model/${i}`, + variant: null, + })); + expect(EfficientModelPoolSchema.safeParse(entries).success).toBe(false); + }); + + it('accepts null pool via nullable wrapper (inherit)', () => { + expect(EfficientModelPoolSchema.nullable().safeParse(null).success).toBe(true); + }); + + it('accepts the maximum of 10 unique entries', () => { + const entries = Array.from({ length: MAX_POOL_ENTRIES }, (_, i) => ({ + model: `model/${i}`, + variant: null as string | null, + })); + expect(EfficientModelPoolSchema.safeParse(entries).success).toBe(true); + }); +}); + +describe('poolEntryKey', () => { + it('distinguishes xhigh, max, and null variants of one model', () => { + const model = 'provider/model'; + const keys = [ + poolEntryKey({ model, variant: 'xhigh' }), + poolEntryKey({ model, variant: 'max' }), + poolEntryKey({ model, variant: null }), + ]; + expect(new Set(keys).size).toBe(3); + }); + + it('round-trips through JSON as a collision-safe encoding', () => { + const entry = { model: 'provider/model', variant: 'xhigh' as string | null }; + const key = poolEntryKey(entry); + expect(JSON.parse(key)).toEqual([entry.model, entry.variant]); + expect(poolEntryKey({ model: entry.model, variant: entry.variant })).toBe(key); + }); }); diff --git a/packages/auto-routing-contracts/src/routing-table.ts b/packages/auto-routing-contracts/src/routing-table.ts index c3c8ce4118..3828b8c73e 100644 --- a/packages/auto-routing-contracts/src/routing-table.ts +++ b/packages/auto-routing-contracts/src/routing-table.ts @@ -2,20 +2,89 @@ import * as z from 'zod'; import { ReasoningEffortSchema } from './reasoning'; import { TaxonomyRouteKeySchema } from './taxonomy'; -export const RankedCandidateSchema = z.object({ +/** Maximum Pool entries in an owner-configured Efficient model pool. */ +export const MAX_POOL_ENTRIES = 10; + +/** + * One concrete managed model plus its canonical catalog variant key. + * `variant` is null only for models that expose no variants; catalog + * enforcement happens at the web boundary, not here. + */ +export const PoolEntrySchema = z.object({ model: z.string().trim().min(1), - // Benchmark accuracy in [0, 1] for this taxonomy route. - accuracy: z.number().min(0).max(1), - // Average observed OpenRouter cost per benchmark case, in USD credits. - avgCostUsd: z.number().nonnegative(), - meetsThreshold: z.boolean(), - // Reasoning effort the model was benchmarked with; serving mirrors it. - // Optional so tables published before this field existed stay valid. - reasoningEffort: ReasoningEffortSchema.nullable().optional(), + variant: z.string().trim().min(1).nullable(), }); +export type PoolEntry = z.infer; + +/** + * Collision-safe canonical key for a Pool entry. Prefer this over delimiter + * concatenation so model/variant values cannot collide across encoding. + */ +export function poolEntryKey(entry: PoolEntry): string { + return JSON.stringify([entry.model, entry.variant]); +} + +function addDuplicatePoolEntryIssues(entries: PoolEntry[], ctx: z.RefinementCtx): void { + const seen = new Set(); + entries.forEach((entry, index) => { + const key = poolEntryKey(entry); + if (seen.has(key)) { + ctx.addIssue({ + code: 'custom', + path: [index], + message: `Duplicate pool entry: ${key}`, + }); + } + seen.add(key); + }); +} + +/** + * Owner-configured Efficient model pool: 1–10 unique exact (model, variant) + * pairs. A null configured pool means inherit; an empty array is invalid. + */ +export const EfficientModelPoolSchema = z + .array(PoolEntrySchema) + .min(1) + .max(MAX_POOL_ENTRIES) + .superRefine((entries, ctx) => { + addDuplicatePoolEntryIssues(entries, ctx); + }); +export type EfficientModelPool = z.infer; + +/** + * Reader precedence for candidate variant identity: prefer `variant`; when + * absent, a valid legacy `reasoningEffort` string may be interpreted as the + * legacy variant key. New writers emit `variant` only; old writers emitted + * `reasoningEffort` only. Both non-null is malformed. + */ +export const RankedCandidateSchema = z + .object({ + model: z.string().trim().min(1), + // Benchmark accuracy in [0, 1] for this taxonomy route. + accuracy: z.number().min(0).max(1), + // Average observed OpenRouter cost per benchmark case, in USD credits. + avgCostUsd: z.number().nonnegative(), + meetsThreshold: z.boolean(), + // Canonical catalog variant key the model was benchmarked with. + // Optional so tables published before this field existed stay valid. + variant: z.string().trim().min(1).nullable().optional(), + // Legacy effort the model was benchmarked with; retained for rolling + // deploy reads of old published tables. Prefer `variant` when present. + reasoningEffort: ReasoningEffortSchema.nullable().optional(), + }) + .superRefine((candidate, ctx) => { + if (candidate.variant != null && candidate.reasoningEffort != null) { + ctx.addIssue({ + code: 'custom', + path: ['variant'], + message: 'Candidate must not set both variant and reasoningEffort; emit variant only', + }); + } + }); export type RankedCandidate = z.infer; -export const RoutingTableSchema = z.object({ +const routingTableFields = { // Benchmark run id. version: z.string().min(1), generatedAt: z.string().min(1), @@ -27,20 +96,43 @@ export const RoutingTableSchema = z.object({ // fresh pick improves accuracy by more than this absolute delta. bestAccuracySwitchThreshold: z.number().min(0).max(1).default(0.05), source: z.enum(['benchmark']), - routes: z.record(z.string(), z.array(RankedCandidateSchema).min(1)).superRefine((routes, ctx) => { - for (const key of Object.keys(routes)) { - if (!TaxonomyRouteKeySchema.safeParse(key).success) { - ctx.addIssue({ - code: 'custom', - path: [key], - message: `Unknown taxonomy route ${key}`, - }); - } +} as const; + +function refineTaxonomyRouteKeys(routes: Record, ctx: z.RefinementCtx): void { + for (const key of Object.keys(routes)) { + if (!TaxonomyRouteKeySchema.safeParse(key).success) { + ctx.addIssue({ + code: 'custom', + path: [key], + message: `Unknown taxonomy route ${key}`, + }); } + } +} + +/** Platform default routing table: every published route has ≥1 candidate. */ +export const RoutingTableSchema = z.object({ + ...routingTableFields, + routes: z.record(z.string(), z.array(RankedCandidateSchema).min(1)).superRefine((routes, ctx) => { + refineTaxonomyRouteKeys(routes, ctx); }), }); export type RoutingTable = z.infer; +/** + * Sparse custom routing table assembled for an owner Efficient model pool. + * Route keys with no graded candidates are omitted (not empty arrays). + * Consumers must return no decision for an omitted route (balanced fallback); + * existing `computeDecision` already does this. + */ +export const CustomRoutingTableSchema = z.object({ + ...routingTableFields, + routes: z.record(z.string(), z.array(RankedCandidateSchema).min(1)).superRefine((routes, ctx) => { + refineTaxonomyRouteKeys(routes, ctx); + }), +}); +export type CustomRoutingTable = z.infer; + export const ROUTING_TABLE_KV_KEY = 'routing_table_v1'; // "Best bang for buck": candidates meeting the accuracy threshold come first, diff --git a/services/auto-routing-benchmark/README.md b/services/auto-routing-benchmark/README.md index 228b8e410c..5de29308ed 100644 --- a/services/auto-routing-benchmark/README.md +++ b/services/auto-routing-benchmark/README.md @@ -28,10 +28,104 @@ All under `/admin`, gated by `Authorization: Bearer ` |---|---| | `GET/PUT /admin/config` | Read / save benchmark config (model lists, thresholds, `benchmarkUserId`, optional `benchmarkOrgId`) | | `GET /admin/runs` | List runs (sweeps stale `running` runs to `failed` first) | -| `POST /admin/runs` | Start a run (`{kind, force}`); returns 409 if one of that kind is already running | -| `GET /admin/routing-table` | Latest published routing table | +| `POST /admin/runs` | Start a **platform** run (`{kind, force}`); returns 409 if one of that kind is already running | +| `GET /admin/routing-table` | Latest published **platform** routing table | | `GET /admin/classifier-winner` | Current classifier winner | | `POST /admin/debug-cli` | Run one ad-hoc prompt through the kilo CLI container (diagnostic) | +| `POST /admin/profiles/register` | Atomically admit missing/stale/retried global Benchmark profiles for an owner (quota 10/24h) | +| `POST /admin/profiles/status` | Current statuses for up to 10 exact Pool entries (may free-admit engine-drifted rows) | +| `POST /admin/custom-routing-table` | Assemble a **sparse** custom table for ready/current entries only (`table: null` → balanced fallback) | + +## Owner pools and Benchmark profiles + +Owners (personal users or organizations) can constrain `kilo-auto/efficient` to +1–10 exact `(model, canonical variant)` **Pool entries**. The pool itself lives +in the auto-routing Durable Object; this worker owns **global Benchmark +profiles** — per-route measurements for each exact pair under the current +benchmark engine identity and decider repetitions. + +### Profile purpose runs vs platform runs + +| | Platform run | Profile run | +|---|---|---| +| Trigger | Admin `POST /admin/runs`, scheduled auto-decider sync | Single-slot drain of pending registry rows | +| `benchmark_runs.purpose` | `platform` (default) | `profile` | +| Model set | Saved admin config | Explicit entry snapshot from pending profiles | +| On completion | Publishes platform routing table / classifier winner; clears platform KV keys | Marks claimed profiles `ready` (or `failed`); **never** replaces the platform artifact | +| Slot | Shares the one-active-decider constraint | Same single slot; never preempts a platform run | + +Rollback: turning off owner pools (or clearing a pool) leaves the platform +default table untouched — profile runs never wrote it. + +### Request ledger and admission + +`POST /admin/profiles/register` evaluates every submitted entry in one D1 batch: + +- Ready/current or already pending/running global profiles are reported without + charging quota. +- Globally new, engine-stale, or explicitly retried-failed profiles are admitted + as `pending` while the owner has fewer than **10** charged admissions in the + rolling 24h window (`profile_request_events`). Over-limit → 429, nothing written. +- Concurrent owners requesting the same exact pair + engine identity dedupe to + one global row. + +Admission is all-or-nothing and returns before the caller's Durable Object +settings write. If that later write fails, the owner's previous pool is kept, +admitted profiles stay pending globally, the single-slot drain still runs them, +and a later retry reports them without re-charging quota. + +### Status transitions and drain + +``` +pending → running → ready + ↘ failed (bounded failure_reason; Retry re-admits) +``` + +- **running**: set when a profile run claims the entries at `startRun`. +- **ready**: set when that profile run completes successfully (`run_id` + provenance points at the measuring run). +- **failed**: set when the run fails (enqueue, timeout sweep, etc.). Only rows + still pointing at that `run_id` transition — a newer pending/ready row is + never clobbered. + +Currency: a profile is current only when `engine_identity`, `repetitions`, and +exact variant match the live decider engine. Stale rows are never returned as +Ready or assembled into custom tables. + +**Single-slot drain** (after any decider terminal state — completion *or* +failure — and when the scheduled handler has no platform start to claim): + +1. If a decider run is already active → log and leave pending rows alone. +2. Else take pending current-engine rows oldest-`requested_at`-first, capped by + the existing container budget (`entries × repetitions ≤ maxConcurrency`). +3. `startRun(purpose: 'profile', entries: snapshot)`. + +**Scheduled auto-decider sync ordering** (platform priority on a free slot): + +1. Sweep stale runs (cleanup only — not a slot claim). +2. Sync auto-decider candidates from the web API. +3. If models changed **and** a config exists → platform `startRun` claims the + free slot; **no** pending-profile drain this cycle (terminal transition + drains later). Profile work never preempts platform start. +4. If no platform start (no change, or no config) → drain pending profiles now + so stranded work recovers. +5. If the slot is already occupied → log/skip; leave pending rows untouched. + +Long waits stay `Benchmarking` in the UI; there is no second timeout state in v1. + +### Sparse custom routing tables + +`POST /admin/custom-routing-table` with 1–10 entries returns +`CustomRoutingTableResponse`: + +- Candidates come only from **ready + current** profiles' provenance + `model_summaries` for that entry's exact pair **and** measuring `run_id` + (no cross-run leakage), ranked with the saved policy knobs via + `rankCandidates`. +- Candidates carry exact `variant` (never `reasoningEffort`). +- Route keys with no graded candidates are **omitted** (not empty arrays). +- If no requested entry is ready/current → `{ table: null }` so the gateway + falls back to balanced. Never fabricate candidates. ## Local development diff --git a/services/auto-routing-benchmark/migrations/0005_fuzzy_senator_kelly.sql b/services/auto-routing-benchmark/migrations/0005_fuzzy_senator_kelly.sql new file mode 100644 index 0000000000..07fad61425 --- /dev/null +++ b/services/auto-routing-benchmark/migrations/0005_fuzzy_senator_kelly.sql @@ -0,0 +1,60 @@ +PRAGMA foreign_keys=OFF;--> statement-breakpoint +CREATE TABLE `__new_case_results` ( + `run_id` text NOT NULL, + `model` text NOT NULL, + `variant` text DEFAULT '' NOT NULL, + `case_id` text NOT NULL, + `route_key` text, + `score` real NOT NULL, + `latency_ms` integer NOT NULL, + `cost_usd` real, + `error` text, + `fallback_reason` text, + `retried` integer, + `exit_code` integer, + `output_prefix` text, + `event_count` integer, + `last_event_types` text, + `rep` integer DEFAULT 0 NOT NULL, + `timed_out` integer DEFAULT 0 NOT NULL, + PRIMARY KEY(`run_id`, `model`, `variant`, `case_id`, `rep`) +); +--> statement-breakpoint +INSERT INTO `__new_case_results`("run_id", "model", "variant", "case_id", "route_key", "score", "latency_ms", "cost_usd", "error", "fallback_reason", "retried", "exit_code", "output_prefix", "event_count", "last_event_types", "rep", "timed_out") SELECT cr."run_id", cr."model", COALESCE(rm."reasoning_effort", ''), cr."case_id", cr."route_key", cr."score", cr."latency_ms", cr."cost_usd", cr."error", cr."fallback_reason", cr."retried", cr."exit_code", cr."output_prefix", cr."event_count", cr."last_event_types", cr."rep", cr."timed_out" FROM `case_results` cr LEFT JOIN `run_models` rm ON rm."run_id" = cr."run_id" AND rm."model" = cr."model";--> statement-breakpoint +DROP TABLE `case_results`;--> statement-breakpoint +ALTER TABLE `__new_case_results` RENAME TO `case_results`;--> statement-breakpoint +PRAGMA foreign_keys=ON;--> statement-breakpoint +CREATE TABLE `__new_model_summaries` ( + `run_id` text NOT NULL, + `model` text NOT NULL, + `variant` text DEFAULT '' NOT NULL, + `route_key` text NOT NULL, + `accuracy` real NOT NULL, + `avg_cost_usd` real, + `avg_latency_ms` real NOT NULL, + `p50_latency_ms` real, + `cases` integer NOT NULL, + `errors` integer NOT NULL, + `p95_latency_ms` real, + `timeouts` integer DEFAULT 0 NOT NULL, + `carried` integer DEFAULT false NOT NULL, + PRIMARY KEY(`run_id`, `model`, `variant`, `route_key`) +); +--> statement-breakpoint +INSERT INTO `__new_model_summaries`("run_id", "model", "variant", "route_key", "accuracy", "avg_cost_usd", "avg_latency_ms", "p50_latency_ms", "cases", "errors", "p95_latency_ms", "timeouts", "carried") SELECT ms."run_id", ms."model", COALESCE(rm."reasoning_effort", ''), ms."route_key", ms."accuracy", ms."avg_cost_usd", ms."avg_latency_ms", ms."p50_latency_ms", ms."cases", ms."errors", ms."p95_latency_ms", ms."timeouts", ms."carried" FROM `model_summaries` ms LEFT JOIN `run_models` rm ON rm."run_id" = ms."run_id" AND rm."model" = ms."model";--> statement-breakpoint +DROP TABLE `model_summaries`;--> statement-breakpoint +ALTER TABLE `__new_model_summaries` RENAME TO `model_summaries`;--> statement-breakpoint +CREATE TABLE `__new_run_models` ( + `run_id` text NOT NULL, + `model` text NOT NULL, + `variant` text DEFAULT '' NOT NULL, + `enqueued` integer NOT NULL, + `reasoning_effort` text, + PRIMARY KEY(`run_id`, `model`, `variant`) +); +--> statement-breakpoint +INSERT INTO `__new_run_models`("run_id", "model", "variant", "enqueued", "reasoning_effort") SELECT "run_id", "model", COALESCE("reasoning_effort", ''), "enqueued", "reasoning_effort" FROM `run_models`;--> statement-breakpoint +DROP TABLE `run_models`;--> statement-breakpoint +ALTER TABLE `__new_run_models` RENAME TO `run_models`;--> statement-breakpoint +ALTER TABLE `routing_table_candidates` ADD `variant` text DEFAULT '' NOT NULL;--> statement-breakpoint +UPDATE `routing_table_candidates` SET `variant` = COALESCE(`reasoning_effort`, '') WHERE `reasoning_effort` IS NOT NULL; diff --git a/services/auto-routing-benchmark/migrations/0006_hard_blizzard.sql b/services/auto-routing-benchmark/migrations/0006_hard_blizzard.sql new file mode 100644 index 0000000000..737752451e --- /dev/null +++ b/services/auto-routing-benchmark/migrations/0006_hard_blizzard.sql @@ -0,0 +1,26 @@ +CREATE TABLE `benchmark_profiles` ( + `model` text NOT NULL, + `variant` text DEFAULT '' NOT NULL, + `engine_identity` text NOT NULL, + `repetitions` integer NOT NULL, + `status` text NOT NULL, + `run_id` text, + `failure_reason` text, + `requested_at` text NOT NULL, + `updated_at` text NOT NULL, + `completed_at` text, + PRIMARY KEY(`model`, `variant`, `engine_identity`, `repetitions`) +); +--> statement-breakpoint +CREATE TABLE `profile_request_events` ( + `id` integer PRIMARY KEY AUTOINCREMENT NOT NULL, + `owner_type` text NOT NULL, + `owner_id` text NOT NULL, + `model` text NOT NULL, + `variant` text DEFAULT '' NOT NULL, + `engine_identity` text NOT NULL, + `repetitions` integer NOT NULL, + `admitted_at` text NOT NULL +); +--> statement-breakpoint +CREATE INDEX `IDX_profile_request_events_owner_admitted` ON `profile_request_events` (`owner_type`,`owner_id`,`admitted_at`); \ No newline at end of file diff --git a/services/auto-routing-benchmark/migrations/0007_complex_makkari.sql b/services/auto-routing-benchmark/migrations/0007_complex_makkari.sql new file mode 100644 index 0000000000..495b185116 --- /dev/null +++ b/services/auto-routing-benchmark/migrations/0007_complex_makkari.sql @@ -0,0 +1 @@ +ALTER TABLE `benchmark_runs` ADD `purpose` text DEFAULT 'platform' NOT NULL; \ No newline at end of file diff --git a/services/auto-routing-benchmark/migrations/meta/0005_snapshot.json b/services/auto-routing-benchmark/migrations/meta/0005_snapshot.json new file mode 100644 index 0000000000..72c3fd94ae --- /dev/null +++ b/services/auto-routing-benchmark/migrations/meta/0005_snapshot.json @@ -0,0 +1,809 @@ +{ + "version": "6", + "dialect": "sqlite", + "id": "deb1ce98-4003-477c-9ac2-9c6f423506a7", + "prevId": "5a8061af-29a3-47f4-bba8-8c84c9ed2901", + "tables": { + "benchmark_config": { + "name": "benchmark_config", + "columns": { + "id": { + "name": "id", + "type": "integer", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "min_accuracy": { + "name": "min_accuracy", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "switch_cost_factor": { + "name": "switch_cost_factor", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "best_accuracy_switch_threshold": { + "name": "best_accuracy_switch_threshold", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0.05 + }, + "max_concurrency": { + "name": "max_concurrency", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "benchmark_user_id": { + "name": "benchmark_user_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "benchmark_org_id": { + "name": "benchmark_org_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "classifier_repetitions": { + "name": "classifier_repetitions", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 1 + }, + "decider_repetitions": { + "name": "decider_repetitions", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 1 + }, + "classifier_max_p95_latency_ms": { + "name": "classifier_max_p95_latency_ms", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "auto_decider_min_cost_usd": { + "name": "auto_decider_min_cost_usd", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 15 + }, + "auto_decider_max_cost_usd": { + "name": "auto_decider_max_cost_usd", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 25 + }, + "updated_at": { + "name": "updated_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "updated_by": { + "name": "updated_by", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "benchmark_runs": { + "name": "benchmark_runs", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "kind": { + "name": "kind", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "started_at": { + "name": "started_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "completed_at": { + "name": "completed_at", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "error": { + "name": "error", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "min_accuracy": { + "name": "min_accuracy", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "switch_cost_factor": { + "name": "switch_cost_factor", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "best_accuracy_switch_threshold": { + "name": "best_accuracy_switch_threshold", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0.05 + }, + "max_concurrency": { + "name": "max_concurrency", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "benchmark_user_id": { + "name": "benchmark_user_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "benchmark_org_id": { + "name": "benchmark_org_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "repetitions": { + "name": "repetitions", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 1 + }, + "classifier_max_p95_latency_ms": { + "name": "classifier_max_p95_latency_ms", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "engine_identity": { + "name": "engine_identity", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + } + }, + "indexes": { + "UQ_benchmark_runs_one_running_per_kind": { + "name": "UQ_benchmark_runs_one_running_per_kind", + "columns": [ + "kind" + ], + "isUnique": true, + "where": "\"benchmark_runs\".\"status\" = 'running'" + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "case_results": { + "name": "case_results", + "columns": { + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "variant": { + "name": "variant", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + }, + "case_id": { + "name": "case_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "route_key": { + "name": "route_key", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "score": { + "name": "score", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "latency_ms": { + "name": "latency_ms", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "cost_usd": { + "name": "cost_usd", + "type": "real", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "error": { + "name": "error", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "fallback_reason": { + "name": "fallback_reason", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "retried": { + "name": "retried", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "exit_code": { + "name": "exit_code", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "output_prefix": { + "name": "output_prefix", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "event_count": { + "name": "event_count", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "last_event_types": { + "name": "last_event_types", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "rep": { + "name": "rep", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0 + }, + "timed_out": { + "name": "timed_out", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0 + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "case_results_run_id_model_variant_case_id_rep_pk": { + "columns": [ + "run_id", + "model", + "variant", + "case_id", + "rep" + ], + "name": "case_results_run_id_model_variant_case_id_rep_pk" + } + }, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "config_auto_decider_exclusions": { + "name": "config_auto_decider_exclusions", + "columns": { + "model": { + "name": "model", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "config_auto_decider_models": { + "name": "config_auto_decider_models", + "columns": { + "model": { + "name": "model", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "reasoning_effort": { + "name": "reasoning_effort", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "avg_attempt_cost_usd": { + "name": "avg_attempt_cost_usd", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "synced_at": { + "name": "synced_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "config_classifier_models": { + "name": "config_classifier_models", + "columns": { + "model": { + "name": "model", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "config_decider_models": { + "name": "config_decider_models", + "columns": { + "model": { + "name": "model", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "reasoning_effort": { + "name": "reasoning_effort", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "model_summaries": { + "name": "model_summaries", + "columns": { + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "variant": { + "name": "variant", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + }, + "route_key": { + "name": "route_key", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "accuracy": { + "name": "accuracy", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "avg_cost_usd": { + "name": "avg_cost_usd", + "type": "real", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "avg_latency_ms": { + "name": "avg_latency_ms", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "p50_latency_ms": { + "name": "p50_latency_ms", + "type": "real", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "cases": { + "name": "cases", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "errors": { + "name": "errors", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "p95_latency_ms": { + "name": "p95_latency_ms", + "type": "real", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "timeouts": { + "name": "timeouts", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0 + }, + "carried": { + "name": "carried", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "model_summaries_run_id_model_variant_route_key_pk": { + "columns": [ + "run_id", + "model", + "variant", + "route_key" + ], + "name": "model_summaries_run_id_model_variant_route_key_pk" + } + }, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "routing_table_candidates": { + "name": "routing_table_candidates", + "columns": { + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "route_key": { + "name": "route_key", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "rank": { + "name": "rank", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "accuracy": { + "name": "accuracy", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "avg_cost_usd": { + "name": "avg_cost_usd", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "meets_threshold": { + "name": "meets_threshold", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "reasoning_effort": { + "name": "reasoning_effort", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "variant": { + "name": "variant", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "routing_table_candidates_run_id_route_key_rank_pk": { + "columns": [ + "run_id", + "route_key", + "rank" + ], + "name": "routing_table_candidates_run_id_route_key_rank_pk" + } + }, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "routing_tables": { + "name": "routing_tables", + "columns": { + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "published_at": { + "name": "published_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "generated_at": { + "name": "generated_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "min_accuracy": { + "name": "min_accuracy", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "switch_cost_factor": { + "name": "switch_cost_factor", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "best_accuracy_switch_threshold": { + "name": "best_accuracy_switch_threshold", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0.05 + }, + "source": { + "name": "source", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "run_models": { + "name": "run_models", + "columns": { + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "variant": { + "name": "variant", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + }, + "enqueued": { + "name": "enqueued", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "reasoning_effort": { + "name": "reasoning_effort", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "run_models_run_id_model_variant_pk": { + "columns": [ + "run_id", + "model", + "variant" + ], + "name": "run_models_run_id_model_variant_pk" + } + }, + "uniqueConstraints": {}, + "checkConstraints": {} + } + }, + "views": {}, + "enums": {}, + "_meta": { + "schemas": {}, + "tables": {}, + "columns": {} + }, + "internal": { + "indexes": {} + } +} \ No newline at end of file diff --git a/services/auto-routing-benchmark/migrations/meta/0006_snapshot.json b/services/auto-routing-benchmark/migrations/meta/0006_snapshot.json new file mode 100644 index 0000000000..34fc7e9eff --- /dev/null +++ b/services/auto-routing-benchmark/migrations/meta/0006_snapshot.json @@ -0,0 +1,977 @@ +{ + "version": "6", + "dialect": "sqlite", + "id": "0c106329-0bcc-4b5c-aa87-e2921771ef98", + "prevId": "deb1ce98-4003-477c-9ac2-9c6f423506a7", + "tables": { + "benchmark_config": { + "name": "benchmark_config", + "columns": { + "id": { + "name": "id", + "type": "integer", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "min_accuracy": { + "name": "min_accuracy", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "switch_cost_factor": { + "name": "switch_cost_factor", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "best_accuracy_switch_threshold": { + "name": "best_accuracy_switch_threshold", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0.05 + }, + "max_concurrency": { + "name": "max_concurrency", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "benchmark_user_id": { + "name": "benchmark_user_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "benchmark_org_id": { + "name": "benchmark_org_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "classifier_repetitions": { + "name": "classifier_repetitions", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 1 + }, + "decider_repetitions": { + "name": "decider_repetitions", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 1 + }, + "classifier_max_p95_latency_ms": { + "name": "classifier_max_p95_latency_ms", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "auto_decider_min_cost_usd": { + "name": "auto_decider_min_cost_usd", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 15 + }, + "auto_decider_max_cost_usd": { + "name": "auto_decider_max_cost_usd", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 25 + }, + "updated_at": { + "name": "updated_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "updated_by": { + "name": "updated_by", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "benchmark_profiles": { + "name": "benchmark_profiles", + "columns": { + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "variant": { + "name": "variant", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + }, + "engine_identity": { + "name": "engine_identity", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "repetitions": { + "name": "repetitions", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "failure_reason": { + "name": "failure_reason", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "requested_at": { + "name": "requested_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "updated_at": { + "name": "updated_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "completed_at": { + "name": "completed_at", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "benchmark_profiles_model_variant_engine_identity_repetitions_pk": { + "columns": [ + "model", + "variant", + "engine_identity", + "repetitions" + ], + "name": "benchmark_profiles_model_variant_engine_identity_repetitions_pk" + } + }, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "benchmark_runs": { + "name": "benchmark_runs", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "kind": { + "name": "kind", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "started_at": { + "name": "started_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "completed_at": { + "name": "completed_at", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "error": { + "name": "error", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "min_accuracy": { + "name": "min_accuracy", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "switch_cost_factor": { + "name": "switch_cost_factor", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "best_accuracy_switch_threshold": { + "name": "best_accuracy_switch_threshold", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0.05 + }, + "max_concurrency": { + "name": "max_concurrency", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "benchmark_user_id": { + "name": "benchmark_user_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "benchmark_org_id": { + "name": "benchmark_org_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "repetitions": { + "name": "repetitions", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 1 + }, + "classifier_max_p95_latency_ms": { + "name": "classifier_max_p95_latency_ms", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "engine_identity": { + "name": "engine_identity", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + } + }, + "indexes": { + "UQ_benchmark_runs_one_running_per_kind": { + "name": "UQ_benchmark_runs_one_running_per_kind", + "columns": [ + "kind" + ], + "isUnique": true, + "where": "\"benchmark_runs\".\"status\" = 'running'" + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "case_results": { + "name": "case_results", + "columns": { + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "variant": { + "name": "variant", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + }, + "case_id": { + "name": "case_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "route_key": { + "name": "route_key", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "score": { + "name": "score", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "latency_ms": { + "name": "latency_ms", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "cost_usd": { + "name": "cost_usd", + "type": "real", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "error": { + "name": "error", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "fallback_reason": { + "name": "fallback_reason", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "retried": { + "name": "retried", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "exit_code": { + "name": "exit_code", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "output_prefix": { + "name": "output_prefix", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "event_count": { + "name": "event_count", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "last_event_types": { + "name": "last_event_types", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "rep": { + "name": "rep", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0 + }, + "timed_out": { + "name": "timed_out", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0 + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "case_results_run_id_model_variant_case_id_rep_pk": { + "columns": [ + "run_id", + "model", + "variant", + "case_id", + "rep" + ], + "name": "case_results_run_id_model_variant_case_id_rep_pk" + } + }, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "config_auto_decider_exclusions": { + "name": "config_auto_decider_exclusions", + "columns": { + "model": { + "name": "model", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "config_auto_decider_models": { + "name": "config_auto_decider_models", + "columns": { + "model": { + "name": "model", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "reasoning_effort": { + "name": "reasoning_effort", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "avg_attempt_cost_usd": { + "name": "avg_attempt_cost_usd", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "synced_at": { + "name": "synced_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "config_classifier_models": { + "name": "config_classifier_models", + "columns": { + "model": { + "name": "model", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "config_decider_models": { + "name": "config_decider_models", + "columns": { + "model": { + "name": "model", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "reasoning_effort": { + "name": "reasoning_effort", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "model_summaries": { + "name": "model_summaries", + "columns": { + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "variant": { + "name": "variant", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + }, + "route_key": { + "name": "route_key", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "accuracy": { + "name": "accuracy", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "avg_cost_usd": { + "name": "avg_cost_usd", + "type": "real", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "avg_latency_ms": { + "name": "avg_latency_ms", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "p50_latency_ms": { + "name": "p50_latency_ms", + "type": "real", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "cases": { + "name": "cases", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "errors": { + "name": "errors", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "p95_latency_ms": { + "name": "p95_latency_ms", + "type": "real", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "timeouts": { + "name": "timeouts", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0 + }, + "carried": { + "name": "carried", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "model_summaries_run_id_model_variant_route_key_pk": { + "columns": [ + "run_id", + "model", + "variant", + "route_key" + ], + "name": "model_summaries_run_id_model_variant_route_key_pk" + } + }, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "profile_request_events": { + "name": "profile_request_events", + "columns": { + "id": { + "name": "id", + "type": "integer", + "primaryKey": true, + "notNull": true, + "autoincrement": true + }, + "owner_type": { + "name": "owner_type", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "owner_id": { + "name": "owner_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "variant": { + "name": "variant", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + }, + "engine_identity": { + "name": "engine_identity", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "repetitions": { + "name": "repetitions", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "admitted_at": { + "name": "admitted_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": { + "IDX_profile_request_events_owner_admitted": { + "name": "IDX_profile_request_events_owner_admitted", + "columns": [ + "owner_type", + "owner_id", + "admitted_at" + ], + "isUnique": false + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "routing_table_candidates": { + "name": "routing_table_candidates", + "columns": { + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "route_key": { + "name": "route_key", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "rank": { + "name": "rank", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "accuracy": { + "name": "accuracy", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "avg_cost_usd": { + "name": "avg_cost_usd", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "meets_threshold": { + "name": "meets_threshold", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "reasoning_effort": { + "name": "reasoning_effort", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "variant": { + "name": "variant", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "routing_table_candidates_run_id_route_key_rank_pk": { + "columns": [ + "run_id", + "route_key", + "rank" + ], + "name": "routing_table_candidates_run_id_route_key_rank_pk" + } + }, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "routing_tables": { + "name": "routing_tables", + "columns": { + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "published_at": { + "name": "published_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "generated_at": { + "name": "generated_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "min_accuracy": { + "name": "min_accuracy", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "switch_cost_factor": { + "name": "switch_cost_factor", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "best_accuracy_switch_threshold": { + "name": "best_accuracy_switch_threshold", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0.05 + }, + "source": { + "name": "source", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "run_models": { + "name": "run_models", + "columns": { + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "variant": { + "name": "variant", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + }, + "enqueued": { + "name": "enqueued", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "reasoning_effort": { + "name": "reasoning_effort", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "run_models_run_id_model_variant_pk": { + "columns": [ + "run_id", + "model", + "variant" + ], + "name": "run_models_run_id_model_variant_pk" + } + }, + "uniqueConstraints": {}, + "checkConstraints": {} + } + }, + "views": {}, + "enums": {}, + "_meta": { + "schemas": {}, + "tables": {}, + "columns": {} + }, + "internal": { + "indexes": {} + } +} \ No newline at end of file diff --git a/services/auto-routing-benchmark/migrations/meta/0007_snapshot.json b/services/auto-routing-benchmark/migrations/meta/0007_snapshot.json new file mode 100644 index 0000000000..f5698f2871 --- /dev/null +++ b/services/auto-routing-benchmark/migrations/meta/0007_snapshot.json @@ -0,0 +1,985 @@ +{ + "version": "6", + "dialect": "sqlite", + "id": "d374401e-a36e-48bc-8a5a-6d98690c356b", + "prevId": "0c106329-0bcc-4b5c-aa87-e2921771ef98", + "tables": { + "benchmark_config": { + "name": "benchmark_config", + "columns": { + "id": { + "name": "id", + "type": "integer", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "min_accuracy": { + "name": "min_accuracy", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "switch_cost_factor": { + "name": "switch_cost_factor", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "best_accuracy_switch_threshold": { + "name": "best_accuracy_switch_threshold", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0.05 + }, + "max_concurrency": { + "name": "max_concurrency", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "benchmark_user_id": { + "name": "benchmark_user_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "benchmark_org_id": { + "name": "benchmark_org_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "classifier_repetitions": { + "name": "classifier_repetitions", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 1 + }, + "decider_repetitions": { + "name": "decider_repetitions", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 1 + }, + "classifier_max_p95_latency_ms": { + "name": "classifier_max_p95_latency_ms", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "auto_decider_min_cost_usd": { + "name": "auto_decider_min_cost_usd", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 15 + }, + "auto_decider_max_cost_usd": { + "name": "auto_decider_max_cost_usd", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 25 + }, + "updated_at": { + "name": "updated_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "updated_by": { + "name": "updated_by", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "benchmark_profiles": { + "name": "benchmark_profiles", + "columns": { + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "variant": { + "name": "variant", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + }, + "engine_identity": { + "name": "engine_identity", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "repetitions": { + "name": "repetitions", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "failure_reason": { + "name": "failure_reason", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "requested_at": { + "name": "requested_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "updated_at": { + "name": "updated_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "completed_at": { + "name": "completed_at", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "benchmark_profiles_model_variant_engine_identity_repetitions_pk": { + "columns": [ + "model", + "variant", + "engine_identity", + "repetitions" + ], + "name": "benchmark_profiles_model_variant_engine_identity_repetitions_pk" + } + }, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "benchmark_runs": { + "name": "benchmark_runs", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "kind": { + "name": "kind", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "started_at": { + "name": "started_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "completed_at": { + "name": "completed_at", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "error": { + "name": "error", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "min_accuracy": { + "name": "min_accuracy", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "switch_cost_factor": { + "name": "switch_cost_factor", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "best_accuracy_switch_threshold": { + "name": "best_accuracy_switch_threshold", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0.05 + }, + "max_concurrency": { + "name": "max_concurrency", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "benchmark_user_id": { + "name": "benchmark_user_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "benchmark_org_id": { + "name": "benchmark_org_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "repetitions": { + "name": "repetitions", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 1 + }, + "classifier_max_p95_latency_ms": { + "name": "classifier_max_p95_latency_ms", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "engine_identity": { + "name": "engine_identity", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + }, + "purpose": { + "name": "purpose", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "'platform'" + } + }, + "indexes": { + "UQ_benchmark_runs_one_running_per_kind": { + "name": "UQ_benchmark_runs_one_running_per_kind", + "columns": [ + "kind" + ], + "isUnique": true, + "where": "\"benchmark_runs\".\"status\" = 'running'" + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "case_results": { + "name": "case_results", + "columns": { + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "variant": { + "name": "variant", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + }, + "case_id": { + "name": "case_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "route_key": { + "name": "route_key", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "score": { + "name": "score", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "latency_ms": { + "name": "latency_ms", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "cost_usd": { + "name": "cost_usd", + "type": "real", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "error": { + "name": "error", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "fallback_reason": { + "name": "fallback_reason", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "retried": { + "name": "retried", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "exit_code": { + "name": "exit_code", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "output_prefix": { + "name": "output_prefix", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "event_count": { + "name": "event_count", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "last_event_types": { + "name": "last_event_types", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "rep": { + "name": "rep", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0 + }, + "timed_out": { + "name": "timed_out", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0 + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "case_results_run_id_model_variant_case_id_rep_pk": { + "columns": [ + "run_id", + "model", + "variant", + "case_id", + "rep" + ], + "name": "case_results_run_id_model_variant_case_id_rep_pk" + } + }, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "config_auto_decider_exclusions": { + "name": "config_auto_decider_exclusions", + "columns": { + "model": { + "name": "model", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "config_auto_decider_models": { + "name": "config_auto_decider_models", + "columns": { + "model": { + "name": "model", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "reasoning_effort": { + "name": "reasoning_effort", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "avg_attempt_cost_usd": { + "name": "avg_attempt_cost_usd", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "synced_at": { + "name": "synced_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "config_classifier_models": { + "name": "config_classifier_models", + "columns": { + "model": { + "name": "model", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "config_decider_models": { + "name": "config_decider_models", + "columns": { + "model": { + "name": "model", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "reasoning_effort": { + "name": "reasoning_effort", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "model_summaries": { + "name": "model_summaries", + "columns": { + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "variant": { + "name": "variant", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + }, + "route_key": { + "name": "route_key", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "accuracy": { + "name": "accuracy", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "avg_cost_usd": { + "name": "avg_cost_usd", + "type": "real", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "avg_latency_ms": { + "name": "avg_latency_ms", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "p50_latency_ms": { + "name": "p50_latency_ms", + "type": "real", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "cases": { + "name": "cases", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "errors": { + "name": "errors", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "p95_latency_ms": { + "name": "p95_latency_ms", + "type": "real", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "timeouts": { + "name": "timeouts", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0 + }, + "carried": { + "name": "carried", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "model_summaries_run_id_model_variant_route_key_pk": { + "columns": [ + "run_id", + "model", + "variant", + "route_key" + ], + "name": "model_summaries_run_id_model_variant_route_key_pk" + } + }, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "profile_request_events": { + "name": "profile_request_events", + "columns": { + "id": { + "name": "id", + "type": "integer", + "primaryKey": true, + "notNull": true, + "autoincrement": true + }, + "owner_type": { + "name": "owner_type", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "owner_id": { + "name": "owner_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "variant": { + "name": "variant", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + }, + "engine_identity": { + "name": "engine_identity", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "repetitions": { + "name": "repetitions", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "admitted_at": { + "name": "admitted_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": { + "IDX_profile_request_events_owner_admitted": { + "name": "IDX_profile_request_events_owner_admitted", + "columns": [ + "owner_type", + "owner_id", + "admitted_at" + ], + "isUnique": false + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "routing_table_candidates": { + "name": "routing_table_candidates", + "columns": { + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "route_key": { + "name": "route_key", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "rank": { + "name": "rank", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "accuracy": { + "name": "accuracy", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "avg_cost_usd": { + "name": "avg_cost_usd", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "meets_threshold": { + "name": "meets_threshold", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "reasoning_effort": { + "name": "reasoning_effort", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "variant": { + "name": "variant", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "routing_table_candidates_run_id_route_key_rank_pk": { + "columns": [ + "run_id", + "route_key", + "rank" + ], + "name": "routing_table_candidates_run_id_route_key_rank_pk" + } + }, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "routing_tables": { + "name": "routing_tables", + "columns": { + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "published_at": { + "name": "published_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "generated_at": { + "name": "generated_at", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "min_accuracy": { + "name": "min_accuracy", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "switch_cost_factor": { + "name": "switch_cost_factor", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "best_accuracy_switch_threshold": { + "name": "best_accuracy_switch_threshold", + "type": "real", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": 0.05 + }, + "source": { + "name": "source", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "run_models": { + "name": "run_models", + "columns": { + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "variant": { + "name": "variant", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false, + "default": "''" + }, + "enqueued": { + "name": "enqueued", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "reasoning_effort": { + "name": "reasoning_effort", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "run_models_run_id_model_variant_pk": { + "columns": [ + "run_id", + "model", + "variant" + ], + "name": "run_models_run_id_model_variant_pk" + } + }, + "uniqueConstraints": {}, + "checkConstraints": {} + } + }, + "views": {}, + "enums": {}, + "_meta": { + "schemas": {}, + "tables": {}, + "columns": {} + }, + "internal": { + "indexes": {} + } +} \ No newline at end of file diff --git a/services/auto-routing-benchmark/migrations/meta/_journal.json b/services/auto-routing-benchmark/migrations/meta/_journal.json index bf98e68c75..1ac59b6af9 100644 --- a/services/auto-routing-benchmark/migrations/meta/_journal.json +++ b/services/auto-routing-benchmark/migrations/meta/_journal.json @@ -36,6 +36,27 @@ "when": 1781886063127, "tag": "0004_stiff_talos", "breakpoints": true + }, + { + "idx": 5, + "version": "6", + "when": 1785248685068, + "tag": "0005_fuzzy_senator_kelly", + "breakpoints": true + }, + { + "idx": 6, + "version": "6", + "when": 1785253984099, + "tag": "0006_hard_blizzard", + "breakpoints": true + }, + { + "idx": 7, + "version": "6", + "when": 1785260908640, + "tag": "0007_complex_makkari", + "breakpoints": true } ] } \ No newline at end of file diff --git a/services/auto-routing-benchmark/src/admin.test.ts b/services/auto-routing-benchmark/src/admin.test.ts index 598cbb290e..9b88914fa8 100644 --- a/services/auto-routing-benchmark/src/admin.test.ts +++ b/services/auto-routing-benchmark/src/admin.test.ts @@ -9,6 +9,7 @@ import { import { app } from './index'; import { computeEngineIdentity } from './run'; import type * as DbModule from './db'; +import type * as ProfilesModule from './profiles'; import { CLASSIFIER_CASES } from './datasets/classifier-cases'; function makeSummary(model: string): BenchmarkModelSummary { @@ -91,12 +92,27 @@ vi.mock('./db', async importOriginal => { getLatestSummariesByModel: vi.fn(), insertRun: vi.fn(), markStaleRunsFailed: vi.fn(), + listStaleRunningDeciderRunIds: vi.fn(), + listPendingCurrentProfiles: vi.fn(), + markProfilesFailedForRun: vi.fn(), + markProfilesRunningForRun: vi.fn(), + markProfilesReadyForRun: vi.fn(), getRunningRun: vi.fn(), existsNewerCompletedRun: vi.fn(), }; }); +vi.mock('./profiles', async importOriginal => { + const actual = await importOriginal(); + return { + ...actual, + registerProfiles: vi.fn(), + lookupProfileStatuses: vi.fn(), + }; +}); + import { + exactPairKey, getConfigRows, getClassifierWinner, getLatestRoutingTable, @@ -104,10 +120,13 @@ import { getRunningRun, existsNewerCompletedRun, insertRun, + listPendingCurrentProfiles, listRuns, + listStaleRunningDeciderRunIds, markStaleRunsFailed, replaceConfig, } from './db'; +import { lookupProfileStatuses, ProfileQuotaExceededError, registerProfiles } from './profiles'; const tokenGet = vi.fn<() => Promise>(); const queueSendBatch = vi.fn(); @@ -173,8 +192,12 @@ beforeEach(() => { vi.mocked(getLatestSummariesByModel).mockResolvedValue(new Map()); vi.mocked(insertRun).mockResolvedValue(undefined); vi.mocked(markStaleRunsFailed).mockResolvedValue(undefined); + vi.mocked(listStaleRunningDeciderRunIds).mockResolvedValue([]); + vi.mocked(listPendingCurrentProfiles).mockResolvedValue([]); vi.mocked(getRunningRun).mockResolvedValue(undefined); vi.mocked(existsNewerCompletedRun).mockResolvedValue(false); + vi.mocked(registerProfiles).mockReset(); + vi.mocked(lookupProfileStatuses).mockReset(); queueSendBatch.mockResolvedValue(undefined); }); @@ -394,6 +417,7 @@ describe('POST /admin/runs', () => { repetitions: 1, classifier_max_p95_latency_ms: 1000, engine_identity: 'v1:deadbeef', + purpose: 'platform', }); const res = await authedPost('/admin/runs', { kind: 'classifier' }); @@ -480,10 +504,11 @@ describe('POST /admin/runs', () => { vi.mocked(getLatestSummariesByModel).mockResolvedValue( new Map([ [ - 'vendor/a', + exactPairKey('vendor/a', null), { engineIdentity: computeEngineIdentity('decider'), repetitions: 1, + variant: null, reasoningEffort: null, summaries: [makeSummary('vendor/a')], }, @@ -504,15 +529,16 @@ describe('POST /admin/runs', () => { config: { ...TEST_CONFIG_ROWS.config, benchmark_user_id: 'user-123' }, deciderModels: [{ model: 'vendor/a', reasoning_effort: null }], }); - // Prior result was measured at reasoning_effort 'high'; current config runs - // it at null, so the carry is invalidated and the model is re-enqueued. + // Prior result was measured at variant 'high'; current config runs it at + // null, so the carry is invalidated and the model is re-enqueued. vi.mocked(getLatestSummariesByModel).mockResolvedValue( new Map([ [ - 'vendor/a', + exactPairKey('vendor/a', 'high'), { engineIdentity: computeEngineIdentity('decider'), repetitions: 1, + variant: 'high', reasoningEffort: 'high', summaries: [makeSummary('vendor/a')], }, @@ -527,6 +553,69 @@ describe('POST /admin/runs', () => { expect(body.enqueuedModels).toBe(1); }); + it('does not carry a prior summary when only the variant differs', async () => { + vi.mocked(getConfigRows).mockResolvedValue({ + ...TEST_CONFIG_ROWS, + config: { ...TEST_CONFIG_ROWS.config, benchmark_user_id: 'user-123' }, + deciderModels: [{ model: 'vendor/a', reasoning_effort: 'high' }], + }); + // Same model/engine/reps but prior was measured at a different variant. + vi.mocked(getLatestSummariesByModel).mockResolvedValue( + new Map([ + [ + exactPairKey('vendor/a', 'low'), + { + engineIdentity: computeEngineIdentity('decider'), + repetitions: 1, + variant: 'low', + reasoningEffort: 'low', + summaries: [makeSummary('vendor/a')], + }, + ], + ]) + ); + + const res = await authedPost('/admin/runs', { kind: 'decider' }); + expect(res.status).toBe(200); + const body = (await res.json()) as { enqueuedModels: number; skippedModels: string[] }; + expect(body.skippedModels).toEqual([]); + expect(body.enqueuedModels).toBe(1); + }); + + it('carries a prior summary when the exact pair matches (legacy effort key)', async () => { + vi.mocked(getConfigRows).mockResolvedValue({ + ...TEST_CONFIG_ROWS, + config: { ...TEST_CONFIG_ROWS.config, benchmark_user_id: 'user-123' }, + // vendor/b has no prior → stays enqueued so we do not hit the all-carried + // finalize path (which needs a real D1 client in unit tests). + deciderModels: [ + { model: 'vendor/a', reasoning_effort: 'high' }, + { model: 'vendor/b', reasoning_effort: null }, + ], + }); + // Legacy rows store variant from reasoning_effort; exact pair high matches. + vi.mocked(getLatestSummariesByModel).mockResolvedValue( + new Map([ + [ + exactPairKey('vendor/a', 'high'), + { + engineIdentity: computeEngineIdentity('decider'), + repetitions: 1, + variant: 'high', + reasoningEffort: 'high', + summaries: [makeSummary('vendor/a')], + }, + ], + ]) + ); + + const res = await authedPost('/admin/runs', { kind: 'decider' }); + expect(res.status).toBe(200); + const body = (await res.json()) as { enqueuedModels: number; skippedModels: string[] }; + expect(body.skippedModels).toEqual(['vendor/a']); + expect(body.enqueuedModels).toBe(1); + }); + it('seeds sharded decider lanes bounded by the container cap', async () => { // Later chunks are chained by processJob within each shard lane. Start // seeds as many lanes as fit under the 100-container cap so the benchmark @@ -684,3 +773,114 @@ describe('GET /admin/classifier-winner', () => { await expect(res.json()).resolves.toEqual({ winner }); }); }); + +// --------------------------------------------------------------------------- +// POST /admin/profiles/register + /admin/profiles/status +// --------------------------------------------------------------------------- + +describe('POST /admin/profiles/register', () => { + it('returns 400 when benchmark config is not set', async () => { + const res = await authedPost('/admin/profiles/register', { + ownerType: 'user', + ownerId: 'u1', + entries: [{ model: 'a/b', variant: null }], + }); + expect(res.status).toBe(400); + await expect(res.json()).resolves.toMatchObject({ + error: expect.stringContaining('benchmark config not set'), + }); + expect(registerProfiles).not.toHaveBeenCalled(); + }); + + it('returns 400 for malformed body', async () => { + // Validation fails before the handler reads config — do not queue a + // mockResolvedValueOnce here or it will leak into later tests. + const res = await authedPost('/admin/profiles/register', { + ownerType: 'user', + ownerId: 'u1', + entries: [], + }); + expect(res.status).toBe(400); + expect(registerProfiles).not.toHaveBeenCalled(); + }); + + it('returns 200 with statuses on successful admission', async () => { + vi.mocked(getConfigRows).mockResolvedValueOnce(TEST_CONFIG_ROWS); + vi.mocked(registerProfiles).mockResolvedValueOnce({ + statuses: [ + { entry: { model: 'a/b', variant: null }, status: 'pending', failureReason: null }, + ], + }); + + const res = await authedPost('/admin/profiles/register', { + ownerType: 'user', + ownerId: 'u1', + entries: [{ model: 'a/b', variant: null }], + retryEntries: [{ model: 'a/b', variant: null }], + }); + expect(res.status).toBe(200); + await expect(res.json()).resolves.toEqual({ + statuses: [ + { entry: { model: 'a/b', variant: null }, status: 'pending', failureReason: null }, + ], + }); + expect(registerProfiles).toHaveBeenCalledWith( + env.BENCH_DB, + expect.objectContaining({ deciderRepetitions: 1 }), + expect.objectContaining({ + ownerType: 'user', + ownerId: 'u1', + entries: [{ model: 'a/b', variant: null }], + retryEntries: [{ model: 'a/b', variant: null }], + }) + ); + }); + + it('returns 429 with retryAt when quota is exceeded', async () => { + vi.mocked(getConfigRows).mockResolvedValueOnce(TEST_CONFIG_ROWS); + vi.mocked(registerProfiles).mockRejectedValueOnce( + new ProfileQuotaExceededError({ + error: + 'Profile benchmark request limit reached. New benchmarks can be requested after 2026-07-29T12:00:00.000Z.', + retryAt: '2026-07-29T12:00:00.000Z', + }) + ); + + const res = await authedPost('/admin/profiles/register', { + ownerType: 'org', + ownerId: 'o1', + entries: [{ model: 'a/b', variant: 'xhigh' }], + }); + expect(res.status).toBe(429); + await expect(res.json()).resolves.toEqual({ + error: + 'Profile benchmark request limit reached. New benchmarks can be requested after 2026-07-29T12:00:00.000Z.', + retryAt: '2026-07-29T12:00:00.000Z', + }); + }); +}); + +describe('POST /admin/profiles/status', () => { + it('returns 400 when benchmark config is not set', async () => { + const res = await authedPost('/admin/profiles/status', { + entries: [{ model: 'a/b', variant: null }], + }); + expect(res.status).toBe(400); + expect(lookupProfileStatuses).not.toHaveBeenCalled(); + }); + + it('returns current statuses', async () => { + vi.mocked(getConfigRows).mockResolvedValueOnce(TEST_CONFIG_ROWS); + vi.mocked(lookupProfileStatuses).mockResolvedValueOnce({ + statuses: [{ entry: { model: 'a/b', variant: null }, status: 'ready', failureReason: null }], + }); + + const res = await authedPost('/admin/profiles/status', { + entries: [{ model: 'a/b', variant: null }], + }); + expect(res.status).toBe(200); + await expect(res.json()).resolves.toEqual({ + statuses: [{ entry: { model: 'a/b', variant: null }, status: 'ready', failureReason: null }], + }); + }); +}); diff --git a/services/auto-routing-benchmark/src/admin.ts b/services/auto-routing-benchmark/src/admin.ts index 15fa9a4d26..2dc76643f1 100644 --- a/services/auto-routing-benchmark/src/admin.ts +++ b/services/auto-routing-benchmark/src/admin.ts @@ -1,6 +1,9 @@ import * as z from 'zod'; import { BenchmarkConfigSchema, + BenchmarkProfileStatusesRequestSchema, + CustomRoutingTableRequestSchema, + RegisterBenchmarkProfilesRequestSchema, resolveBenchmarkIdentity, StartBenchmarkRunRequestSchema, type BenchmarkRun, @@ -9,6 +12,14 @@ import { zodJsonValidator } from '@kilocode/worker-utils'; import type { Hono } from 'hono'; import { getBenchmarkConfig, saveBenchmarkConfig } from './config'; import { debugRunCli } from './cli-runner'; +import { assembleCustomRoutingTable } from './custom-routing-table'; +import { + lookupProfileStatuses, + ProfileConfigMissingError, + ProfileQuotaExceededError, + ProfileValidationError, + registerProfiles, +} from './profiles'; import { BenchmarkRunConfigError, fetchBenchmarkUserToken, @@ -112,4 +123,102 @@ export function registerAdminRoutes(app: Hono): void { return c.json(result); } ); + + // Atomically admit missing/stale/retried Benchmark profiles for an owner. + app.post( + '/admin/profiles/register', + zodJsonValidator(RegisterBenchmarkProfilesRequestSchema, { + errorMessage: 'Invalid profile register request', + }), + async c => { + const body = c.req.valid('json'); + const config = await getBenchmarkConfig(c.env.BENCH_DB); + if (!config) { + return c.json( + { + error: + 'benchmark config not set: save it in the admin panel before registering profiles', + }, + 400 + ); + } + try { + const result = await registerProfiles(c.env.BENCH_DB, config, { + ownerType: body.ownerType, + ownerId: body.ownerId, + entries: body.entries, + retryEntries: body.retryEntries, + }); + return c.json(result); + } catch (error) { + if (error instanceof ProfileQuotaExceededError) { + return c.json(error.quota, 429); + } + if (error instanceof ProfileValidationError || error instanceof ProfileConfigMissingError) { + return c.json({ error: error.message }, 400); + } + throw error; + } + } + ); + + // Current per-entry Benchmark-profile statuses (may free-admit stale rows). + app.post( + '/admin/profiles/status', + zodJsonValidator(BenchmarkProfileStatusesRequestSchema, { + errorMessage: 'Invalid profile status request', + }), + async c => { + const body = c.req.valid('json'); + const config = await getBenchmarkConfig(c.env.BENCH_DB); + if (!config) { + return c.json( + { + error: + 'benchmark config not set: save it in the admin panel before looking up profiles', + }, + 400 + ); + } + try { + return c.json( + await lookupProfileStatuses(c.env.BENCH_DB, config, { entries: body.entries }) + ); + } catch (error) { + if (error instanceof ProfileValidationError || error instanceof ProfileConfigMissingError) { + return c.json({ error: error.message }, 400); + } + throw error; + } + } + ); + + // Sparse custom routing table for ready/current pool entries only. + app.post( + '/admin/custom-routing-table', + zodJsonValidator(CustomRoutingTableRequestSchema, { + errorMessage: 'Invalid custom routing table request', + }), + async c => { + const body = c.req.valid('json'); + const config = await getBenchmarkConfig(c.env.BENCH_DB); + if (!config) { + return c.json( + { + error: + 'benchmark config not set: save it in the admin panel before assembling a custom routing table', + }, + 400 + ); + } + try { + return c.json(await assembleCustomRoutingTable(c.env.BENCH_DB, config, body.entries)); + } catch (error) { + if (error instanceof ProfileValidationError || error instanceof ProfileConfigMissingError) { + return c.json({ error: error.message }, 400); + } + throw error; + } + } + ); } diff --git a/services/auto-routing-benchmark/src/auto-decider-sync.test.ts b/services/auto-routing-benchmark/src/auto-decider-sync.test.ts index e13904441b..aebcb47b9a 100644 --- a/services/auto-routing-benchmark/src/auto-decider-sync.test.ts +++ b/services/auto-routing-benchmark/src/auto-decider-sync.test.ts @@ -12,6 +12,11 @@ vi.mock('./db', async importOriginal => { getLatestSummariesByModel: vi.fn(), insertRun: vi.fn(), markStaleRunsFailed: vi.fn(), + listStaleRunningDeciderRunIds: vi.fn(), + listPendingCurrentProfiles: vi.fn(), + markProfilesFailedForRun: vi.fn(), + markProfilesRunningForRun: vi.fn(), + markProfilesReadyForRun: vi.fn(), }; }); @@ -20,6 +25,8 @@ import { getLatestSummariesByModel, getRunningRun, insertRun, + listPendingCurrentProfiles, + listStaleRunningDeciderRunIds, markStaleRunsFailed, replaceAutoDeciderModels, } from './db'; @@ -85,6 +92,8 @@ describe('syncAutoDeciderModels', () => { }); vi.mocked(replaceAutoDeciderModels).mockResolvedValue(undefined); vi.mocked(markStaleRunsFailed).mockResolvedValue(undefined); + vi.mocked(listStaleRunningDeciderRunIds).mockResolvedValue([]); + vi.mocked(listPendingCurrentProfiles).mockResolvedValue([]); vi.mocked(getRunningRun).mockResolvedValue(undefined); vi.mocked(getLatestSummariesByModel).mockResolvedValue(new Map()); vi.mocked(insertRun).mockResolvedValue(undefined); @@ -129,6 +138,7 @@ describe('syncAutoDeciderModels', () => { repetitions: 1, classifier_max_p95_latency_ms: null, engine_identity: 'v1:test', + purpose: 'platform', }); const result = await syncAutoDeciderModels(env, { fetchImpl }); @@ -143,4 +153,94 @@ describe('syncAutoDeciderModels', () => { }); expect(insertRun).not.toHaveBeenCalled(); }); + + it('drains stranded pending profiles when the slot is free and no model change', async () => { + // No effective model change after sync. + fetchImpl.mockResolvedValue( + new Response( + JSON.stringify({ + candidates: [{ id: 'auto/existing', avgAttemptCostUsd: 18 }], + }), + { status: 200, headers: { 'content-type': 'application/json' } } + ) + ); + vi.mocked(getConfigRows).mockResolvedValue({ + config, + classifierModels: ['classifier/model'], + deciderModels: [{ model: 'manual/model', reasoning_effort: null }], + autoDeciderModels: [ + { + model: 'auto/existing', + reasoning_effort: 'high', + avg_attempt_cost_usd: 18, + synced_at: '2026-06-01T00:00:00.000Z', + }, + ], + excludedAutoDeciderModels: [], + }); + vi.mocked(listPendingCurrentProfiles).mockResolvedValue([ + { model: 'stranded/m', variant: 'high', requested_at: '2026-06-01T00:00:00.000Z' }, + ]); + + const result = await syncAutoDeciderModels(env, { fetchImpl }); + + expect(result.startedRun).toBe(false); + expect(result.profileDrainRunId).toMatch(/^profile-/); + expect(insertRun).toHaveBeenCalledOnce(); + expect(vi.mocked(insertRun).mock.calls[0][1].purpose).toBe('profile'); + }); + + it('platform start claims free slot over pending profiles when models changed', async () => { + // Free slot + pending profiles + changed models → PLATFORM run starts; + // no profile drain this cycle (terminal transition drains later). + vi.mocked(listPendingCurrentProfiles).mockResolvedValue([ + { model: 'pending/m', variant: 'high', requested_at: '2026-06-01T00:00:00.000Z' }, + ]); + + const result = await syncAutoDeciderModels(env, { fetchImpl }); + + expect(result.startedRun).toBe(true); + expect(result.runId).toBeTruthy(); + expect(result.profileDrainRunId).toBeNull(); + expect(insertRun).toHaveBeenCalledOnce(); + expect(vi.mocked(insertRun).mock.calls[0][1].purpose).toBe('platform'); + // Drain must not have claimed pending entries this cycle. + expect(listPendingCurrentProfiles).not.toHaveBeenCalled(); + }); + + it('occupied slot skips without draining pending profiles', async () => { + vi.mocked(getRunningRun).mockResolvedValue({ + id: 'decider-active', + kind: 'decider', + status: 'running', + started_at: '2026-06-01T00:00:00.000Z', + completed_at: null, + error: null, + min_accuracy: 0.7, + switch_cost_factor: 3, + best_accuracy_switch_threshold: 0.05, + max_concurrency: 100, + benchmark_user_id: 'user-123', + benchmark_org_id: null, + repetitions: 1, + classifier_max_p95_latency_ms: null, + engine_identity: 'v1:test', + purpose: 'platform', + }); + vi.mocked(listPendingCurrentProfiles).mockResolvedValue([ + { model: 'pending/m', variant: 'high', requested_at: '2026-06-01T00:00:00.000Z' }, + ]); + + const result = await syncAutoDeciderModels(env, { fetchImpl }); + + expect(result).toMatchObject({ + startedRun: false, + skippedReason: 'active-run', + activeRunId: 'decider-active', + profileDrainRunId: null, + }); + expect(insertRun).not.toHaveBeenCalled(); + // Slot occupied → leave pending rows untouched (no drain attempt). + expect(listPendingCurrentProfiles).not.toHaveBeenCalled(); + }); }); diff --git a/services/auto-routing-benchmark/src/auto-decider-sync.ts b/services/auto-routing-benchmark/src/auto-decider-sync.ts index 972b6284f5..431f0ff986 100644 --- a/services/auto-routing-benchmark/src/auto-decider-sync.ts +++ b/services/auto-routing-benchmark/src/auto-decider-sync.ts @@ -6,7 +6,7 @@ import { } from '@kilocode/auto-routing-contracts'; import { getBenchmarkConfig, mapConfigRows } from './config'; import { getConfigRows, replaceAutoDeciderModels, type ConfigAutoDeciderModelRow } from './db'; -import { RunAlreadyActiveError, startRun } from './run'; +import { drainPendingProfileBatch, RunAlreadyActiveError, startRun, sweepStaleRuns } from './run'; type SyncOptions = { fetchImpl?: typeof fetch; @@ -20,6 +20,8 @@ export type AutoDeciderSyncResult = { runId: string | null; skippedReason?: 'active-run'; activeRunId?: string; + /** Profile batch started when no platform start claimed the free slot. */ + profileDrainRunId?: string | null; }; function modelKey(model: BenchmarkDeciderModel): string { @@ -64,12 +66,34 @@ async function fetchAutoDeciderCandidates( return parsed.data.candidates; } +async function tryDrainPendingProfiles(env: Env): Promise { + try { + const drained = await drainPendingProfileBatch(env); + return drained?.runId ?? null; + } catch (error) { + console.warn( + JSON.stringify({ + event: 'benchmark_profile_drain_error', + afterAutoDeciderSync: true, + error: error instanceof Error ? error.message : String(error), + }) + ); + return null; + } +} + export async function syncAutoDeciderModels( env: Env, options: SyncOptions = {} ): Promise { const fetchImpl = options.fetchImpl ?? fetch; const syncedAt = (options.now ?? new Date()).toISOString(); + + // Stale-run cleanup only — not a slot claim. Pending-profile drain is deferred + // until after the platform-start decision so profile work never preempts a + // free slot that scheduled platform start needs (plan change 6 / res #5). + await sweepStaleRuns(env.BENCH_DB); + const beforeRows = await getConfigRows(env.BENCH_DB); const beforeConfig = mapConfigRows( beforeRows.config, @@ -110,16 +134,35 @@ export async function syncAutoDeciderModels( ); const diff = diffModels(beforeConfig?.deciderModels ?? [], afterConfig?.deciderModels ?? []); const changed = diff.added.length > 0 || diff.removed.length > 0; + const hasConfig = Boolean(await getBenchmarkConfig(env.BENCH_DB)); + const platformStartNeeded = changed && hasConfig; - if (!changed || !(await getBenchmarkConfig(env.BENCH_DB))) { - return { addedModels: diff.added, removedModels: diff.removed, startedRun: false, runId: null }; + if (!platformStartNeeded) { + // No platform start this cycle → drain now so stranded pending work recovers. + const profileDrainRunId = await tryDrainPendingProfiles(env); + return { + addedModels: diff.added, + removedModels: diff.removed, + startedRun: false, + runId: null, + profileDrainRunId, + }; } + // Platform scheduled start takes the free slot first; profile work never + // preempts. Occupied slot → log/skip that cycle (never fail job); leave + // pending rows untouched (no drain while the slot is held). let run: Awaited>; try { run = await startRun(env, 'decider'); } catch (error) { if (error instanceof RunAlreadyActiveError) { + console.log( + JSON.stringify({ + event: 'auto_decider_sync_skipped_active_run', + activeRunId: error.activeRunId, + }) + ); return { addedModels: diff.added, removedModels: diff.removed, @@ -127,14 +170,19 @@ export async function syncAutoDeciderModels( runId: null, skippedReason: 'active-run', activeRunId: error.activeRunId, + profileDrainRunId: null, }; } throw error; } + + // Platform run claimed the slot. Do NOT drain this cycle — the run's terminal + // transition (or a later cron with no platform start) drains pending work. return { addedModels: diff.added, removedModels: diff.removed, startedRun: true, runId: run.runId, + profileDrainRunId: null, }; } diff --git a/services/auto-routing-benchmark/src/cli-runner.test.ts b/services/auto-routing-benchmark/src/cli-runner.test.ts index a60b7ef6ee..010be8b16c 100644 --- a/services/auto-routing-benchmark/src/cli-runner.test.ts +++ b/services/auto-routing-benchmark/src/cli-runner.test.ts @@ -50,6 +50,54 @@ describe('runDeciderCaseViaCli', () => { orgId: 'org-123', }); }); + + it('passes canonical variant as the container body variant field', async () => { + const fetch = vi.fn(async (_request: Request) => + Response.json({ + exitCode: 0, + durationMs: 10, + stdoutLines: [], + stderrTail: '', + }) + ); + const { env } = createEnv(fetch); + + await runDeciderCaseViaCli(env, { + instanceName: 'run:model:0', + model: 'vendor/model', + benchCase, + kiloToken: 'kilo-user-token', + kiloApiUrl: 'http://host.docker.internal:3000', + variant: 'xhigh', + }); + + const request = fetch.mock.calls[0]?.[0]; + await expect(readJsonBody(request)).resolves.toMatchObject({ variant: 'xhigh' }); + }); + + it('accepts legacy reasoningEffort alias for the same body field', async () => { + const fetch = vi.fn(async (_request: Request) => + Response.json({ + exitCode: 0, + durationMs: 10, + stdoutLines: [], + stderrTail: '', + }) + ); + const { env } = createEnv(fetch); + + await runDeciderCaseViaCli(env, { + instanceName: 'run:model:0', + model: 'vendor/model', + benchCase, + kiloToken: 'kilo-user-token', + kiloApiUrl: 'http://host.docker.internal:3000', + reasoningEffort: 'high', + }); + + const request = fetch.mock.calls[0]?.[0]; + await expect(readJsonBody(request)).resolves.toMatchObject({ variant: 'high' }); + }); }); describe('warmUpCliContainer', () => { diff --git a/services/auto-routing-benchmark/src/cli-runner.ts b/services/auto-routing-benchmark/src/cli-runner.ts index 18ad73b9be..798e2f44cf 100644 --- a/services/auto-routing-benchmark/src/cli-runner.ts +++ b/services/auto-routing-benchmark/src/cli-runner.ts @@ -58,10 +58,18 @@ export async function runDeciderCaseViaCli( kiloToken: string; kiloApiUrl: string; orgId?: string | null; + /** + * Canonical catalog variant passed as the container body's `variant` + * field (unchanged body shape). Accepts the legacy `reasoningEffort` + * alias so existing call sites keep working. + */ + variant?: string | null; + /** @deprecated Prefer `variant`. Same value; kept for call-site compatibility. */ reasoningEffort?: string | null; } ): Promise { - const { instanceName, model, benchCase, kiloToken, kiloApiUrl, orgId, reasoningEffort } = params; + const { instanceName, model, benchCase, kiloToken, kiloApiUrl, orgId } = params; + const variant = params.variant !== undefined ? params.variant : (params.reasoningEffort ?? null); const stub = env.BENCH_RUNNER.get(env.BENCH_RUNNER.idFromName(instanceName)); const prompt = `${benchCase.systemPrompt}\n\n${benchCase.userPrompt}${FINAL_ANSWER_SUFFIX}`; @@ -77,7 +85,7 @@ export async function runDeciderCaseViaCli( kiloApiUrl, orgId: orgId ?? null, timeoutMs: DECIDER_CLI_TIMEOUT_MS, - variant: reasoningEffort ?? null, + variant: variant ?? null, }), }) ); diff --git a/services/auto-routing-benchmark/src/custom-routing-table.test.ts b/services/auto-routing-benchmark/src/custom-routing-table.test.ts new file mode 100644 index 0000000000..69efbc65b3 --- /dev/null +++ b/services/auto-routing-benchmark/src/custom-routing-table.test.ts @@ -0,0 +1,230 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type { TaxonomyRouteKey } from '@kilocode/auto-routing-contracts'; +import { TAXONOMY_ROUTE_KEYS } from '@kilocode/auto-routing-contracts'; +import type * as DbModule from './db'; +import type { BenchmarkModelSummaryWithRun } from './db'; +import { buildCustomRoutingTable, computeCustomRoutingTableVersion } from './routing-table-builder'; + +vi.mock('./db', async importOriginal => { + const actual = await importOriginal(); + return { + ...actual, + listReadyCurrentProfilesForEntries: vi.fn(), + getSummariesForRuns: vi.fn(), + }; +}); + +import { getSummariesForRuns, listReadyCurrentProfilesForEntries } from './db'; +import { assembleCustomRoutingTable } from './custom-routing-table'; +import { computeEngineIdentity } from './run'; + +function summary( + model: string, + variant: string | null, + routeKey: TaxonomyRouteKey, + accuracy = 0.9, + avgCostUsd: number | null = 0.001, + runId = 'run-default' +): BenchmarkModelSummaryWithRun { + return { + model, + variant, + routeKey, + accuracy, + avgCostUsd, + avgLatencyMs: 100, + p50LatencyMs: 90, + p95LatencyMs: 120, + cases: 5, + errors: 0, + timeouts: 0, + runId, + }; +} + +describe('buildCustomRoutingTable', () => { + const base = { + generatedAt: '2026-06-01T00:00:00.000Z', + minAccuracy: 0.7, + switchCostFactor: 3, + bestAccuracySwitchThreshold: 0.05, + }; + + it('assembles two variants of one model as distinct candidates with variant', () => { + const readyEntries = [ + { entry: { model: 'vendor/m', variant: 'xhigh' }, runId: 'run-x' }, + { entry: { model: 'vendor/m', variant: 'max' }, runId: 'run-m' }, + ]; + const routeKey = 'implementation/code_generation' as TaxonomyRouteKey; + const summaries = [ + summary('vendor/m', 'xhigh', routeKey, 0.9, 0.002, 'run-x'), + summary('vendor/m', 'max', routeKey, 0.85, 0.001, 'run-m'), + ]; + const table = buildCustomRoutingTable({ ...base, readyEntries, summaries }); + expect(table).not.toBeNull(); + const cands = table!.routes[routeKey]!; + expect(cands).toHaveLength(2); + expect(cands.map(c => c.variant).sort()).toEqual(['max', 'xhigh']); + expect(cands.every(c => c.reasoningEffort === undefined)).toBe(true); + }); + + it('omits routes with no graded candidates', () => { + const readyEntries = [{ entry: { model: 'vendor/m', variant: null }, runId: 'run-1' }]; + const onlyOneRoute = 'implementation/code_generation' as TaxonomyRouteKey; + const summaries = [summary('vendor/m', null, onlyOneRoute, 0.9, 0.001, 'run-1')]; + const table = buildCustomRoutingTable({ ...base, readyEntries, summaries }); + expect(table).not.toBeNull(); + expect(Object.keys(table!.routes)).toEqual([onlyOneRoute]); + for (const key of TAXONOMY_ROUTE_KEYS) { + if (key !== onlyOneRoute) { + expect(table!.routes[key]).toBeUndefined(); + } + } + }); + + it('unready entries contribute nothing; empty → null', () => { + expect( + buildCustomRoutingTable({ + ...base, + readyEntries: [], + summaries: [summary('vendor/m', null, 'implementation/code_generation')], + }) + ).toBeNull(); + }); + + it('ignores summaries for pairs not in readyEntries (stale/unready)', () => { + const readyEntries = [{ entry: { model: 'vendor/ready', variant: 'high' }, runId: 'run-r' }]; + const summaries = [ + summary('vendor/ready', 'high', 'implementation/code_generation', 0.9, 0.001, 'run-r'), + summary('vendor/stale', 'high', 'implementation/code_generation', 0.9, 0.001, 'run-r'), + ]; + const table = buildCustomRoutingTable({ ...base, readyEntries, summaries }); + const cands = table!.routes['implementation/code_generation']!; + expect(cands).toHaveLength(1); + expect(cands[0].model).toBe('vendor/ready'); + }); + + it('binds each ready entry to its provenance run only (no cross-run leakage)', () => { + // R1 measured entry A and also graded B's pair; B is ready on R2 with + // fresher metrics. Assembly for B must use R2 only — never R1's stale row. + const routeKey = 'implementation/code_generation' as TaxonomyRouteKey; + const readyEntries = [ + { entry: { model: 'vendor/a', variant: 'high' }, runId: 'run-r1' }, + { entry: { model: 'vendor/b', variant: 'max' }, runId: 'run-r2' }, + ]; + const summaries = [ + summary('vendor/a', 'high', routeKey, 0.9, 0.002, 'run-r1'), + // Stale B metrics from R1 (higher accuracy would win if not filtered) + summary('vendor/b', 'max', routeKey, 0.99, 0.0001, 'run-r1'), + // True provenance metrics for B from R2 + summary('vendor/b', 'max', routeKey, 0.8, 0.003, 'run-r2'), + ]; + const table = buildCustomRoutingTable({ ...base, readyEntries, summaries }); + expect(table).not.toBeNull(); + const cands = table!.routes[routeKey]!; + expect(cands).toHaveLength(2); + const b = cands.find(c => c.model === 'vendor/b'); + expect(b).toMatchObject({ + model: 'vendor/b', + variant: 'max', + accuracy: 0.8, + avgCostUsd: 0.003, + }); + // Deterministic version shape unchanged + expect(table!.version).toMatch(/^custom-[0-9a-f]{8}$/); + expect(table!.version).toBe( + computeCustomRoutingTableVersion([ + { runId: 'run-r1', model: 'vendor/a', variant: 'high' }, + { runId: 'run-r2', model: 'vendor/b', variant: 'max' }, + ]) + ); + }); + + it('excludes candidates with no cost signal or zero cases', () => { + const readyEntries = [ + { entry: { model: 'a', variant: null }, runId: 'r1' }, + { entry: { model: 'b', variant: null }, runId: 'r2' }, + ]; + const routeKey = 'implementation/code_generation' as TaxonomyRouteKey; + const summaries = [ + summary('a', null, routeKey, 0.9, null, 'r1'), + { ...summary('b', null, routeKey, 0.8, 0.001, 'r2'), cases: 0 }, + ]; + expect(buildCustomRoutingTable({ ...base, readyEntries, summaries })).toBeNull(); + }); + + it('produces a deterministic version for identical inputs', () => { + const readyEntries = [ + { entry: { model: 'm', variant: 'a' }, runId: 'run-1' }, + { entry: { model: 'm', variant: 'b' }, runId: 'run-2' }, + ]; + const summaries = [ + summary('m', 'a', 'implementation/code_generation', 0.9, 0.001, 'run-1'), + summary('m', 'b', 'implementation/code_generation', 0.9, 0.001, 'run-2'), + ]; + const t1 = buildCustomRoutingTable({ ...base, readyEntries, summaries }); + const t2 = buildCustomRoutingTable({ + ...base, + generatedAt: '2099-01-01T00:00:00.000Z', + readyEntries: [...readyEntries].reverse(), + summaries: [...summaries].reverse(), + }); + expect(t1!.version).toBe(t2!.version); + expect(t1!.version).toMatch(/^custom-[0-9a-f]{8}$/); + }); + + it('computeCustomRoutingTableVersion is order-independent', () => { + const a = [ + { runId: 'r1', model: 'm', variant: 'x' as string | null }, + { runId: 'r2', model: 'n', variant: null }, + ]; + expect(computeCustomRoutingTableVersion(a)).toBe( + computeCustomRoutingTableVersion([...a].reverse()) + ); + }); +}); + +describe('assembleCustomRoutingTable', () => { + const config = { + deciderRepetitions: 1, + minAccuracy: 0.7, + switchCostFactor: 3, + bestAccuracySwitchThreshold: 0.05, + }; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('returns table null when no ready current profiles', async () => { + vi.mocked(listReadyCurrentProfilesForEntries).mockResolvedValue([]); + const result = await assembleCustomRoutingTable({} as D1Database, config, [ + { model: 'm', variant: 'x' }, + ]); + expect(result).toEqual({ table: null }); + expect(getSummariesForRuns).not.toHaveBeenCalled(); + }); + + it('loads provenance summaries and assembles with current engine filter', async () => { + const engine = computeEngineIdentity('decider'); + vi.mocked(listReadyCurrentProfilesForEntries).mockResolvedValue([ + { model: 'vendor/m', variant: 'xhigh', run_id: 'run-x' }, + { model: 'vendor/m', variant: 'max', run_id: 'run-m' }, + ]); + vi.mocked(getSummariesForRuns).mockResolvedValue([ + summary('vendor/m', 'xhigh', 'implementation/code_generation', 0.9, 0.002, 'run-x'), + summary('vendor/m', 'max', 'implementation/code_generation', 0.8, 0.001, 'run-m'), + ]); + const result = await assembleCustomRoutingTable({} as D1Database, config, [ + { model: 'vendor/m', variant: 'xhigh' }, + { model: 'vendor/m', variant: 'max' }, + ]); + expect(listReadyCurrentProfilesForEntries).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ engineIdentity: engine, repetitions: 1 }), + expect.any(Array) + ); + expect(result.table).not.toBeNull(); + expect(result.table!.routes['implementation/code_generation']).toHaveLength(2); + }); +}); diff --git a/services/auto-routing-benchmark/src/custom-routing-table.ts b/services/auto-routing-benchmark/src/custom-routing-table.ts new file mode 100644 index 0000000000..21569a398a --- /dev/null +++ b/services/auto-routing-benchmark/src/custom-routing-table.ts @@ -0,0 +1,77 @@ +import type { + BenchmarkConfig, + CustomRoutingTable, + CustomRoutingTableResponse, + PoolEntry, +} from '@kilocode/auto-routing-contracts'; +import { MAX_POOL_ENTRIES, poolEntryKey } from '@kilocode/auto-routing-contracts'; +import { getSummariesForRuns, listReadyCurrentProfilesForEntries } from './db'; +import { variantFromStorage } from './reasoning-effort'; +import { buildCustomRoutingTable } from './routing-table-builder'; +import { currentProfileContextFromConfig, ProfileValidationError } from './profiles'; + +function assertUniqueEntries(entries: readonly PoolEntry[]): void { + if (entries.length < 1 || entries.length > MAX_POOL_ENTRIES) { + throw new ProfileValidationError( + `entries must contain between 1 and ${MAX_POOL_ENTRIES} unique pool entries` + ); + } + const seen = new Set(); + for (const entry of entries) { + const key = poolEntryKey(entry); + if (seen.has(key)) { + throw new ProfileValidationError(`Duplicate pool entry: ${key}`); + } + seen.add(key); + } +} + +/** + * Assemble a sparse custom routing table for the requested exact Pool entries + * from ready+current Benchmark profiles and their provenance-run summaries. + * Returns `{ table: null }` when no requested entry is ready/current. + */ +export async function assembleCustomRoutingTable( + db: D1Database, + config: Pick< + BenchmarkConfig, + 'deciderRepetitions' | 'minAccuracy' | 'switchCostFactor' | 'bestAccuracySwitchThreshold' + >, + entries: readonly PoolEntry[], + options: { now?: Date } = {} +): Promise { + assertUniqueEntries(entries); + const current = currentProfileContextFromConfig(config); + const readyRows = await listReadyCurrentProfilesForEntries(db, current, entries); + + const readyEntries = readyRows + .filter( + (row): row is typeof row & { run_id: string } => row.run_id != null && row.run_id !== '' + ) + .map(row => ({ + entry: { + model: row.model, + variant: variantFromStorage(row.variant), + } satisfies PoolEntry, + runId: row.run_id, + })); + + if (readyEntries.length === 0) { + return { table: null }; + } + + const runIds = [...new Set(readyEntries.map(r => r.runId))]; + const summaries = await getSummariesForRuns(db, runIds); + const generatedAt = (options.now ?? new Date()).toISOString(); + + const table: CustomRoutingTable | null = buildCustomRoutingTable({ + generatedAt, + minAccuracy: config.minAccuracy, + switchCostFactor: config.switchCostFactor, + bestAccuracySwitchThreshold: config.bestAccuracySwitchThreshold, + readyEntries, + summaries, + }); + + return { table }; +} diff --git a/services/auto-routing-benchmark/src/db-replace-summaries.test.ts b/services/auto-routing-benchmark/src/db-replace-summaries.test.ts index 56e1219260..a33d4d7425 100644 --- a/services/auto-routing-benchmark/src/db-replace-summaries.test.ts +++ b/services/auto-routing-benchmark/src/db-replace-summaries.test.ts @@ -52,8 +52,9 @@ describe('replaceModelSummaries', () => { await replaceModelSummaries({} as D1Database, 'run-1', summaries); expect(mocks.insertValues).toHaveBeenCalledTimes(2); + // 13 bind values per row × 7 rows = 91 < 100 D1 variable ceiling. expect(mocks.insertValues.mock.calls.map(([rows]) => (rows as unknown[]).length)).toEqual([ - 8, 2, + 7, 3, ]); expect(mocks.batch).toHaveBeenCalledTimes(1); expect(mocks.batch.mock.calls[0]?.[0]).toHaveLength(3); @@ -96,7 +97,7 @@ describe('insertRun', () => { .map(([rows]) => rows) .filter(Array.isArray) .map(rows => rows.length); - expect(carriedInsertSizes).toEqual([8, 2]); + expect(carriedInsertSizes).toEqual([7, 3]); expect(mocks.batch).toHaveBeenCalledTimes(1); expect(mocks.batch.mock.calls[0]?.[0]).toHaveLength(3); }); diff --git a/services/auto-routing-benchmark/src/db-save-routing-table.test.ts b/services/auto-routing-benchmark/src/db-save-routing-table.test.ts index d5fd4b7f53..82a9c50f24 100644 --- a/services/auto-routing-benchmark/src/db-save-routing-table.test.ts +++ b/services/auto-routing-benchmark/src/db-save-routing-table.test.ts @@ -1,28 +1,31 @@ -import { describe, expect, it, vi } from 'vitest'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; import type { RankedCandidate, RoutingTable } from '@kilocode/auto-routing-contracts'; -const mockState = vi.hoisted(() => ({ - batchCalls: [] as Array>, -})); +const mocks = vi.hoisted(() => { + const batch = vi.fn(async (_stmts: unknown[]) => []); + const where = vi.fn(() => ({ kind: 'delete' })); + const deleteFrom = vi.fn(() => ({ where })); + const onConflictDoUpdate = vi.fn((args: unknown) => ({ kind: 'upsert', args })); + const insertValues = vi.fn((values: unknown) => ({ + kind: 'insert', + values, + onConflictDoUpdate, + })); + const insertInto = vi.fn(() => ({ values: insertValues })); + + return { batch, deleteFrom, insertInto, insertValues, onConflictDoUpdate, where }; +}); vi.mock('drizzle-orm/d1', () => ({ drizzle: vi.fn(() => ({ - delete: vi.fn(() => ({ - where: vi.fn(() => ({ kind: 'delete' })), - })), - insert: vi.fn(() => ({ - values: vi.fn((values: unknown) => ({ - kind: 'insert', - values, - onConflictDoUpdate: vi.fn(() => ({ kind: 'upsert', values })), - })), - })), - batch: vi.fn(async (stmts: Array<{ kind: string; values?: unknown }>) => { - mockState.batchCalls.push(stmts); - }), + batch: mocks.batch, + delete: mocks.deleteFrom, + insert: mocks.insertInto, })), })); +import { saveRoutingTable } from './db'; + const candidate = (model: string): RankedCandidate => ({ model, accuracy: 0.9, @@ -32,9 +35,16 @@ const candidate = (model: string): RankedCandidate => ({ }); describe('saveRoutingTable', () => { - it('chunks routing candidate inserts to stay under D1 variable limits', async () => { - const { saveRoutingTable } = await import('./db'); + beforeEach(() => { + mocks.batch.mockClear(); + mocks.deleteFrom.mockClear(); + mocks.insertInto.mockClear(); + mocks.insertValues.mockClear(); + mocks.onConflictDoUpdate.mockClear(); + mocks.where.mockClear(); + }); + it('chunks routing candidate inserts to stay under D1 variable limits', async () => { const table: RoutingTable = { version: 'run-large-routing-table', generatedAt: '2026-06-16T18:00:00.000Z', @@ -53,7 +63,8 @@ describe('saveRoutingTable', () => { await saveRoutingTable({} as D1Database, table, '2026-06-16T18:01:00.000Z'); - const [batch] = mockState.batchCalls; + expect(mocks.batch).toHaveBeenCalledTimes(1); + const batch = mocks.batch.mock.calls[0]?.[0] as Array<{ kind: string; values?: unknown }>; expect(batch).toBeDefined(); const candidateInsertSizes = batch .filter(stmt => stmt.kind === 'insert') diff --git a/services/auto-routing-benchmark/src/db-schema.ts b/services/auto-routing-benchmark/src/db-schema.ts index d11e35eae3..bbee012fec 100644 --- a/services/auto-routing-benchmark/src/db-schema.ts +++ b/services/auto-routing-benchmark/src/db-schema.ts @@ -1,6 +1,18 @@ import { sql } from 'drizzle-orm'; -import { integer, primaryKey, real, sqliteTable, text, uniqueIndex } from 'drizzle-orm/sqlite-core'; -import type { BenchmarkKind, BenchmarkRunStatus } from '@kilocode/auto-routing-contracts'; +import { + index, + integer, + primaryKey, + real, + sqliteTable, + text, + uniqueIndex, +} from 'drizzle-orm/sqlite-core'; +import type { + BenchmarkKind, + BenchmarkProfileStatus, + BenchmarkRunStatus, +} from '@kilocode/auto-routing-contracts'; // Migrations are generated via `pnpm db:generate` (drizzle-kit) and applied // via wrangler d1 migrations apply. @@ -66,6 +78,10 @@ export const benchmarkRuns = sqliteTable( // dataset, grading, or CLI/image pinning re-benchmark instead of pairing // current serving config with measurements taken under different conditions. engine_identity: text('engine_identity').notNull().default(''), + // 'platform' (default): publishes the default routing table / classifier winner. + // 'profile': measures exact Pool entries for the global Benchmark-profile + // registry and never replaces the platform artifact. + purpose: text('purpose').notNull().default('platform'), }, table => [ // At most one running run per kind — the atomic backstop for the @@ -82,11 +98,15 @@ export const runModels = sqliteTable( { run_id: text('run_id').notNull(), model: text('model').notNull(), + // Canonical variant key at the D1 boundary. '' means null/default variant. + // Application code converts '' ↔ null at the edges. + variant: text('variant').notNull().default(''), // enqueued=false means the model was skipped (had prior results). enqueued: integer('enqueued', { mode: 'boolean' }).notNull(), + // Legacy mirror of the platform effort key; kept for rollback/provenance. reasoning_effort: text('reasoning_effort'), }, - table => [primaryKey({ columns: [table.run_id, table.model] })] + table => [primaryKey({ columns: [table.run_id, table.model, table.variant] })] ); export const modelSummaries = sqliteTable( @@ -94,6 +114,8 @@ export const modelSummaries = sqliteTable( { run_id: text('run_id').notNull(), model: text('model').notNull(), + // Canonical variant key at the D1 boundary. '' means null/default variant. + variant: text('variant').notNull().default(''), route_key: text('route_key').notNull(), accuracy: real('accuracy').notNull(), avg_cost_usd: real('avg_cost_usd'), @@ -106,7 +128,7 @@ export const modelSummaries = sqliteTable( // carried=true rows are prior-run summaries copied in at startRun for skipped models. carried: integer('carried', { mode: 'boolean' }).notNull().default(false), }, - table => [primaryKey({ columns: [table.run_id, table.model, table.route_key] })] + table => [primaryKey({ columns: [table.run_id, table.model, table.variant, table.route_key] })] ); export const caseResults = sqliteTable( @@ -114,6 +136,8 @@ export const caseResults = sqliteTable( { run_id: text('run_id').notNull(), model: text('model').notNull(), + // Canonical variant key at the D1 boundary. '' means null/default variant. + variant: text('variant').notNull().default(''), case_id: text('case_id').notNull(), route_key: text('route_key'), score: real('score').notNull(), @@ -128,14 +152,16 @@ export const caseResults = sqliteTable( output_prefix: text('output_prefix'), event_count: integer('event_count'), last_event_types: text('last_event_types'), - // Repetition index (0-based); together with run_id/model/case_id forms the PK. + // Repetition index (0-based); together with run_id/model/variant/case_id forms the PK. rep: integer('rep').notNull().default(0), // 1 when the case was killed by the wall-clock timeout, 0 otherwise. timed_out: integer('timed_out').notNull().default(0), }, // The composite PK's leftmost column already serves run_id-prefix lookups // (count/fetch by run); no separate run_id index is needed. - table => [primaryKey({ columns: [table.run_id, table.model, table.case_id, table.rep] })] + table => [ + primaryKey({ columns: [table.run_id, table.model, table.variant, table.case_id, table.rep] }), + ] ); export const routingTables = sqliteTable('routing_tables', { @@ -161,7 +187,65 @@ export const routingTableCandidates = sqliteTable( // cost signal, so every published candidate has one). avg_cost_usd: real('avg_cost_usd').notNull(), meets_threshold: integer('meets_threshold', { mode: 'boolean' }).notNull(), + // Legacy effort key; platform table JSON still reads this column. reasoning_effort: text('reasoning_effort'), + // Exact-pair identity mirror ('' = null). Keeps rows self-describing; no PK change. + variant: text('variant').notNull().default(''), }, table => [primaryKey({ columns: [table.run_id, table.route_key, table.rank] })] ); + +/** + * Global Benchmark-profile registry: one current row per exact Pool entry per + * engine identity + repetitions. Old-engine rows remain as history; currency + * is decided by matching the live decider engine identity and repetitions. + */ +export const benchmarkProfiles = sqliteTable( + 'benchmark_profiles', + { + model: text('model').notNull(), + // Canonical variant key at the D1 boundary. '' means null/default variant. + variant: text('variant').notNull().default(''), + engine_identity: text('engine_identity').notNull(), + repetitions: integer('repetitions').notNull(), + status: text('status').$type().notNull(), + // Provenance of the run that measured / is measuring this profile. + run_id: text('run_id'), + failure_reason: text('failure_reason'), + requested_at: text('requested_at').notNull(), + updated_at: text('updated_at').notNull(), + completed_at: text('completed_at'), + }, + table => [ + primaryKey({ + columns: [table.model, table.variant, table.engine_identity, table.repetitions], + }), + ] +); + +/** + * Rolling owner admission ledger for the 24h profile request quota. One row + * per charged admission (new under current engine, or explicit failed retry). + * Stale engine-drift re-admissions are free and do not insert here. + */ +export const profileRequestEvents = sqliteTable( + 'profile_request_events', + { + id: integer('id').primaryKey({ autoIncrement: true }), + owner_type: text('owner_type').notNull(), + owner_id: text('owner_id').notNull(), + model: text('model').notNull(), + variant: text('variant').notNull().default(''), + engine_identity: text('engine_identity').notNull(), + repetitions: integer('repetitions').notNull(), + admitted_at: text('admitted_at').notNull(), + }, + table => [ + // Lookup window: count an owner's admissions in the preceding 24 hours. + index('IDX_profile_request_events_owner_admitted').on( + table.owner_type, + table.owner_id, + table.admitted_at + ), + ] +); diff --git a/services/auto-routing-benchmark/src/db.test.ts b/services/auto-routing-benchmark/src/db.test.ts index d5d284d67b..b143a67fce 100644 --- a/services/auto-routing-benchmark/src/db.test.ts +++ b/services/auto-routing-benchmark/src/db.test.ts @@ -1,18 +1,54 @@ import { describe, it, expect } from 'vitest'; import { RoutingTableSchema } from '@kilocode/auto-routing-contracts'; import type { RankedCandidate, RoutingTable } from '@kilocode/auto-routing-contracts'; -import { mapRunRow, mapSummaryRow, routingTableToRows, rowsToRoutingTable } from './db'; +import { + exactPairKey, + mapRunRow, + mapSummaryRow, + routingTableToRows, + rowsToRoutingTable, +} from './db'; import type { BenchmarkModelSummary } from '@kilocode/auto-routing-contracts'; +import { + variantFromReasoningEffort, + variantFromStorage, + variantToStorage, +} from './reasoning-effort'; // --------------------------------------------------------------------------- // mapSummaryRow // --------------------------------------------------------------------------- +describe('variant storage helpers', () => { + it('round-trips null and empty as the D1 null/default convention', () => { + expect(variantToStorage(null)).toBe(''); + expect(variantToStorage(undefined)).toBe(''); + expect(variantToStorage('xhigh')).toBe('xhigh'); + expect(variantFromStorage('')).toBeNull(); + expect(variantFromStorage(null)).toBeNull(); + expect(variantFromStorage('xhigh')).toBe('xhigh'); + }); + + it('maps platform reasoningEffort to the stored variant value', () => { + expect(variantFromReasoningEffort('high')).toBe('high'); + expect(variantFromReasoningEffort(null)).toBeNull(); + }); + + it('exactPairKey distinguishes variants of the same model', () => { + expect(exactPairKey('m', 'a')).not.toBe(exactPairKey('m', 'b')); + // Application code normalizes '' → null before calling exactPairKey; the + // helper itself only coalesces nullish (undefined/null), not empty string. + expect(exactPairKey('m', null)).toBe(exactPairKey('m', undefined)); + expect(exactPairKey('m', variantFromStorage(''))).toBe(exactPairKey('m', null)); + }); +}); + describe('mapSummaryRow', () => { it('maps snake_case columns to camelCase BenchmarkModelSummary', () => { const row = { run_id: 'run-1', model: 'openai/gpt-4o', + variant: 'high', route_key: 'implementation/code_generation', accuracy: 0.92, avg_cost_usd: 0.0015, @@ -27,6 +63,7 @@ describe('mapSummaryRow', () => { const result = mapSummaryRow(row); expect(result).toEqual({ model: 'openai/gpt-4o', + variant: 'high', routeKey: 'implementation/code_generation', accuracy: 0.92, avgCostUsd: 0.0015, @@ -39,10 +76,11 @@ describe('mapSummaryRow', () => { }); }); - it('handles null avg_cost_usd and p50_latency_ms', () => { + it('maps stored empty variant to null', () => { const row = { - run_id: 'run-2', - model: 'anthropic/claude-3-haiku', + run_id: 'run-legacy', + model: 'openai/gpt-4o', + variant: '', route_key: '*', accuracy: 0.85, avg_cost_usd: null, @@ -55,6 +93,7 @@ describe('mapSummaryRow', () => { carried: false, }; const result = mapSummaryRow(row); + expect(result.variant).toBeNull(); expect(result.avgCostUsd).toBeNull(); expect(result.p50LatencyMs).toBeNull(); expect(result.p95LatencyMs).toBeNull(); @@ -86,6 +125,7 @@ describe('mapRunRow', () => { repetitions: 1, classifier_max_p95_latency_ms: null, engine_identity: 'v1:deadbeef', + purpose: 'platform', }; const summaries: BenchmarkModelSummary[] = [ { @@ -129,6 +169,7 @@ describe('mapRunRow', () => { repetitions: 1, classifier_max_p95_latency_ms: null, engine_identity: 'v1:deadbeef', + purpose: 'platform', }; const result = mapRunRow(runRow, []); expect(result.summaries).toEqual([]); @@ -203,4 +244,31 @@ describe('rowsToRoutingTable', () => { expect(reassembled.routes['implementation/code_generation']?.[0]?.model).toBe('model-a'); expect(reassembled.routes['implementation/code_generation']?.[1]?.model).toBe('model-b'); }); + + it('platform candidate rows store variant mirror but reassemble without variant key', () => { + const withEffort: RoutingTable = { + ...sampleTable, + routes: { + 'implementation/code_generation': [ + { + model: 'model-a', + accuracy: 0.9, + avgCostUsd: 0.001, + meetsThreshold: true, + reasoningEffort: 'high', + }, + ], + }, + }; + const { candidateRows } = routingTableToRows(withEffort, '2026-06-01T11:00:00.000Z'); + expect(candidateRows[0]?.reasoning_effort).toBe('high'); + expect(candidateRows[0]?.variant).toBe('high'); + const reassembled = rowsToRoutingTable( + routingTableToRows(withEffort, '2026-06-01T11:00:00.000Z').tableRow, + candidateRows + ); + const cand = reassembled.routes['implementation/code_generation']?.[0]; + expect(cand?.reasoningEffort).toBe('high'); + expect(cand && 'variant' in cand ? cand.variant : undefined).toBeUndefined(); + }); }); diff --git a/services/auto-routing-benchmark/src/db.ts b/services/auto-routing-benchmark/src/db.ts index 92239291e4..0b93cf626f 100644 --- a/services/auto-routing-benchmark/src/db.ts +++ b/services/auto-routing-benchmark/src/db.ts @@ -6,12 +6,15 @@ import type { RankedCandidate, RoutingTable, } from '@kilocode/auto-routing-contracts'; +import { poolEntryKey, RoutingTableSchema } from '@kilocode/auto-routing-contracts'; +import type { PoolEntry } from '@kilocode/auto-routing-contracts'; +import { BENCHMARK_PROFILE_FAILURE_REASON_MAX_LENGTH } from '@kilocode/auto-routing-contracts'; import type { BatchItem } from 'drizzle-orm/batch'; -import { RoutingTableSchema } from '@kilocode/auto-routing-contracts'; -import { and, count, desc, eq, gt, inArray, lt, ne } from 'drizzle-orm'; +import { and, asc, count, desc, eq, gt, inArray, lt, ne, or } from 'drizzle-orm'; import { drizzle } from 'drizzle-orm/d1'; import { benchmarkConfig, + benchmarkProfiles, benchmarkRuns, caseResults, configAutoDeciderExclusions, @@ -24,7 +27,11 @@ import { runModels, } from './db-schema'; import { pickClassifierWinner } from './winner'; -import { parsePersistedReasoningEffort } from './reasoning-effort'; +import { + parsePersistedReasoningEffort, + variantFromStorage, + variantToStorage, +} from './reasoning-effort'; export type CaseResultRow = typeof caseResults.$inferSelect; export type RunRow = typeof benchmarkRuns.$inferSelect; @@ -34,13 +41,13 @@ export type ConfigAutoDeciderModelRow = typeof configAutoDeciderModels.$inferSel type ModelSummaryRow = typeof modelSummaries.$inferSelect; // D1 rejects statements with too many bound variables. A model summary insert -// binds 12 values per row, so 8 rows keeps each INSERT below the 100-variable -// ceiling while still batching the delete plus inserts together. -const MODEL_SUMMARY_INSERT_BATCH_SIZE = 8; +// binds 13 values per row (including variant), so 7 rows keeps each INSERT +// below the 100-variable ceiling while still batching the delete plus inserts. +const MODEL_SUMMARY_INSERT_BATCH_SIZE = 7; -// Routing table candidates bind 8 values per row. Keep each INSERT comfortably -// under D1's 100-variable ceiling; publishing is infrequent, so smaller -// statements are preferable to risking a skipped routing-table update. +// Routing table candidates bind 9 values per row (including variant). Keep each +// INSERT comfortably under D1's 100-variable ceiling; publishing is infrequent, +// so smaller statements are preferable to risking a skipped routing-table update. const ROUTING_TABLE_CANDIDATE_INSERT_BATCH_SIZE = 10; // --------------------------------------------------------------------------- @@ -50,6 +57,7 @@ const ROUTING_TABLE_CANDIDATE_INSERT_BATCH_SIZE = 10; export function mapSummaryRow(row: ModelSummaryRow): BenchmarkModelSummary { return { model: row.model, + variant: variantFromStorage(row.variant), routeKey: row.route_key as BenchmarkModelSummary['routeKey'], accuracy: row.accuracy, avgCostUsd: row.avg_cost_usd, @@ -62,6 +70,19 @@ export function mapSummaryRow(row: ModelSummaryRow): BenchmarkModelSummary { }; } +/** + * Summary row with provenance run identity. Used by custom-table assembly so + * candidates are bound to the measuring run, not collapsed across runs. + */ +export type BenchmarkModelSummaryWithRun = BenchmarkModelSummary & { runId: string }; + +export function mapSummaryRowWithRun(row: ModelSummaryRow): BenchmarkModelSummaryWithRun { + return { + ...mapSummaryRow(row), + runId: row.run_id, + }; +} + export function mapRunRow(row: RunRow, summaries: BenchmarkModelSummary[]): BenchmarkRun { return { id: row.id, @@ -173,6 +194,9 @@ export async function replaceAutoDeciderModels( // Runs // --------------------------------------------------------------------------- +/** Run purpose: platform publishes default artifacts; profile updates the registry only. */ +export type BenchmarkRunPurpose = 'platform' | 'profile'; + export async function insertRun( db: D1Database, run: { @@ -188,6 +212,7 @@ export async function insertRun( repetitions: number; classifier_max_p95_latency_ms: number | null; engine_identity: string; + purpose?: BenchmarkRunPurpose; }, models: RunModelRow[], carriedSummaries: BenchmarkModelSummary[] @@ -207,6 +232,7 @@ export async function insertRun( repetitions: run.repetitions, classifier_max_p95_latency_ms: run.classifier_max_p95_latency_ms, engine_identity: run.engine_identity, + purpose: run.purpose ?? 'platform', }); if (models.length === 0 && carriedSummaries.length === 0) { @@ -227,6 +253,7 @@ export async function insertRun( summaryChunk.map(s => ({ run_id: run.id, model: s.model, + variant: variantToStorage(s.variant), route_key: s.routeKey, accuracy: s.accuracy, avg_cost_usd: s.avgCostUsd, @@ -267,7 +294,13 @@ export async function upsertCaseResult(db: D1Database, row: CaseResultRow): Prom .insert(caseResults) .values(row) .onConflictDoUpdate({ - target: [caseResults.run_id, caseResults.model, caseResults.case_id, caseResults.rep], + target: [ + caseResults.run_id, + caseResults.model, + caseResults.variant, + caseResults.case_id, + caseResults.rep, + ], set: { route_key: row.route_key, score: row.score, @@ -301,9 +334,17 @@ export async function getCaseResults(db: D1Database, runId: string): Promise> { if (params.caseIds.length === 0) return new Set(); + const storedVariant = variantToStorage(params.variant); const rows = await drizzle(db) .select({ case_id: caseResults.case_id }) .from(caseResults) @@ -311,6 +352,7 @@ export async function getExistingCaseResultIds( and( eq(caseResults.run_id, params.runId), eq(caseResults.model, params.model), + eq(caseResults.variant, storedVariant), eq(caseResults.rep, params.rep), inArray(caseResults.case_id, params.caseIds) ) @@ -346,6 +388,7 @@ export async function replaceModelSummaries( summaryChunk.map(s => ({ run_id: runId, model: s.model, + variant: variantToStorage(s.variant), route_key: s.routeKey, accuracy: s.accuracy, avg_cost_usd: s.avgCostUsd, @@ -467,24 +510,253 @@ export async function markRunFailed(db: D1Database, runId: string, error: string .where(and(eq(benchmarkRuns.id, runId), eq(benchmarkRuns.status, 'running'))); } +function boundProfileFailureReason(reason: string): string { + if (reason.length <= BENCHMARK_PROFILE_FAILURE_REASON_MAX_LENGTH) return reason; + return reason.slice(0, BENCHMARK_PROFILE_FAILURE_REASON_MAX_LENGTH); +} + +/** + * Mark Benchmark-profile rows claimed by this run as running. Only touches rows + * still pending (or already running for this run_id) under the exact PK — never + * clobbers ready/failed or a different run's claim. + */ +export async function markProfilesRunningForRun( + db: D1Database, + runId: string, + entries: readonly PoolEntry[], + current: { engineIdentity: string; repetitions: number }, + nowIso: string = new Date().toISOString() +): Promise { + if (entries.length === 0) return; + const orm = drizzle(db); + for (const entry of entries) { + await orm + .update(benchmarkProfiles) + .set({ + status: 'running', + run_id: runId, + failure_reason: null, + updated_at: nowIso, + completed_at: null, + }) + .where( + and( + eq(benchmarkProfiles.model, entry.model), + eq(benchmarkProfiles.variant, variantToStorage(entry.variant)), + eq(benchmarkProfiles.engine_identity, current.engineIdentity), + eq(benchmarkProfiles.repetitions, current.repetitions), + or( + eq(benchmarkProfiles.status, 'pending'), + and(eq(benchmarkProfiles.status, 'running'), eq(benchmarkProfiles.run_id, runId)) + ) + ) + ); + } +} + +/** + * Production UPDATE for ready transition. Exported for honest SQLite tests of + * the run_id + status='running' no-clobber guard. + */ +export function markProfilesReadyForRunStatement( + orm: ReturnType, + runId: string, + nowIso: string +) { + return orm + .update(benchmarkProfiles) + .set({ + status: 'ready', + failure_reason: null, + updated_at: nowIso, + completed_at: nowIso, + }) + .where(and(eq(benchmarkProfiles.run_id, runId), eq(benchmarkProfiles.status, 'running'))); +} + +/** + * Transition profiles claimed by this run to ready. Only rows still pointing at + * this run_id and still running are updated — never a newer pending/ready row. + */ +export async function markProfilesReadyForRun( + db: D1Database, + runId: string, + nowIso: string = new Date().toISOString() +): Promise { + await markProfilesReadyForRunStatement(drizzle(db), runId, nowIso); +} + +/** + * Production UPDATE for failed transition. Exported for honest SQLite tests of + * the run_id + status='running' no-clobber guard. + */ +export function markProfilesFailedForRunStatement( + orm: ReturnType, + runId: string, + failureReason: string, + nowIso: string +) { + return orm + .update(benchmarkProfiles) + .set({ + status: 'failed', + failure_reason: boundProfileFailureReason(failureReason), + updated_at: nowIso, + completed_at: nowIso, + }) + .where(and(eq(benchmarkProfiles.run_id, runId), eq(benchmarkProfiles.status, 'running'))); +} + +/** + * Transition profiles claimed by this run to failed with a bounded reason. + * Only rows still pointing at this run_id and still running are updated. + */ +export async function markProfilesFailedForRun( + db: D1Database, + runId: string, + failureReason: string, + nowIso: string = new Date().toISOString() +): Promise { + await markProfilesFailedForRunStatement(drizzle(db), runId, failureReason, nowIso); +} + +/** + * Pending current-engine profile rows, oldest request first. Used by the single + * decider-slot drain to claim the next batch. + */ +export async function listPendingCurrentProfiles( + db: D1Database, + current: { engineIdentity: string; repetitions: number } +): Promise> { + return drizzle(db) + .select({ + model: benchmarkProfiles.model, + variant: benchmarkProfiles.variant, + requested_at: benchmarkProfiles.requested_at, + }) + .from(benchmarkProfiles) + .where( + and( + eq(benchmarkProfiles.status, 'pending'), + eq(benchmarkProfiles.engine_identity, current.engineIdentity), + eq(benchmarkProfiles.repetitions, current.repetitions) + ) + ) + .orderBy( + asc(benchmarkProfiles.requested_at), + asc(benchmarkProfiles.model), + asc(benchmarkProfiles.variant) + ); +} + +/** + * Ready current-engine profiles for the given exact pairs (for custom table assembly). + */ +export async function listReadyCurrentProfilesForEntries( + db: D1Database, + current: { engineIdentity: string; repetitions: number }, + entries: readonly PoolEntry[] +): Promise> { + if (entries.length === 0) return []; + const models = [...new Set(entries.map(e => e.model))]; + const rows = await drizzle(db) + .select({ + model: benchmarkProfiles.model, + variant: benchmarkProfiles.variant, + run_id: benchmarkProfiles.run_id, + engine_identity: benchmarkProfiles.engine_identity, + repetitions: benchmarkProfiles.repetitions, + status: benchmarkProfiles.status, + }) + .from(benchmarkProfiles) + .where( + and( + inArray(benchmarkProfiles.model, models), + eq(benchmarkProfiles.status, 'ready'), + eq(benchmarkProfiles.engine_identity, current.engineIdentity), + eq(benchmarkProfiles.repetitions, current.repetitions) + ) + ); + const wanted = new Set(entries.map(e => exactPairKey(e.model, e.variant))); + return rows + .filter(row => wanted.has(exactPairKey(row.model, variantFromStorage(row.variant)))) + .map(row => ({ + model: row.model, + variant: row.variant, + run_id: row.run_id, + })); +} + +/** + * Load model_summaries for the given run ids (custom table assembly). + * Each row carries its provenance `runId` so assembly can bind candidates to + * the measuring run for each ready entry (no cross-run leakage). + */ +export async function getSummariesForRuns( + db: D1Database, + runIds: readonly string[] +): Promise { + if (runIds.length === 0) return []; + const rows = await drizzle(db) + .select() + .from(modelSummaries) + .where(inArray(modelSummaries.run_id, [...runIds])); + return rows.map(mapSummaryRowWithRun); +} + +/** + * Running decider run ids older than the stale threshold (for profile fail-over + * after the bulk stale sweep). + */ +export async function listStaleRunningDeciderRunIds( + db: D1Database, + olderThanIso: string +): Promise { + const rows = await drizzle(db) + .select({ id: benchmarkRuns.id }) + .from(benchmarkRuns) + .where( + and( + eq(benchmarkRuns.kind, 'decider'), + eq(benchmarkRuns.status, 'running'), + lt(benchmarkRuns.started_at, olderThanIso) + ) + ); + return rows.map(r => r.id); +} + // --------------------------------------------------------------------------- // Latest summaries per model (for skip logic and classifier winner) // --------------------------------------------------------------------------- -// What the most recent completed run measured for a model, plus the -// benchmark identity it was measured under. startRun carries these summaries -// into a new run only when the identity (engine + repetitions + the model's -// reasoning_effort) still matches; otherwise the model is re-benchmarked. +// What the most recent completed run measured for an exact Pool entry, plus +// the benchmark identity it was measured under. startRun carries these +// summaries into a new run only when the identity (engine + repetitions + +// exact variant) still matches; otherwise the entry is re-benchmarked. export type PriorModelResult = { engineIdentity: string; repetitions: number; + /** Canonical variant (null = default). Prefer this for carry matching. */ + variant: string | null; + /** + * Legacy effort mirror from run_models. Still exposed so older callers and + * legacy-row tests can inspect it; carry matching uses `variant`. + */ reasoningEffort: string | null; summaries: BenchmarkModelSummary[]; }; -// Latest summaries per model for a benchmark kind: for each model, all routes -// from the most recent COMPLETED run that included it (mixing routes across -// runs would pair incomparable numbers). +/** + * Canonical map key for an exact (model, variant) pair. Uses the shared + * poolEntryKey contract helper so keys stay collision-safe. + */ +export function exactPairKey(model: string, variant: string | null | undefined): string { + return poolEntryKey({ model, variant: variant ?? null }); +} + +// Latest summaries per exact Pool entry for a benchmark kind: for each +// (model, variant), all routes from the most recent COMPLETED run that +// included that pair (mixing routes across runs would pair incomparable numbers). export async function getLatestSummariesByModel( db: D1Database, kind: BenchmarkKind @@ -493,6 +765,7 @@ export async function getLatestSummariesByModel( .select({ run_id: modelSummaries.run_id, model: modelSummaries.model, + variant: modelSummaries.variant, route_key: modelSummaries.route_key, accuracy: modelSummaries.accuracy, avg_cost_usd: modelSummaries.avg_cost_usd, @@ -511,31 +784,41 @@ export async function getLatestSummariesByModel( .innerJoin(benchmarkRuns, eq(benchmarkRuns.id, modelSummaries.run_id)) .leftJoin( runModels, - and(eq(runModels.run_id, modelSummaries.run_id), eq(runModels.model, modelSummaries.model)) + and( + eq(runModels.run_id, modelSummaries.run_id), + eq(runModels.model, modelSummaries.model), + eq(runModels.variant, modelSummaries.variant) + ) ) .where(and(eq(benchmarkRuns.kind, kind), eq(benchmarkRuns.status, 'completed'))) .orderBy(desc(benchmarkRuns.started_at)); - const latestRunByModel = new Map(); + const latestRunByPair = new Map(); for (const row of results) { - if (!latestRunByModel.has(row.model)) latestRunByModel.set(row.model, row.run_id); + const pairKey = exactPairKey(row.model, variantFromStorage(row.variant)); + if (!latestRunByPair.has(pairKey)) latestRunByPair.set(pairKey, row.run_id); } - const byModel = new Map(); + const byPair = new Map(); for (const row of results) { - if (latestRunByModel.get(row.model) !== row.run_id) continue; - const existing = byModel.get(row.model); + const appVariant = variantFromStorage(row.variant); + const pairKey = exactPairKey(row.model, appVariant); + if (latestRunByPair.get(pairKey) !== row.run_id) continue; + const existing = byPair.get(pairKey); if (existing) { existing.summaries.push(mapSummaryRow(row)); } else { - byModel.set(row.model, { + byPair.set(pairKey, { engineIdentity: row.engine_identity, repetitions: row.repetitions, - reasoningEffort: row.reasoning_effort, + variant: appVariant, + // Prefer the summary/run_models variant; fall back to legacy effort for + // rows migrated before variant was written independently. + reasoningEffort: row.reasoning_effort ?? appVariant, summaries: [mapSummaryRow(row)], }); } } - return byModel; + return byPair; } // --------------------------------------------------------------------------- @@ -562,6 +845,11 @@ export function routingTableToRows( const candidateRows: RoutingTableCandidateRow[] = []; for (const [routeKey, candidates] of Object.entries(table.routes)) { candidates.forEach((c, rank) => { + // Platform tables emit reasoningEffort only; mirror that effort key into + // the self-describing variant column ('' when null). Custom sparse tables + // (later slice) will emit variant and leave reasoning_effort null. + const effortKey = c.reasoningEffort ?? null; + const variantKey = c.variant ?? effortKey; candidateRows.push({ run_id: table.version, route_key: routeKey, @@ -570,7 +858,8 @@ export function routingTableToRows( accuracy: c.accuracy, avg_cost_usd: c.avgCostUsd, meets_threshold: c.meetsThreshold, - reasoning_effort: c.reasoningEffort ?? null, + reasoning_effort: effortKey, + variant: variantToStorage(variantKey), }); }); } @@ -589,6 +878,8 @@ export function rowsToRoutingTable( }); for (const row of sorted) { routeMap[row.route_key] ??= []; + // Platform artifact compatibility: keep reading reasoning_effort exactly + // as today so published JSON shape stays effort-based during rolling deploys. routeMap[row.route_key].push({ model: row.model, accuracy: row.accuracy, diff --git a/services/auto-routing-benchmark/src/profile-runs.test.ts b/services/auto-routing-benchmark/src/profile-runs.test.ts new file mode 100644 index 0000000000..0614b99611 --- /dev/null +++ b/services/auto-routing-benchmark/src/profile-runs.test.ts @@ -0,0 +1,532 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type * as DbModule from './db'; +import type * as CliRunnerModule from './cli-runner'; + +vi.mock('./cli-runner', async importOriginal => { + const actual = await importOriginal(); + return { + ...actual, + destroyDeciderCliContainer: vi.fn().mockResolvedValue(undefined), + runDeciderCaseViaCli: vi.fn(), + warmUpCliContainer: vi.fn().mockResolvedValue(undefined), + }; +}); + +vi.mock('./db', async importOriginal => { + const actual = await importOriginal(); + return { + ...actual, + getConfigRows: vi.fn(), + getLatestSummariesByModel: vi.fn(), + getRunningRun: vi.fn(), + insertRun: vi.fn(), + markStaleRunsFailed: vi.fn(), + listStaleRunningDeciderRunIds: vi.fn(), + listPendingCurrentProfiles: vi.fn(), + markProfilesFailedForRun: vi.fn(), + markProfilesReadyForRun: vi.fn(), + markProfilesRunningForRun: vi.fn(), + markRunCompleted: vi.fn(), + markRunFailed: vi.fn(), + countCaseResults: vi.fn(), + getCaseResults: vi.fn(), + getExistingCaseResultIds: vi.fn(), + getSummaries: vi.fn(), + replaceModelSummaries: vi.fn(), + saveRoutingTable: vi.fn(), + existsNewerCompletedRun: vi.fn(), + getRunWithModels: vi.fn(), + }; +}); + +import { + countCaseResults, + existsNewerCompletedRun, + getCaseResults, + getConfigRows, + getExistingCaseResultIds, + getLatestSummariesByModel, + getRunningRun, + getRunWithModels, + getSummaries, + insertRun, + listPendingCurrentProfiles, + listStaleRunningDeciderRunIds, + markProfilesFailedForRun, + markProfilesReadyForRun, + markProfilesRunningForRun, + markRunCompleted, + markRunFailed, + markStaleRunsFailed, + replaceModelSummaries, + saveRoutingTable, +} from './db'; +import { + drainPendingProfileBatch, + failRunAndDrain, + processJob, + startRun, + sweepStaleRunsAndDrain, +} from './run'; + +const queueSendBatch = vi.fn(); +const kvDelete = vi.fn(); + +const configRow = { + id: 1 as const, + min_accuracy: 0.7, + switch_cost_factor: 3, + best_accuracy_switch_threshold: 0.05, + max_concurrency: 4, + benchmark_user_id: 'user-1', + benchmark_org_id: null, + classifier_repetitions: 1, + decider_repetitions: 1, + classifier_max_p95_latency_ms: 1000, + auto_decider_min_cost_usd: 15, + auto_decider_max_cost_usd: 25, + updated_at: '2026-06-01T00:00:00.000Z', + updated_by: null, +}; + +const env = { + BENCH_DB: {} as D1Database, + BENCH_QUEUE: { sendBatch: queueSendBatch }, + AUTO_ROUTING_CONFIG: { delete: kvDelete }, + INTERNAL_API_SECRET_PROD: { get: async () => 'secret' }, + KILO_WEB_API_BASE_URL: 'https://app.test', + KILO_CLI_API_URL: 'https://api.test', +} as unknown as Env; + +function mockConfig(overrides: Partial = {}) { + vi.mocked(getConfigRows).mockResolvedValue({ + config: { ...configRow, ...overrides }, + // mapConfigRows requires non-empty classifier + decider lists. + classifierModels: ['classifier/a'], + deciderModels: [ + { model: 'platform/a', reasoning_effort: null }, + { model: 'platform/b', reasoning_effort: 'high' }, + ], + autoDeciderModels: [], + excludedAutoDeciderModels: [], + }); +} + +beforeEach(() => { + vi.clearAllMocks(); + mockConfig(); + vi.mocked(getRunningRun).mockResolvedValue(undefined); + vi.mocked(getLatestSummariesByModel).mockResolvedValue(new Map()); + vi.mocked(insertRun).mockResolvedValue(undefined); + vi.mocked(markStaleRunsFailed).mockResolvedValue(undefined); + vi.mocked(listStaleRunningDeciderRunIds).mockResolvedValue([]); + vi.mocked(listPendingCurrentProfiles).mockResolvedValue([]); + vi.mocked(markProfilesRunningForRun).mockResolvedValue(undefined); + vi.mocked(markProfilesReadyForRun).mockResolvedValue(undefined); + vi.mocked(markProfilesFailedForRun).mockResolvedValue(undefined); + vi.mocked(markRunCompleted).mockResolvedValue(undefined); + vi.mocked(markRunFailed).mockResolvedValue(undefined); + vi.mocked(countCaseResults).mockResolvedValue(0); + vi.mocked(existsNewerCompletedRun).mockResolvedValue(false); + vi.mocked(replaceModelSummaries).mockResolvedValue(undefined); + vi.mocked(saveRoutingTable).mockResolvedValue(undefined); + vi.mocked(getCaseResults).mockResolvedValue([]); + vi.mocked(getSummaries).mockResolvedValue([]); + vi.mocked(getExistingCaseResultIds).mockResolvedValue(new Set()); + queueSendBatch.mockResolvedValue(undefined); + kvDelete.mockResolvedValue(undefined); +}); + +describe('startRun — profile purpose', () => { + it('starts a profile run with an explicit two-variants-of-one-model snapshot', async () => { + const entries = [ + { model: 'vendor/m', variant: 'xhigh' }, + { model: 'vendor/m', variant: 'max' }, + ]; + const result = await startRun(env, 'decider', { purpose: 'profile', entries }); + + expect(result.runId).toMatch(/^profile-/); + expect(result.enqueuedModels).toBe(2); + expect(insertRun).toHaveBeenCalledOnce(); + const [, runArg, modelRows] = vi.mocked(insertRun).mock.calls[0]; + expect(runArg.purpose).toBe('profile'); + expect(runArg.kind).toBe('decider'); + expect(modelRows).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + model: 'vendor/m', + variant: 'xhigh', + enqueued: true, + reasoning_effort: null, + }), + expect.objectContaining({ + model: 'vendor/m', + variant: 'max', + enqueued: true, + reasoning_effort: null, + }), + ]) + ); + expect(markProfilesRunningForRun).toHaveBeenCalledWith( + env.BENCH_DB, + result.runId, + entries, + expect.objectContaining({ repetitions: 1 }) + ); + expect(queueSendBatch).toHaveBeenCalled(); + }); + + it('platform runs still write purpose platform and publish-path identity from config', async () => { + const result = await startRun(env, 'decider'); + expect(result.runId).toMatch(/^decider-/); + const [, runArg, modelRows] = vi.mocked(insertRun).mock.calls[0]; + expect(runArg.purpose).toBe('platform'); + expect(modelRows.map(m => m.model).sort()).toEqual(['platform/a', 'platform/b']); + expect(markProfilesRunningForRun).not.toHaveBeenCalled(); + }); + + it('rejects profile runs that are not decider or lack entries', async () => { + await expect(startRun(env, 'classifier', { purpose: 'profile', entries: [] })).rejects.toThrow( + /profile runs must be decider/ + ); + await expect(startRun(env, 'decider', { purpose: 'profile' })).rejects.toThrow( + /non-empty entries/ + ); + }); + + it('rejects and marks run failed when markProfilesRunningForRun rejects', async () => { + const entries = [{ model: 'vendor/m', variant: 'xhigh' }]; + vi.mocked(markProfilesRunningForRun).mockRejectedValueOnce(new Error('claim failed')); + + await expect(startRun(env, 'decider', { purpose: 'profile', entries })).rejects.toThrow( + 'claim failed' + ); + + // Run was marked failed via failRunAndDrain + expect(markRunFailed).toHaveBeenCalledWith( + env.BENCH_DB, + expect.stringMatching(/^profile-/), + 'claim failed' + ); + // No queue work enqueued + expect(queueSendBatch).not.toHaveBeenCalled(); + }); +}); + +describe('profile completion transitions', () => { + function mockProfileRunState(runId: string) { + vi.mocked(getRunWithModels).mockResolvedValue({ + run: { + id: runId, + kind: 'decider', + status: 'running', + started_at: '2026-06-01T00:00:00.000Z', + completed_at: null, + error: null, + min_accuracy: 0.7, + switch_cost_factor: 3, + best_accuracy_switch_threshold: 0.05, + max_concurrency: 4, + benchmark_user_id: 'user-1', + benchmark_org_id: null, + repetitions: 1, + classifier_max_p95_latency_ms: null, + engine_identity: 'v1:test', + purpose: 'profile', + }, + models: [ + { + run_id: runId, + model: 'vendor/m', + variant: 'xhigh', + enqueued: true, + reasoning_effort: null, + }, + ], + }); + } + + it('marks profiles ready on completion and does not publish platform table', async () => { + const runId = 'profile-complete-1'; + mockProfileRunState(runId); + const { DECIDER_CASES } = await import('./datasets/decider-cases'); + vi.mocked(countCaseResults).mockResolvedValue(DECIDER_CASES.length); + vi.mocked(getCaseResults).mockResolvedValue( + DECIDER_CASES.map(c => ({ + run_id: runId, + model: 'vendor/m', + variant: 'xhigh', + case_id: c.id, + route_key: 'implementation/code_generation', + score: 1, + latency_ms: 10, + cost_usd: 0.001, + error: null, + fallback_reason: null, + retried: null, + exit_code: 0, + output_prefix: 'ok', + event_count: 1, + last_event_types: 'x', + rep: 0, + timed_out: 0, + })) + ); + vi.mocked(getSummaries).mockResolvedValue([]); + vi.mocked(listPendingCurrentProfiles).mockResolvedValue([]); + // Case already present → skip CLI; chunk 9999 → no next chunk → finalize. + vi.mocked(getExistingCaseResultIds).mockResolvedValue(new Set([DECIDER_CASES[0].id])); + + await processJob(env, { + runId, + kind: 'decider', + model: 'vendor/m', + variant: 'xhigh', + caseIds: [DECIDER_CASES[0].id], + chunk: 9999, + shard: 0, + shardCount: 1, + rep: 0, + }); + + expect(markRunCompleted).toHaveBeenCalledWith(env.BENCH_DB, runId); + expect(markProfilesReadyForRun).toHaveBeenCalledWith(env.BENCH_DB, runId); + expect(saveRoutingTable).not.toHaveBeenCalled(); + expect(kvDelete).not.toHaveBeenCalled(); + }); + + // No-clobber SQL guard (run_id + status='running') is covered honestly in + // profile-transition-sql.test.ts with real node:sqlite execution. + + it('failRunAndDrain marks profiles failed and attempts drain', async () => { + vi.mocked(listPendingCurrentProfiles).mockResolvedValue([ + { model: 'vendor/next', variant: 'high', requested_at: '2026-06-01T00:00:00.000Z' }, + ]); + // After fail, slot free → drain starts profile run + vi.mocked(getRunningRun).mockResolvedValue(undefined); + + await failRunAndDrain(env, 'profile-fail-1', 'enqueue failed'); + + expect(markRunFailed).toHaveBeenCalledWith(env.BENCH_DB, 'profile-fail-1', 'enqueue failed'); + expect(markProfilesFailedForRun).toHaveBeenCalledWith( + env.BENCH_DB, + 'profile-fail-1', + 'enqueue failed' + ); + // Drain claimed pending and started a profile run + expect(insertRun).toHaveBeenCalled(); + const [, runArg] = vi.mocked(insertRun).mock.calls[0]; + expect(runArg.purpose).toBe('profile'); + }); +}); + +describe('drainPendingProfileBatch', () => { + it('starts the oldest pending batch after slot is free', async () => { + vi.mocked(listPendingCurrentProfiles).mockResolvedValue([ + { model: 'm/old', variant: 'a', requested_at: '2026-06-01T00:00:00.000Z' }, + { model: 'm/new', variant: 'b', requested_at: '2026-06-02T00:00:00.000Z' }, + ]); + const result = await drainPendingProfileBatch(env); + expect(result).not.toBeNull(); + expect(result!.entryCount).toBe(2); + const [, , modelRows] = vi.mocked(insertRun).mock.calls[0]; + expect(modelRows[0].model).toBe('m/old'); + expect(modelRows[1].model).toBe('m/new'); + }); + + it('leaves pending untouched when the decider slot is occupied', async () => { + vi.mocked(getRunningRun).mockResolvedValue({ + id: 'decider-busy', + kind: 'decider', + status: 'running', + started_at: '2026-06-01T00:00:00.000Z', + completed_at: null, + error: null, + min_accuracy: 0.7, + switch_cost_factor: 3, + best_accuracy_switch_threshold: 0.05, + max_concurrency: 4, + benchmark_user_id: 'u', + benchmark_org_id: null, + repetitions: 1, + classifier_max_p95_latency_ms: null, + engine_identity: 'v1:x', + purpose: 'platform', + }); + vi.mocked(listPendingCurrentProfiles).mockResolvedValue([ + { model: 'm/pending', variant: '', requested_at: '2026-06-01T00:00:00.000Z' }, + ]); + const result = await drainPendingProfileBatch(env); + expect(result).toBeNull(); + expect(insertRun).not.toHaveBeenCalled(); + }); + + it('caps the batch by the container concurrency budget', async () => { + // max_concurrency=4, repetitions=2 → max entries = floor(4/2)=2 + mockConfig({ max_concurrency: 4, decider_repetitions: 2 }); + vi.mocked(listPendingCurrentProfiles).mockResolvedValue([ + { model: 'm/1', variant: '', requested_at: '2026-06-01T00:00:00.000Z' }, + { model: 'm/2', variant: '', requested_at: '2026-06-01T00:00:01.000Z' }, + { model: 'm/3', variant: '', requested_at: '2026-06-01T00:00:02.000Z' }, + { model: 'm/4', variant: '', requested_at: '2026-06-01T00:00:03.000Z' }, + ]); + const result = await drainPendingProfileBatch(env); + expect(result!.entryCount).toBe(2); + const [, , modelRows] = vi.mocked(insertRun).mock.calls[0]; + expect(modelRows).toHaveLength(2); + expect(modelRows.map(m => m.model)).toEqual(['m/1', 'm/2']); + }); + + it('fires after a failed decider run via failRunAndDrain', async () => { + vi.mocked(listPendingCurrentProfiles).mockResolvedValue([ + { model: 'm/stranded', variant: 'x', requested_at: '2026-06-01T00:00:00.000Z' }, + ]); + await failRunAndDrain(env, 'decider-failed', 'container gone'); + expect(insertRun).toHaveBeenCalledOnce(); + expect(vi.mocked(insertRun).mock.calls[0][1].purpose).toBe('profile'); + }); +}); + +describe('sweepStaleRunsAndDrain', () => { + it('fails profile claims for stale runs and drains pending work', async () => { + vi.mocked(listStaleRunningDeciderRunIds).mockResolvedValue(['stale-profile-1']); + vi.mocked(listPendingCurrentProfiles).mockResolvedValue([ + { model: 'm/after-stale', variant: '', requested_at: '2026-06-01T00:00:00.000Z' }, + ]); + const result = await sweepStaleRunsAndDrain(env); + expect(result.staleRunIds).toEqual(['stale-profile-1']); + expect(markProfilesFailedForRun).toHaveBeenCalledWith( + env.BENCH_DB, + 'stale-profile-1', + 'timed out' + ); + expect(result.drained?.entryCount).toBe(1); + }); + + it('does not throw when slot is occupied after sweep', async () => { + vi.mocked(getRunningRun).mockResolvedValue({ + id: 'still-running', + kind: 'decider', + status: 'running', + started_at: '2026-06-01T00:00:00.000Z', + completed_at: null, + error: null, + min_accuracy: 0.7, + switch_cost_factor: 3, + best_accuracy_switch_threshold: 0.05, + max_concurrency: 4, + benchmark_user_id: 'u', + benchmark_org_id: null, + repetitions: 1, + classifier_max_p95_latency_ms: null, + engine_identity: 'v1:x', + purpose: 'platform', + }); + vi.mocked(listPendingCurrentProfiles).mockResolvedValue([ + { model: 'm/p', variant: '', requested_at: '2026-06-01T00:00:00.000Z' }, + ]); + const result = await sweepStaleRunsAndDrain(env); + expect(result.drained).toBeNull(); + }); +}); + +describe('platform run completion still publishes', () => { + it('publishes routing table for platform decider finalize', async () => { + const runId = 'decider-platform-1'; + const { DECIDER_CASES } = await import('./datasets/decider-cases'); + const { TAXONOMY_ROUTE_KEYS } = await import('@kilocode/auto-routing-contracts'); + vi.mocked(getRunWithModels).mockResolvedValue({ + run: { + id: runId, + kind: 'decider', + status: 'running', + started_at: '2026-06-01T00:00:00.000Z', + completed_at: null, + error: null, + min_accuracy: 0.7, + switch_cost_factor: 3, + best_accuracy_switch_threshold: 0.05, + max_concurrency: 4, + benchmark_user_id: 'user-1', + benchmark_org_id: null, + repetitions: 1, + classifier_max_p95_latency_ms: null, + engine_identity: 'v1:test', + purpose: 'platform', + }, + models: [ + { + run_id: runId, + model: 'platform/a', + variant: '', + enqueued: true, + reasoning_effort: null, + }, + ], + }); + vi.mocked(countCaseResults).mockResolvedValue(DECIDER_CASES.length); + vi.mocked(getCaseResults).mockResolvedValue( + DECIDER_CASES.map(c => ({ + run_id: runId, + model: 'platform/a', + variant: '', + case_id: c.id, + route_key: 'implementation/code_generation', + score: 1, + latency_ms: 10, + cost_usd: 0.001, + error: null, + fallback_reason: null, + retried: null, + exit_code: 0, + output_prefix: 'ok', + event_count: 1, + last_event_types: 'x', + rep: 0, + timed_out: 0, + })) + ); + // Summaries covering every taxonomy route so platform table validates. + vi.mocked(getSummaries).mockResolvedValue( + TAXONOMY_ROUTE_KEYS.map(routeKey => ({ + model: 'platform/a', + variant: null, + routeKey, + accuracy: 0.9, + avgCostUsd: 0.001, + avgLatencyMs: 100, + p50LatencyMs: 90, + p95LatencyMs: 120, + cases: 5, + errors: 0, + timeouts: 0, + })) + ); + vi.mocked(getExistingCaseResultIds).mockResolvedValue(new Set([DECIDER_CASES[0].id])); + + await processJob(env, { + runId, + kind: 'decider', + model: 'platform/a', + variant: null, + caseIds: [DECIDER_CASES[0].id], + chunk: 9999, + shard: 0, + shardCount: 1, + rep: 0, + }); + + expect(markRunCompleted).toHaveBeenCalled(); + expect(saveRoutingTable).toHaveBeenCalled(); + expect(markProfilesReadyForRun).not.toHaveBeenCalled(); + const table = vi.mocked(saveRoutingTable).mock.calls[0][1]; + // Platform artifact: candidates use reasoningEffort, not variant + const firstRoute = Object.values(table.routes)[0]; + expect(firstRoute[0]).toMatchObject({ model: 'platform/a' }); + expect( + firstRoute[0].reasoningEffort === null || + firstRoute[0].reasoningEffort === undefined || + typeof firstRoute[0].reasoningEffort === 'string' + ).toBe(true); + }); +}); diff --git a/services/auto-routing-benchmark/src/profile-transition-sql.test.ts b/services/auto-routing-benchmark/src/profile-transition-sql.test.ts new file mode 100644 index 0000000000..1d1d659c7f --- /dev/null +++ b/services/auto-routing-benchmark/src/profile-transition-sql.test.ts @@ -0,0 +1,283 @@ +/** + * Real SQLite execution of production profile ready/failed transition SQL. + * Intentionally does NOT mock drizzle — these tests fail if the run_id + + * status='running' no-clobber guard is missing or weakened. + */ +import { DatabaseSync } from 'node:sqlite'; +import { beforeEach, describe, expect, it } from 'vitest'; +import { drizzle } from 'drizzle-orm/d1'; +import { markProfilesFailedForRunStatement, markProfilesReadyForRunStatement } from './db'; + +const SCHEMA_SQL = ` + CREATE TABLE benchmark_profiles ( + model TEXT NOT NULL, + variant TEXT NOT NULL DEFAULT '', + engine_identity TEXT NOT NULL, + repetitions INTEGER NOT NULL, + status TEXT NOT NULL, + run_id TEXT, + failure_reason TEXT, + requested_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + completed_at TEXT, + PRIMARY KEY (model, variant, engine_identity, repetitions) + ); +`; + +type SqlQuery = { sql: string; params: unknown[] }; + +function toSql(stmt: { toSQL: () => SqlQuery }): SqlQuery { + return stmt.toSQL(); +} + +/** node:sqlite Statement.run overloads confuse TS with mixed string/number args. */ +function execSql(db: DatabaseSync, sqlText: string, params: readonly unknown[] = []): void { + db.prepare(sqlText).run(...(params as never[])); +} + +function runQuery(db: DatabaseSync, q: SqlQuery): void { + execSql(db, q.sql, q.params); +} + +function insertProfile( + db: DatabaseSync, + row: { + model: string; + variant?: string; + engine_identity?: string; + repetitions?: number; + status: string; + run_id: string | null; + failure_reason?: string | null; + requested_at?: string; + updated_at?: string; + completed_at?: string | null; + } +): void { + execSql( + db, + `INSERT INTO benchmark_profiles + (model, variant, engine_identity, repetitions, status, run_id, failure_reason, requested_at, updated_at, completed_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + [ + row.model, + row.variant ?? 'xhigh', + row.engine_identity ?? 'v-test:engine', + row.repetitions ?? 1, + row.status, + row.run_id, + row.failure_reason ?? null, + row.requested_at ?? '2026-07-28T10:00:00.000Z', + row.updated_at ?? '2026-07-28T10:00:00.000Z', + row.completed_at ?? null, + ] + ); +} + +type ProfileRow = { + model: string; + status: string; + run_id: string | null; + failure_reason: string | null; + completed_at: string | null; +}; + +function allProfiles(db: DatabaseSync): ProfileRow[] { + return db + .prepare( + 'SELECT model, status, run_id, failure_reason, completed_at FROM benchmark_profiles ORDER BY model' + ) + .all() as ProfileRow[]; +} + +describe('markProfilesReadyForRunStatement (real SQLite)', () => { + let db: DatabaseSync; + const orm = drizzle({} as D1Database); + const now = '2026-07-28T12:00:00.000Z'; + + beforeEach(() => { + db = new DatabaseSync(':memory:'); + db.exec(SCHEMA_SQL); + }); + + it('transitions running rows owned by this run_id to ready', () => { + insertProfile(db, { model: 'm/owned', status: 'running', run_id: 'run-old' }); + + runQuery(db, toSql(markProfilesReadyForRunStatement(orm, 'run-old', now))); + + expect(allProfiles(db)).toEqual([ + { + model: 'm/owned', + status: 'ready', + run_id: 'run-old', + failure_reason: null, + completed_at: now, + }, + ]); + }); + + it('does not update a newer pending row for the same entry (different run_id)', () => { + // Same PK identity: entry was re-admitted pending under a new request while + // an older run finishes. The newer pending row must not be clobbered. + insertProfile(db, { + model: 'm/same', + status: 'pending', + run_id: 'run-newer', + requested_at: '2026-07-28T11:00:00.000Z', + updated_at: '2026-07-28T11:00:00.000Z', + }); + + runQuery(db, toSql(markProfilesReadyForRunStatement(orm, 'run-old', now))); + + expect(allProfiles(db)).toEqual([ + { + model: 'm/same', + status: 'pending', + run_id: 'run-newer', + failure_reason: null, + completed_at: null, + }, + ]); + }); + + it('does not overwrite a ready row when a different run completes', () => { + insertProfile(db, { + model: 'm/ready', + status: 'ready', + run_id: 'run-winner', + completed_at: '2026-07-28T09:00:00.000Z', + }); + + runQuery(db, toSql(markProfilesReadyForRunStatement(orm, 'run-other', now))); + + expect(allProfiles(db)).toEqual([ + { + model: 'm/ready', + status: 'ready', + run_id: 'run-winner', + failure_reason: null, + completed_at: '2026-07-28T09:00:00.000Z', + }, + ]); + }); +}); + +describe('markProfilesFailedForRunStatement (real SQLite)', () => { + let db: DatabaseSync; + const orm = drizzle({} as D1Database); + const now = '2026-07-28T12:00:00.000Z'; + + beforeEach(() => { + db = new DatabaseSync(':memory:'); + db.exec(SCHEMA_SQL); + }); + + it('transitions running rows owned by this run_id to failed', () => { + insertProfile(db, { model: 'm/owned', status: 'running', run_id: 'run-old' }); + + runQuery(db, toSql(markProfilesFailedForRunStatement(orm, 'run-old', 'enqueue failed', now))); + + expect(allProfiles(db)).toEqual([ + { + model: 'm/owned', + status: 'failed', + run_id: 'run-old', + failure_reason: 'enqueue failed', + completed_at: now, + }, + ]); + }); + + it('does not update a newer pending row when an older run fails', () => { + insertProfile(db, { + model: 'm/same', + status: 'pending', + run_id: 'run-newer', + requested_at: '2026-07-28T11:00:00.000Z', + updated_at: '2026-07-28T11:00:00.000Z', + }); + + runQuery(db, toSql(markProfilesFailedForRunStatement(orm, 'run-old', 'boom', now))); + + expect(allProfiles(db)).toEqual([ + { + model: 'm/same', + status: 'pending', + run_id: 'run-newer', + failure_reason: null, + completed_at: null, + }, + ]); + }); + + it('does not overwrite a ready row when a different run fails', () => { + insertProfile(db, { + model: 'm/ready', + status: 'ready', + run_id: 'run-winner', + completed_at: '2026-07-28T09:00:00.000Z', + }); + + runQuery(db, toSql(markProfilesFailedForRunStatement(orm, 'run-other', 'late fail', now))); + + expect(allProfiles(db)).toEqual([ + { + model: 'm/ready', + status: 'ready', + run_id: 'run-winner', + failure_reason: null, + completed_at: '2026-07-28T09:00:00.000Z', + }, + ]); + }); + + it('running-scoped update only affects this run_id among mixed rows', () => { + insertProfile(db, { model: 'm/a', status: 'running', run_id: 'run-old' }); + insertProfile(db, { + model: 'm/b', + status: 'pending', + run_id: null, + requested_at: '2026-07-28T11:00:00.000Z', + }); + insertProfile(db, { + model: 'm/c', + status: 'ready', + run_id: 'run-other', + completed_at: '2026-07-28T08:00:00.000Z', + }); + insertProfile(db, { model: 'm/d', status: 'running', run_id: 'run-other' }); + + runQuery(db, toSql(markProfilesFailedForRunStatement(orm, 'run-old', 'timeout', now))); + + expect(allProfiles(db)).toEqual([ + { + model: 'm/a', + status: 'failed', + run_id: 'run-old', + failure_reason: 'timeout', + completed_at: now, + }, + { + model: 'm/b', + status: 'pending', + run_id: null, + failure_reason: null, + completed_at: null, + }, + { + model: 'm/c', + status: 'ready', + run_id: 'run-other', + failure_reason: null, + completed_at: '2026-07-28T08:00:00.000Z', + }, + { + model: 'm/d', + status: 'running', + run_id: 'run-other', + failure_reason: null, + completed_at: null, + }, + ]); + }); +}); diff --git a/services/auto-routing-benchmark/src/profiles-sql.test.ts b/services/auto-routing-benchmark/src/profiles-sql.test.ts new file mode 100644 index 0000000000..3c765152e6 --- /dev/null +++ b/services/auto-routing-benchmark/src/profiles-sql.test.ts @@ -0,0 +1,270 @@ +/** + * Real SQLite execution of production admission SQL builders. + * Intentionally does NOT mock drizzle — these tests fail if charged-event + * column counts drift or the NOT EXISTS guard is missing. + */ +import { DatabaseSync } from 'node:sqlite'; +import { beforeEach, describe, expect, it } from 'vitest'; +import { drizzle } from 'drizzle-orm/d1'; +import { + buildPendingUpsertValues, + chargedEventInsertStatement, + pendingRegisterUpsertStatement, + pendingStatusInsertStatement, +} from './profiles'; + +const SCHEMA_SQL = ` + CREATE TABLE benchmark_profiles ( + model TEXT NOT NULL, + variant TEXT NOT NULL DEFAULT '', + engine_identity TEXT NOT NULL, + repetitions INTEGER NOT NULL, + status TEXT NOT NULL, + run_id TEXT, + failure_reason TEXT, + requested_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + completed_at TEXT, + PRIMARY KEY (model, variant, engine_identity, repetitions) + ); + CREATE TABLE profile_request_events ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + owner_type TEXT NOT NULL, + owner_id TEXT NOT NULL, + model TEXT NOT NULL, + variant TEXT NOT NULL DEFAULT '', + engine_identity TEXT NOT NULL, + repetitions INTEGER NOT NULL, + admitted_at TEXT NOT NULL + ); +`; + +type SqlQuery = { sql: string; params: unknown[] }; + +function toSql(stmt: { toSQL: () => SqlQuery }): SqlQuery { + return stmt.toSQL(); +} + +/** node:sqlite Statement.run overloads confuse TS with mixed string/number args. */ +function execSql(db: DatabaseSync, sqlText: string, params: readonly unknown[] = []): void { + db.prepare(sqlText).run(...(params as never[])); +} + +function runQuery(db: DatabaseSync, q: SqlQuery): void { + execSql(db, q.sql, q.params); +} + +const eventValues = { + owner_type: 'user', + owner_id: 'owner-a', + model: 'openai/gpt-4o', + variant: 'xhigh', + engine_identity: 'v-test:engine', + repetitions: 1, + admitted_at: '2026-07-28T12:00:00.000Z', +}; + +describe('chargedEventInsertStatement (real SQLite)', () => { + let db: DatabaseSync; + // drizzle/d1 only needs a stub for toSQL compilation; we execute via node:sqlite. + const orm = drizzle({} as D1Database); + + beforeEach(() => { + db = new DatabaseSync(':memory:'); + db.exec(SCHEMA_SQL); + }); + + it('emits column-aligned INSERT...SELECT with NOT EXISTS guard', () => { + const q = toSql(chargedEventInsertStatement(orm, eventValues)); + expect(q.sql.toLowerCase()).toContain('insert into "profile_request_events"'); + // Full table column list includes autoincrement id. + expect(q.sql).toMatch( + /"id".*"owner_type".*"owner_id".*"model".*"variant".*"engine_identity".*"repetitions".*"admitted_at"/s + ); + expect(q.sql).toMatch(/null\s+as\s+id/i); + expect(q.sql.toLowerCase()).toContain('where not exists'); + expect(q.sql.toLowerCase()).toContain(`status" in ('pending', 'running', 'ready')`); + }); + + it('inserts a charged event when no active current profile row exists', () => { + runQuery(db, toSql(chargedEventInsertStatement(orm, eventValues))); + const rows = db.prepare('SELECT * FROM profile_request_events').all() as Array<{ + id: number; + owner_id: string; + model: string; + }>; + expect(rows).toHaveLength(1); + expect(rows[0]).toMatchObject({ + id: 1, + owner_id: 'owner-a', + model: 'openai/gpt-4o', + }); + }); + + it('skips the charge when an active (pending) current profile row exists', () => { + execSql( + db, + `INSERT INTO benchmark_profiles + (model, variant, engine_identity, repetitions, status, run_id, failure_reason, requested_at, updated_at, completed_at) + VALUES (?, ?, ?, ?, 'pending', NULL, NULL, ?, ?, NULL)`, + [ + eventValues.model, + eventValues.variant, + eventValues.engine_identity, + eventValues.repetitions, + '2026-07-28T11:00:00.000Z', + '2026-07-28T11:00:00.000Z', + ] + ); + + runQuery(db, toSql(chargedEventInsertStatement(orm, eventValues))); + expect(db.prepare('SELECT COUNT(*) AS n FROM profile_request_events').get()).toEqual({ n: 0 }); + }); + + it('skips the charge for running and ready active rows', () => { + for (const status of ['running', 'ready'] as const) { + db.exec('DELETE FROM benchmark_profiles; DELETE FROM profile_request_events;'); + execSql( + db, + `INSERT INTO benchmark_profiles + (model, variant, engine_identity, repetitions, status, run_id, failure_reason, requested_at, updated_at, completed_at) + VALUES (?, ?, ?, ?, ?, 'run-1', NULL, ?, ?, NULL)`, + [ + eventValues.model, + eventValues.variant, + eventValues.engine_identity, + eventValues.repetitions, + status, + '2026-07-28T11:00:00.000Z', + '2026-07-28T11:00:00.000Z', + ] + ); + runQuery(db, toSql(chargedEventInsertStatement(orm, eventValues))); + expect(db.prepare('SELECT COUNT(*) AS n FROM profile_request_events').get()).toEqual({ + n: 0, + }); + } + }); + + it('charges when the current row is failed (retry path)', () => { + execSql( + db, + `INSERT INTO benchmark_profiles + (model, variant, engine_identity, repetitions, status, run_id, failure_reason, requested_at, updated_at, completed_at) + VALUES (?, ?, ?, ?, 'failed', 'old', 'boom', ?, ?, NULL)`, + [ + eventValues.model, + eventValues.variant, + eventValues.engine_identity, + eventValues.repetitions, + '2026-07-28T11:00:00.000Z', + '2026-07-28T11:00:00.000Z', + ] + ); + + runQuery(db, toSql(chargedEventInsertStatement(orm, eventValues))); + expect(db.prepare('SELECT COUNT(*) AS n FROM profile_request_events').get()).toEqual({ n: 1 }); + }); +}); + +describe('pendingRegisterUpsertStatement (real SQLite)', () => { + let db: DatabaseSync; + const orm = drizzle({} as D1Database); + const values = buildPendingUpsertValues( + { model: 'm/a', variant: 'xhigh' }, + { engineIdentity: 'v-test:engine', repetitions: 1 }, + '2026-07-28T12:00:00.000Z' + ); + + beforeEach(() => { + db = new DatabaseSync(':memory:'); + db.exec(SCHEMA_SQL); + }); + + it('inserts a new pending row', () => { + runQuery(db, toSql(pendingRegisterUpsertStatement(orm, values))); + const row = db.prepare('SELECT status, run_id FROM benchmark_profiles').get() as { + status: string; + run_id: string | null; + }; + expect(row).toEqual({ status: 'pending', run_id: null }); + }); + + it('transitions failed → pending and clears provenance', () => { + execSql( + db, + `INSERT INTO benchmark_profiles + (model, variant, engine_identity, repetitions, status, run_id, failure_reason, requested_at, updated_at, completed_at) + VALUES (?, ?, ?, ?, 'failed', 'old-run', 'boom', 't0', 't0', 't1')`, + [values.model, values.variant, values.engine_identity, values.repetitions] + ); + + runQuery(db, toSql(pendingRegisterUpsertStatement(orm, values))); + const row = db + .prepare('SELECT status, run_id, failure_reason, completed_at FROM benchmark_profiles') + .get(); + expect(row).toEqual({ + status: 'pending', + run_id: null, + failure_reason: null, + completed_at: null, + }); + }); + + it('does not regress running or ready rows', () => { + for (const status of ['running', 'ready'] as const) { + db.exec('DELETE FROM benchmark_profiles;'); + execSql( + db, + `INSERT INTO benchmark_profiles + (model, variant, engine_identity, repetitions, status, run_id, failure_reason, requested_at, updated_at, completed_at) + VALUES (?, ?, ?, ?, ?, 'live-run', NULL, 't0', 't0', NULL)`, + [values.model, values.variant, values.engine_identity, values.repetitions, status] + ); + + runQuery(db, toSql(pendingRegisterUpsertStatement(orm, values))); + const row = db.prepare('SELECT status, run_id FROM benchmark_profiles').get(); + expect(row).toEqual({ status, run_id: 'live-run' }); + } + }); +}); + +describe('pendingStatusInsertStatement (real SQLite)', () => { + let db: DatabaseSync; + const orm = drizzle({} as D1Database); + const values = buildPendingUpsertValues( + { model: 'stale/m', variant: null }, + { engineIdentity: 'v-test:engine', repetitions: 1 }, + '2026-07-28T12:00:00.000Z' + ); + + beforeEach(() => { + db = new DatabaseSync(':memory:'); + db.exec(SCHEMA_SQL); + }); + + it('inserts when absent and does nothing on conflict with a current row', () => { + runQuery(db, toSql(pendingStatusInsertStatement(orm, values))); + expect(db.prepare('SELECT status FROM benchmark_profiles').get()).toEqual({ + status: 'pending', + }); + + execSql(db, `UPDATE benchmark_profiles SET status = 'ready', run_id = 'winner'`); + runQuery( + db, + toSql( + pendingStatusInsertStatement(orm, { + ...values, + status: 'pending', + run_id: null, + requested_at: '2026-07-28T13:00:00.000Z', + updated_at: '2026-07-28T13:00:00.000Z', + }) + ) + ); + expect(db.prepare('SELECT status, run_id FROM benchmark_profiles').get()).toEqual({ + status: 'ready', + run_id: 'winner', + }); + }); +}); diff --git a/services/auto-routing-benchmark/src/profiles.test.ts b/services/auto-routing-benchmark/src/profiles.test.ts new file mode 100644 index 0000000000..0448fd19cb --- /dev/null +++ b/services/auto-routing-benchmark/src/profiles.test.ts @@ -0,0 +1,1016 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type { PoolEntry } from '@kilocode/auto-routing-contracts'; +import type * as DbSchemaModule from './db-schema'; +import type * as DrizzleOrmModule from 'drizzle-orm'; +import type { ProfileRow } from './profiles'; +import type * as RunModule from './run'; + +// --------------------------------------------------------------------------- +// In-memory D1/drizzle stand-in for admission + status tests. +// Honestly models ON CONFLICT WHERE (failed-only), onConflictDoNothing, +// INSERT...SELECT...WHERE NOT EXISTS charge guards, and post-batch reads. +// --------------------------------------------------------------------------- + +type EventRow = { + id: number; + owner_type: string; + owner_id: string; + model: string; + variant: string; + engine_identity: string; + repetitions: number; + admitted_at: string; +}; + +type PendingStmt = + | { kind: 'event'; values: Omit; guarded: boolean } + | { + kind: 'profile'; + values: ProfileRow; + onConflict: 'update_if_failed' | 'do_nothing' | 'replace'; + }; + +const store = vi.hoisted(() => { + const profiles = new Map(); + const events: EventRow[] = []; + let nextEventId = 1; + + function profilePk(row: { + model: string; + variant: string; + engine_identity: string; + repetitions: number; + }): string { + return JSON.stringify([row.model, row.variant, row.engine_identity, row.repetitions]); + } + + return { + profiles, + events, + nextEventId: () => nextEventId++, + profilePk, + reset() { + profiles.clear(); + events.length = 0; + nextEventId = 1; + }, + seedProfile(row: ProfileRow) { + profiles.set(profilePk(row), { ...row }); + }, + seedEvent(row: Omit) { + events.push({ ...row, id: nextEventId++ }); + }, + }; +}); + +vi.mock('drizzle-orm/d1', () => { + type WhereClause = + | { type: 'models'; models: string[] } + | { + type: 'events'; + ownerType: string; + ownerId: string; + windowStart: string; + }; + + function createSelectBuilder(table: 'profiles' | 'events') { + let where: WhereClause | null = null; + let orderAsc = false; + + const builder = { + from() { + return builder; + }, + where(clause: WhereClause) { + where = clause; + return builder; + }, + orderBy() { + orderAsc = true; + return builder; + }, + then(resolve: (value: unknown) => unknown, reject?: (err: unknown) => unknown) { + return Promise.resolve(builder.execute()).then(resolve, reject); + }, + async execute() { + if (table === 'profiles') { + const all = [...store.profiles.values()]; + if (where?.type === 'models') { + const set = new Set(where.models); + return all.filter(r => set.has(r.model)); + } + return all; + } + let rows = [...store.events]; + if (where && where.type === 'events') { + const eventWhere = where; + rows = rows.filter( + e => + e.owner_type === eventWhere.ownerType && + e.owner_id === eventWhere.ownerId && + e.admitted_at >= eventWhere.windowStart + ); + } + if (orderAsc) { + rows.sort((a, b) => a.admitted_at.localeCompare(b.admitted_at) || a.id - b.id); + } + return rows; + }, + }; + return builder; + } + + function profileValuesFromRow(row: Record): ProfileRow { + return { + model: String(row.model), + variant: typeof row.variant === 'string' ? row.variant : '', + engine_identity: String(row.engine_identity), + repetitions: Number(row.repetitions), + status: row.status as ProfileRow['status'], + run_id: (row.run_id as string | null) ?? null, + failure_reason: (row.failure_reason as string | null) ?? null, + requested_at: String(row.requested_at), + updated_at: String(row.updated_at), + completed_at: (row.completed_at as string | null) ?? null, + }; + } + + function eventValuesFromRow(row: Record): Omit { + return { + owner_type: String(row.owner_type), + owner_id: String(row.owner_id), + model: String(row.model), + variant: typeof row.variant === 'string' ? row.variant : '', + engine_identity: String(row.engine_identity), + repetitions: Number(row.repetitions), + admitted_at: String(row.admitted_at), + }; + } + + function wrapStmt(stmt: PendingStmt) { + return Object.assign(Promise.resolve(stmt), { __stmt: stmt }); + } + + function createOrm() { + return { + select() { + return { + from(table: { _: { name: string } } | unknown) { + const name = + table && typeof table === 'object' && '_' in table + ? String((table as { _: { name: string } })._.name) + : ''; + const tagged = table as { __tableName?: string }; + if ( + tagged.__tableName === 'profile_request_events' || + name.includes('profile_request') + ) { + return createSelectBuilder('events'); + } + return createSelectBuilder('profiles'); + }, + }; + }, + insert(table: { __tableName?: string }) { + const isEvents = table.__tableName === 'profile_request_events'; + return { + values(values: unknown) { + const row = values as Record; + if (isEvents) { + const stmt: PendingStmt = { + kind: 'event', + values: eventValuesFromRow(row), + guarded: false, + }; + return Object.assign(wrapStmt(stmt), { + async execute() { + applyStmt(stmt); + }, + }); + } + const profileValues = profileValuesFromRow(row); + return { + onConflictDoUpdate(config?: { where?: { type?: string; value?: string } }) { + // Register path: WHERE status = 'failed' + const whereFailed = + !!config?.where && + (config.where as { type?: string; value?: string }).type === 'eq' && + (config.where as { value?: string }).value === 'failed'; + const stmt: PendingStmt = { + kind: 'profile', + values: profileValues, + onConflict: whereFailed ? 'update_if_failed' : 'replace', + }; + return wrapStmt(stmt); + }, + onConflictDoNothing() { + const stmt: PendingStmt = { + kind: 'profile', + values: profileValues, + onConflict: 'do_nothing', + }; + return wrapStmt(stmt); + }, + }; + }, + /** + * INSERT ... SELECT ... WHERE NOT EXISTS — charged event guard. + * Production Object.assigns __stmt onto the returned value. + */ + select(_query: unknown) { + // Placeholder object; chargedEventInsertStatement Object.assigns __stmt. + return {}; + }, + }; + }, + async batch(stmts: Array<{ __stmt?: PendingStmt } | PendingStmt>) { + for (const s of stmts) { + const stmt = (s as { __stmt?: PendingStmt }).__stmt ?? (s as PendingStmt); + if (!stmt || typeof stmt !== 'object' || !('kind' in stmt)) { + throw new Error('batch received untagged statement'); + } + applyStmt(stmt); + } + }, + }; + } + + function hasActiveCurrentProfile(values: { + model: string; + variant: string; + engine_identity: string; + repetitions: number; + }): boolean { + const existing = store.profiles.get(store.profilePk(values)); + return ( + !!existing && + (existing.status === 'pending' || + existing.status === 'running' || + existing.status === 'ready') + ); + } + + function applyStmt(stmt: PendingStmt) { + if (stmt.kind === 'event') { + if (stmt.guarded && hasActiveCurrentProfile(stmt.values)) { + // INSERT...SELECT...WHERE NOT EXISTS → zero rows. + return; + } + store.events.push({ ...stmt.values, id: store.nextEventId() }); + return; + } + + const pk = store.profilePk(stmt.values); + const existing = store.profiles.get(pk); + if (!existing) { + store.profiles.set(pk, { ...stmt.values }); + return; + } + if (stmt.onConflict === 'do_nothing') { + return; + } + if (stmt.onConflict === 'update_if_failed') { + if (existing.status !== 'failed') { + // WHERE status = 'failed' not satisfied — leave row intact. + return; + } + } + store.profiles.set(pk, { ...stmt.values }); + } + + return { + drizzle: vi.fn(() => createOrm()), + }; +}); + +// Tag schema tables so the mock can distinguish them. +vi.mock('./db-schema', async importOriginal => { + const actual = await importOriginal(); + Object.assign(actual.benchmarkProfiles, { __tableName: 'benchmark_profiles' }); + Object.assign(actual.profileRequestEvents, { __tableName: 'profile_request_events' }); + return actual; +}); + +// Stabilize engine identity for currency tests that swap only repetitions. +vi.mock('./run', async importOriginal => { + const actual = await importOriginal(); + return { + ...actual, + computeEngineIdentity: vi.fn(() => 'v-test:engine'), + }; +}); + +// drizzle operators: return plain descriptors the mock can read. +vi.mock('drizzle-orm', async importOriginal => { + const actual = await importOriginal(); + return { + ...actual, + inArray: (_col: unknown, values: string[]) => ({ type: 'models' as const, models: values }), + and: (...parts: unknown[]) => { + const asEq = parts.filter( + (p): p is { type: 'eq'; value: string } => + !!p && typeof p === 'object' && (p as { type?: string }).type === 'eq' + ); + const gtePart = parts.find( + (p): p is { type: 'gte'; value: string } => + !!p && typeof p === 'object' && (p as { type?: string }).type === 'gte' + ); + if (asEq.length >= 2 && gtePart) { + return { + type: 'events' as const, + ownerType: asEq[0]!.value, + ownerId: asEq[1]!.value, + windowStart: gtePart.value, + }; + } + return { type: 'and', parts }; + }, + eq: (_col: unknown, value: unknown) => ({ type: 'eq' as const, value: String(value) }), + gte: (_col: unknown, value: unknown) => ({ type: 'gte' as const, value: String(value) }), + asc: (col: unknown) => col, + }; +}); + +import { + buildPendingUpsertValues, + chargedEventInsertStatement, + classifyProfileAdmission, + computeQuotaRetryAt, + isCurrentBenchmarkProfile, + lookupProfileStatuses, + MISSING_PROFILE_FAILURE_REASON, + pendingRegisterUpsertStatement, + pendingStatusInsertStatement, + PROFILE_ADMISSION_LIMIT, + PROFILE_ADMISSION_WINDOW_MS, + ProfileQuotaExceededError, + registerProfiles, +} from './profiles'; +import { variantToStorage } from './reasoning-effort'; +import { computeEngineIdentity } from './run'; + +const ENGINE = 'v-test:engine'; +const REPS = 1; +const NOW = new Date('2026-07-28T12:00:00.000Z'); + +function entry(model: string, variant: string | null = null): PoolEntry { + return { model, variant }; +} + +function profileRow( + partial: Partial & Pick +): ProfileRow { + return { + variant: '', + engine_identity: ENGINE, + repetitions: REPS, + run_id: null, + failure_reason: null, + requested_at: '2026-07-27T00:00:00.000Z', + updated_at: '2026-07-27T00:00:00.000Z', + completed_at: null, + ...partial, + }; +} + +/** + * Build production admission statements (real builders) and apply them through + * the mock batch path. Exercises chargedEventInsertStatement / + * pendingRegisterUpsertStatement / pendingStatusInsertStatement so a missing + * __stmt tag or broken builder wiring fails these race tests. + */ +async function applyProductionAdmissionBatch(opts: { + chargedEvents?: Array<{ + ownerType: string; + ownerId: string; + entry: PoolEntry; + }>; + registerUpserts?: PoolEntry[]; + statusInserts?: PoolEntry[]; + nowIso?: string; +}): Promise { + const { drizzle } = await import('drizzle-orm/d1'); + const orm = drizzle({} as D1Database); + const nowIso = opts.nowIso ?? NOW.toISOString(); + const current = { engineIdentity: ENGINE, repetitions: REPS }; + const stmts: Array<{ __stmt?: PendingStmt }> = []; + + for (const ev of opts.chargedEvents ?? []) { + stmts.push( + chargedEventInsertStatement(orm, { + owner_type: ev.ownerType, + owner_id: ev.ownerId, + model: ev.entry.model, + variant: variantToStorage(ev.entry.variant), + engine_identity: ENGINE, + repetitions: REPS, + admitted_at: nowIso, + }) as { __stmt?: PendingStmt } + ); + } + for (const entry of opts.registerUpserts ?? []) { + stmts.push( + pendingRegisterUpsertStatement(orm, buildPendingUpsertValues(entry, current, nowIso)) as { + __stmt?: PendingStmt; + } + ); + } + for (const entry of opts.statusInserts ?? []) { + stmts.push( + pendingStatusInsertStatement(orm, buildPendingUpsertValues(entry, current, nowIso)) as { + __stmt?: PendingStmt; + } + ); + } + + await (orm as unknown as { batch: (s: typeof stmts) => Promise }).batch(stmts); +} + +const config = { deciderRepetitions: REPS }; + +beforeEach(() => { + store.reset(); + vi.mocked(computeEngineIdentity).mockReturnValue(ENGINE); +}); + +// --------------------------------------------------------------------------- +// Currency predicate +// --------------------------------------------------------------------------- + +describe('isCurrentBenchmarkProfile', () => { + const current = { engineIdentity: ENGINE, repetitions: REPS }; + const e = entry('m/a', 'xhigh'); + + it('is current only when engine identity, repetitions, and exact entry match', () => { + expect( + isCurrentBenchmarkProfile( + { + model: 'm/a', + variant: 'xhigh', + engine_identity: ENGINE, + repetitions: REPS, + }, + current, + e + ) + ).toBe(true); + }); + + it('treats an engine-identity change as stale', () => { + expect( + isCurrentBenchmarkProfile( + { + model: 'm/a', + variant: 'xhigh', + engine_identity: 'v-old:engine', + repetitions: REPS, + }, + current, + e + ) + ).toBe(false); + }); + + it('treats a repetitions change as stale', () => { + expect( + isCurrentBenchmarkProfile( + { + model: 'm/a', + variant: 'xhigh', + engine_identity: ENGINE, + repetitions: 3, + }, + current, + e + ) + ).toBe(false); + }); + + it('treats a same-engine/repetitions row of another variant as not current', () => { + expect( + isCurrentBenchmarkProfile( + { + model: 'm/a', + variant: 'max', + engine_identity: ENGINE, + repetitions: REPS, + }, + current, + e + ) + ).toBe(false); + }); + + it('maps null entry variant to empty storage variant', () => { + expect( + isCurrentBenchmarkProfile( + { + model: 'm/a', + variant: '', + engine_identity: ENGINE, + repetitions: REPS, + }, + current, + entry('m/a', null) + ) + ).toBe(true); + }); +}); + +// --------------------------------------------------------------------------- +// Pure admission classification +// --------------------------------------------------------------------------- + +describe('classifyProfileAdmission', () => { + const current = { engineIdentity: ENGINE, repetitions: REPS }; + const e = entry('m'); + + it('reports ready/pending/running without admission', () => { + for (const status of ['ready', 'pending', 'running'] as const) { + expect( + classifyProfileAdmission([profileRow({ model: 'm', status })], current, e, false) + ).toEqual({ kind: 'report', status }); + } + }); + + it('reports failed without charge when not retried', () => { + expect( + classifyProfileAdmission( + [profileRow({ model: 'm', status: 'failed', failure_reason: 'boom' })], + current, + e, + false + ) + ).toEqual({ kind: 'report', status: 'failed', failureReason: 'boom' }); + }); + + it('charges admission for explicit failed retry', () => { + expect( + classifyProfileAdmission( + [profileRow({ model: 'm', status: 'failed', failure_reason: 'boom' })], + current, + e, + true + ) + ).toEqual({ kind: 'admit', charged: true }); + }); + + it('charges for a never-seen pair', () => { + expect(classifyProfileAdmission([], current, e, false)).toEqual({ + kind: 'admit', + charged: true, + }); + }); + + it('free-admits when only a stale (old engine) row exists', () => { + expect( + classifyProfileAdmission( + [ + profileRow({ + model: 'm', + status: 'ready', + engine_identity: 'v-old:engine', + }), + ], + current, + e, + false + ) + ).toEqual({ kind: 'admit', charged: false }); + }); + + it('free-admits when only a stale (old repetitions) row exists', () => { + expect( + classifyProfileAdmission( + [profileRow({ model: 'm', status: 'ready', repetitions: 5 })], + current, + e, + false + ) + ).toEqual({ kind: 'admit', charged: false }); + }); +}); + +describe('computeQuotaRetryAt', () => { + it('is oldest in-window event + 24h', () => { + const oldest = '2026-07-28T01:00:00.000Z'; + expect(computeQuotaRetryAt(oldest)).toBe('2026-07-29T01:00:00.000Z'); + }); +}); + +// --------------------------------------------------------------------------- +// registerProfiles (atomic admission) +// --------------------------------------------------------------------------- + +describe('registerProfiles', () => { + it('admits a missing pair once; a second owner is not charged', async () => { + const entries = [entry('openai/gpt-4o', 'xhigh')]; + + const first = await registerProfiles({} as D1Database, config, { + ownerType: 'user', + ownerId: 'owner-a', + entries, + now: NOW, + }); + expect(first.statuses).toEqual([{ entry: entries[0], status: 'pending', failureReason: null }]); + expect(store.events).toHaveLength(1); + expect(store.events[0]?.owner_id).toBe('owner-a'); + expect(store.profiles.size).toBe(1); + + const second = await registerProfiles({} as D1Database, config, { + ownerType: 'org', + ownerId: 'owner-b', + entries, + now: NOW, + }); + expect(second.statuses).toEqual([ + { entry: entries[0], status: 'pending', failureReason: null }, + ]); + // Second owner sees pending and is not charged. + expect(store.events).toHaveLength(1); + expect(store.profiles.size).toBe(1); + }); + + it('blocks a second concurrent charge when the batch observes a committed pending row', async () => { + const entries = [entry('race/model', 'xhigh')]; + + // First owner commits a charged pending admission. + await registerProfiles({} as D1Database, config, { + ownerType: 'user', + ownerId: 'owner-a', + entries, + now: NOW, + }); + expect(store.events).toHaveLength(1); + const firstRunId = [...store.profiles.values()][0]?.run_id ?? null; + + // Concurrent second batch after a stale pre-read that classified "charge": + // real builders → event guard + failed-only upsert observe committed pending. + await applyProductionAdmissionBatch({ + chargedEvents: [{ ownerType: 'org', ownerId: 'owner-b', entry: entries[0]! }], + registerUpserts: entries, + }); + + expect(store.events).toHaveLength(1); + const row = [...store.profiles.values()].find(r => r.model === 'race/model'); + expect(row?.status).toBe('pending'); + expect(row?.run_id).toBe(firstRunId); + + // Fresh register after concurrent commit reports true pending, no new charge. + const second = await registerProfiles({} as D1Database, config, { + ownerType: 'org', + ownerId: 'owner-b', + entries, + now: NOW, + }); + expect(second.statuses[0]?.status).toBe('pending'); + expect(store.events).toHaveLength(1); + }); + + it('does not charge ready, pending, or running entries', async () => { + store.seedProfile(profileRow({ model: 'a/ready', status: 'ready', variant: '' })); + store.seedProfile(profileRow({ model: 'a/pending', status: 'pending', variant: '' })); + store.seedProfile(profileRow({ model: 'a/running', status: 'running', variant: '' })); + + const result = await registerProfiles({} as D1Database, config, { + ownerType: 'user', + ownerId: 'u1', + entries: [entry('a/ready'), entry('a/pending'), entry('a/running')], + now: NOW, + }); + + expect(result.statuses.map(s => s.status)).toEqual(['ready', 'pending', 'running']); + expect(store.events).toHaveLength(0); + }); + + it('does not regress a running current row when a stale pre-read would admit', async () => { + store.seedProfile( + profileRow({ + model: 'live/running', + status: 'running', + run_id: 'active-run', + requested_at: '2026-07-28T10:00:00.000Z', + }) + ); + + // Stale-read batch via real builders: charged event + failed-only upsert. + await applyProductionAdmissionBatch({ + chargedEvents: [{ ownerType: 'user', ownerId: 'stale-reader', entry: entry('live/running') }], + registerUpserts: [entry('live/running')], + }); + + const row = store.profiles.get( + store.profilePk(profileRow({ model: 'live/running', status: 'running' })) + ); + expect(row?.status).toBe('running'); + expect(row?.run_id).toBe('active-run'); + expect(store.events).toHaveLength(0); + + const result = await registerProfiles({} as D1Database, config, { + ownerType: 'user', + ownerId: 'u1', + entries: [entry('live/running')], + now: NOW, + }); + expect(result.statuses[0]?.status).toBe('running'); + expect(store.events).toHaveLength(0); + }); + + it('does not regress a ready current row when a stale pre-read would admit', async () => { + store.seedProfile( + profileRow({ + model: 'live/ready', + status: 'ready', + run_id: 'done-run', + completed_at: '2026-07-28T09:00:00.000Z', + }) + ); + + await applyProductionAdmissionBatch({ + registerUpserts: [entry('live/ready')], + }); + + const row = store.profiles.get( + store.profilePk(profileRow({ model: 'live/ready', status: 'ready' })) + ); + expect(row?.status).toBe('ready'); + expect(row?.run_id).toBe('done-run'); + expect(row?.completed_at).toBe('2026-07-28T09:00:00.000Z'); + }); + + it('reports failed without charge when not in retryEntries', async () => { + store.seedProfile( + profileRow({ + model: 'a/fail', + status: 'failed', + failure_reason: 'container crash', + }) + ); + + const result = await registerProfiles({} as D1Database, config, { + ownerType: 'user', + ownerId: 'u1', + entries: [entry('a/fail')], + now: NOW, + }); + + expect(result.statuses).toEqual([ + { + entry: entry('a/fail'), + status: 'failed', + failureReason: 'container crash', + }, + ]); + expect(store.events).toHaveLength(0); + expect( + store.profiles.get(store.profilePk(profileRow({ model: 'a/fail', status: 'failed' })))?.status + ).toBe('failed'); + }); + + it('charges and resets failed entry when listed in retryEntries', async () => { + store.seedProfile( + profileRow({ + model: 'a/fail', + status: 'failed', + failure_reason: 'container crash', + run_id: 'old-run', + }) + ); + + const result = await registerProfiles({} as D1Database, config, { + ownerType: 'user', + ownerId: 'u1', + entries: [entry('a/fail')], + retryEntries: [entry('a/fail')], + now: NOW, + }); + + expect(result.statuses[0]?.status).toBe('pending'); + expect(store.events).toHaveLength(1); + const row = [...store.profiles.values()].find(r => r.model === 'a/fail'); + expect(row).toMatchObject({ + status: 'pending', + failure_reason: null, + run_id: null, + engine_identity: ENGINE, + repetitions: REPS, + }); + }); + + it('free-admits stale engine rows without a request event and keeps history', async () => { + const stale = profileRow({ + model: 'stale/m', + status: 'ready', + engine_identity: 'v-old:engine', + completed_at: '2026-01-01T00:00:00.000Z', + }); + store.seedProfile(stale); + + const result = await registerProfiles({} as D1Database, config, { + ownerType: 'user', + ownerId: 'u1', + entries: [entry('stale/m')], + now: NOW, + }); + + expect(result.statuses[0]?.status).toBe('pending'); + expect(store.events).toHaveLength(0); + // Old row preserved. + expect(store.profiles.get(store.profilePk(stale))?.status).toBe('ready'); + // New current pending row inserted. + const current = [...store.profiles.values()].find( + r => r.model === 'stale/m' && r.engine_identity === ENGINE + ); + expect(current?.status).toBe('pending'); + }); + + it('allows 10 charged admissions in 24h and rejects the 11th with nothing written', async () => { + const owner = { ownerType: 'user' as const, ownerId: 'quota-user' }; + // Seed 10 prior charged events in the window. + for (let i = 0; i < PROFILE_ADMISSION_LIMIT; i++) { + store.seedEvent({ + owner_type: owner.ownerType, + owner_id: owner.ownerId, + model: `prior/${i}`, + variant: '', + engine_identity: ENGINE, + repetitions: REPS, + admitted_at: new Date(NOW.getTime() - (PROFILE_ADMISSION_LIMIT - i) * 60_000).toISOString(), + }); + } + const oldest = store.events[0]!.admitted_at; + const expectedRetryAt = new Date( + Date.parse(oldest) + PROFILE_ADMISSION_WINDOW_MS + ).toISOString(); + + const profilesBefore = store.profiles.size; + const eventsBefore = store.events.length; + + await expect( + registerProfiles({} as D1Database, config, { + ...owner, + entries: [entry('over/quota')], + now: NOW, + }) + ).rejects.toBeInstanceOf(ProfileQuotaExceededError); + + try { + await registerProfiles({} as D1Database, config, { + ...owner, + entries: [entry('over/quota')], + now: NOW, + }); + } catch (err) { + const quotaErr = err as ProfileQuotaExceededError; + expect(quotaErr.quota.retryAt).toBe(expectedRetryAt); + expect(quotaErr.quota.error).toContain(expectedRetryAt); + } + + // All-or-nothing: no new events, no profile upsert. + expect(store.events).toHaveLength(eventsBefore); + expect(store.profiles.size).toBe(profilesBefore); + expect([...store.profiles.keys()].some(k => k.includes('over/quota'))).toBe(false); + + // After the retryAt instant, admission succeeds. + const after = new Date(Date.parse(expectedRetryAt) + 1); + const result = await registerProfiles({} as D1Database, config, { + ...owner, + entries: [entry('over/quota')], + now: after, + }); + expect(result.statuses[0]?.status).toBe('pending'); + expect(store.events.some(e => e.model === 'over/quota')).toBe(true); + }); + + it('rejects the whole batch when any charged admission would exceed quota', async () => { + const owner = { ownerType: 'user' as const, ownerId: 'batch-quota' }; + for (let i = 0; i < 9; i++) { + store.seedEvent({ + owner_type: owner.ownerType, + owner_id: owner.ownerId, + model: `prior/${i}`, + variant: '', + engine_identity: ENGINE, + repetitions: REPS, + admitted_at: new Date(NOW.getTime() - 60_000).toISOString(), + }); + } + // 9 existing + 2 new charged = 11 → reject both, write nothing. + await expect( + registerProfiles({} as D1Database, config, { + ...owner, + entries: [entry('new/a'), entry('new/b')], + now: NOW, + }) + ).rejects.toBeInstanceOf(ProfileQuotaExceededError); + + expect(store.profiles.size).toBe(0); + expect(store.events).toHaveLength(9); + }); +}); + +// --------------------------------------------------------------------------- +// lookupProfileStatuses +// --------------------------------------------------------------------------- + +describe('lookupProfileStatuses', () => { + it('returns current statuses without charging', async () => { + store.seedProfile(profileRow({ model: 'a/ready', status: 'ready' })); + store.seedProfile(profileRow({ model: 'a/fail', status: 'failed', failure_reason: 'nope' })); + + const result = await lookupProfileStatuses({} as D1Database, config, { + entries: [entry('a/ready'), entry('a/fail')], + now: NOW, + }); + + expect(result.statuses).toEqual([ + { entry: entry('a/ready'), status: 'ready', failureReason: null }, + { entry: entry('a/fail'), status: 'failed', failureReason: 'nope' }, + ]); + expect(store.events).toHaveLength(0); + }); + + it('inserts a free pending row for stale profiles and reports pending', async () => { + const stale = profileRow({ + model: 'stale/m', + status: 'ready', + engine_identity: 'v-old:engine', + }); + store.seedProfile(stale); + + const result = await lookupProfileStatuses({} as D1Database, config, { + entries: [entry('stale/m')], + now: NOW, + }); + + expect(result.statuses[0]).toEqual({ + entry: entry('stale/m'), + status: 'pending', + failureReason: null, + }); + expect(store.events).toHaveLength(0); + expect(store.profiles.get(store.profilePk(stale))?.status).toBe('ready'); + const current = [...store.profiles.values()].find( + r => r.model === 'stale/m' && r.engine_identity === ENGINE + ); + expect(current?.status).toBe('pending'); + }); + + it('preserves a concurrent current row on stale status-lookup insert race', async () => { + const stale = profileRow({ + model: 'race/stale', + status: 'ready', + engine_identity: 'v-old:engine', + }); + store.seedProfile(stale); + // Concurrent writer already created the current ready row. + store.seedProfile( + profileRow({ + model: 'race/stale', + status: 'ready', + run_id: 'winner-run', + completed_at: '2026-07-28T08:00:00.000Z', + }) + ); + + // Lookup sees current ready → reports it, no charge. + const withCurrent = await lookupProfileStatuses({} as D1Database, config, { + entries: [entry('race/stale')], + now: NOW, + }); + expect(withCurrent.statuses[0]?.status).toBe('ready'); + expect( + store.profiles.get( + store.profilePk(profileRow({ model: 'race/stale', status: 'ready', run_id: 'winner-run' })) + )?.run_id + ).toBe('winner-run'); + + // Explicit do_nothing race via real status builder: must not clobber ready. + await applyProductionAdmissionBatch({ + statusInserts: [entry('race/stale')], + }); + const current = store.profiles.get( + store.profilePk(profileRow({ model: 'race/stale', status: 'ready' })) + ); + expect(current?.status).toBe('ready'); + expect(current?.run_id).toBe('winner-run'); + expect(store.events).toHaveLength(0); + }); + + it('reports missing as failed without admitting', async () => { + const result = await lookupProfileStatuses({} as D1Database, config, { + entries: [entry('never/seen')], + now: NOW, + }); + + expect(result.statuses).toEqual([ + { + entry: entry('never/seen'), + status: 'failed', + failureReason: MISSING_PROFILE_FAILURE_REASON, + }, + ]); + expect(store.profiles.size).toBe(0); + expect(store.events).toHaveLength(0); + }); +}); diff --git a/services/auto-routing-benchmark/src/profiles.ts b/services/auto-routing-benchmark/src/profiles.ts new file mode 100644 index 0000000000..f67e87fe03 --- /dev/null +++ b/services/auto-routing-benchmark/src/profiles.ts @@ -0,0 +1,555 @@ +import type { + BenchmarkConfig, + BenchmarkProfileEntryStatus, + BenchmarkProfileQuotaError, + BenchmarkProfileStatus, + BenchmarkProfileStatusesResponse, + PoolEntry, +} from '@kilocode/auto-routing-contracts'; +import { + BENCHMARK_PROFILE_FAILURE_REASON_MAX_LENGTH, + MAX_POOL_ENTRIES, + poolEntryKey, +} from '@kilocode/auto-routing-contracts'; +import type { BatchItem } from 'drizzle-orm/batch'; +import { and, asc, eq, gte, inArray, sql } from 'drizzle-orm'; +import { drizzle } from 'drizzle-orm/d1'; +import { benchmarkProfiles, profileRequestEvents } from './db-schema'; +import { variantFromStorage, variantToStorage } from './reasoning-effort'; +import { computeEngineIdentity } from './run'; + +/** Rolling 24h owner admission limit (charged new/retry only). */ +export const PROFILE_ADMISSION_LIMIT = 10; +export const PROFILE_ADMISSION_WINDOW_MS = 24 * 60 * 60 * 1000; + +export const MISSING_PROFILE_FAILURE_REASON = 'Benchmark profile missing; retry to request it.'; + +export type ProfileOwnerType = 'user' | 'org'; + +export type ProfileRow = typeof benchmarkProfiles.$inferSelect; + +export type CurrentProfileContext = { + engineIdentity: string; + repetitions: number; +}; + +/** + * Shared currency predicate: a profile row is current only when it was measured + * (or is being measured) under the live decider engine identity and the saved + * deciderRepetitions, for the same exact Pool entry (model + variant). Used by + * status lookup and (later) custom table assembly — never treat a stale row as + * ready. + */ +export function isCurrentBenchmarkProfile( + row: Pick, + current: CurrentProfileContext, + entry: PoolEntry +): boolean { + return ( + row.model === entry.model && + row.variant === variantToStorage(entry.variant) && + row.engine_identity === current.engineIdentity && + row.repetitions === current.repetitions + ); +} + +export function currentProfileContextFromConfig( + config: Pick +): CurrentProfileContext { + return { + engineIdentity: computeEngineIdentity('decider'), + repetitions: config.deciderRepetitions, + }; +} + +export type AdmissionClass = + | { kind: 'report'; status: BenchmarkProfileStatus; failureReason?: string | null } + | { kind: 'admit'; charged: boolean }; + +/** + * Pure classification of one entry against existing registry rows for that + * exact pair. `retry` is true when the owner listed this entry in retryEntries. + */ +export function classifyProfileAdmission( + entryRows: readonly ProfileRow[], + current: CurrentProfileContext, + entry: PoolEntry, + retry: boolean +): AdmissionClass { + const currentRow = entryRows.find(row => isCurrentBenchmarkProfile(row, current, entry)); + if (currentRow) { + if ( + currentRow.status === 'ready' || + currentRow.status === 'pending' || + currentRow.status === 'running' + ) { + return { kind: 'report', status: currentRow.status }; + } + // failed + if (retry) { + return { kind: 'admit', charged: true }; + } + return { + kind: 'report', + status: 'failed', + failureReason: currentRow.failure_reason, + }; + } + + // No current row: stale history (free) or never seen (charged). + if (entryRows.length > 0) { + return { kind: 'admit', charged: false }; + } + return { kind: 'admit', charged: true }; +} + +export function computeQuotaRetryAt( + oldestInWindowAdmittedAt: string, + nowMs: number = Date.now() +): string { + const oldestMs = Date.parse(oldestInWindowAdmittedAt); + if (Number.isNaN(oldestMs)) { + return new Date(nowMs + PROFILE_ADMISSION_WINDOW_MS).toISOString(); + } + return new Date(oldestMs + PROFILE_ADMISSION_WINDOW_MS).toISOString(); +} + +export function boundFailureReason(reason: string | null | undefined): string | null { + if (reason == null) return null; + if (reason.length <= BENCHMARK_PROFILE_FAILURE_REASON_MAX_LENGTH) return reason; + return reason.slice(0, BENCHMARK_PROFILE_FAILURE_REASON_MAX_LENGTH); +} + +function assertUniqueEntries(entries: readonly PoolEntry[]): void { + if (entries.length < 1 || entries.length > MAX_POOL_ENTRIES) { + throw new ProfileValidationError( + `entries must contain between 1 and ${MAX_POOL_ENTRIES} unique pool entries` + ); + } + const seen = new Set(); + for (const entry of entries) { + const key = poolEntryKey(entry); + if (seen.has(key)) { + throw new ProfileValidationError(`Duplicate pool entry: ${key}`); + } + seen.add(key); + } +} + +export class ProfileValidationError extends Error { + constructor(message: string) { + super(message); + this.name = 'ProfileValidationError'; + } +} + +export class ProfileQuotaExceededError extends Error { + readonly quota: BenchmarkProfileQuotaError; + + constructor(quota: BenchmarkProfileQuotaError) { + super(quota.error); + this.name = 'ProfileQuotaExceededError'; + this.quota = quota; + } +} + +export class ProfileConfigMissingError extends Error { + constructor( + message = 'benchmark config not set: save it in the admin panel before registering profiles' + ) { + super(message); + this.name = 'ProfileConfigMissingError'; + } +} + +function entryKeyFromRow(row: Pick): string { + return poolEntryKey({ model: row.model, variant: variantFromStorage(row.variant) }); +} + +async function loadRowsForEntries( + db: D1Database, + entries: readonly PoolEntry[] +): Promise { + const models = [...new Set(entries.map(e => e.model))]; + if (models.length === 0) return []; + const orm = drizzle(db); + const rows = await orm + .select() + .from(benchmarkProfiles) + .where(inArray(benchmarkProfiles.model, models)); + const wanted = new Set(entries.map(poolEntryKey)); + return rows.filter(row => wanted.has(entryKeyFromRow(row))); +} + +function groupRowsByEntry(rows: readonly ProfileRow[]): Map { + const map = new Map(); + for (const row of rows) { + const key = entryKeyFromRow(row); + const list = map.get(key); + if (list) list.push(row); + else map.set(key, [row]); + } + return map; +} + +function statusFromRow(entry: PoolEntry, row: ProfileRow): BenchmarkProfileEntryStatus { + return { + entry, + status: row.status, + failureReason: row.status === 'failed' ? boundFailureReason(row.failure_reason) : null, + }; +} + +function pendingStatus(entry: PoolEntry): BenchmarkProfileEntryStatus { + return { entry, status: 'pending', failureReason: null }; +} + +function missingStatus(entry: PoolEntry): BenchmarkProfileEntryStatus { + return { + entry, + status: 'failed', + failureReason: MISSING_PROFILE_FAILURE_REASON, + }; +} + +export function buildPendingUpsertValues( + entry: PoolEntry, + current: CurrentProfileContext, + nowIso: string +): typeof benchmarkProfiles.$inferInsert { + return { + model: entry.model, + variant: variantToStorage(entry.variant), + engine_identity: current.engineIdentity, + repetitions: current.repetitions, + status: 'pending', + run_id: null, + failure_reason: null, + requested_at: nowIso, + updated_at: nowIso, + completed_at: null, + }; +} + +const PROFILE_PK_TARGET = [ + benchmarkProfiles.model, + benchmarkProfiles.variant, + benchmarkProfiles.engine_identity, + benchmarkProfiles.repetitions, +] as const; + +/** + * Register-path upsert: insert a pending current row, or only transition an + * existing current row from `failed` → `pending`. Never regresses + * pending/running/ready or clears their provenance. + */ +export function pendingRegisterUpsertStatement( + orm: ReturnType, + values: typeof benchmarkProfiles.$inferInsert +) { + return orm + .insert(benchmarkProfiles) + .values(values) + .onConflictDoUpdate({ + target: [...PROFILE_PK_TARGET], + set: { + status: 'pending', + run_id: null, + failure_reason: null, + requested_at: values.requested_at, + updated_at: values.updated_at, + completed_at: null, + }, + // SQLite ON CONFLICT DO UPDATE WHERE — only rewrite failed rows. + where: eq(benchmarkProfiles.status, 'failed'), + }); +} + +/** + * Status-lookup free re-admission: insert if absent; on conflict with a + * concurrently created current row, do nothing (never clobber). + */ +export function pendingStatusInsertStatement( + orm: ReturnType, + values: typeof benchmarkProfiles.$inferInsert +) { + return orm + .insert(benchmarkProfiles) + .values(values) + .onConflictDoNothing({ + target: [...PROFILE_PK_TARGET], + }); +} + +export type ChargedEventInsertValues = { + owner_type: string; + owner_id: string; + model: string; + variant: string; + engine_identity: string; + repetitions: number; + admitted_at: string; +}; + +/** + * Charged ledger insert guarded inside the batch: only write an event when no + * active (pending/running/ready) current profile row exists yet. Statement + * order places these BEFORE profile upserts so the guard sees pre-batch state + * within the transaction; SQLite writer serialization makes a concurrent + * batch's guard observe the first batch's committed row. + * + * SQL shape: INSERT ... SELECT NULL AS id, WHERE NOT EXISTS ( + * SELECT 1 FROM benchmark_profiles WHERE pk AND status IN (...)) + * + * Drizzle's insert().select always emits the full table column list including + * autoincrement `id`, so the SELECT must provide 8 expressions (NULL AS id) + * or SQLite/D1 rejects with "7 values for 8 columns". + */ +export function chargedEventInsertStatement( + orm: ReturnType, + values: ChargedEventInsertValues +) { + // __stmt tags the batch item for the in-memory test stand-in; real D1 + // executes the SELECT...WHERE NOT EXISTS guard. Literal SELECT (no FROM) + // is expressed as SQL because drizzle's typed select builder requires .from(). + return Object.assign( + orm.insert(profileRequestEvents).select(sql` + SELECT + NULL AS id, + ${values.owner_type} AS owner_type, + ${values.owner_id} AS owner_id, + ${values.model} AS model, + ${values.variant} AS variant, + ${values.engine_identity} AS engine_identity, + ${values.repetitions} AS repetitions, + ${values.admitted_at} AS admitted_at + WHERE NOT EXISTS ( + SELECT 1 FROM ${benchmarkProfiles} + WHERE ${benchmarkProfiles.model} = ${values.model} + AND ${benchmarkProfiles.variant} = ${values.variant} + AND ${benchmarkProfiles.engine_identity} = ${values.engine_identity} + AND ${benchmarkProfiles.repetitions} = ${values.repetitions} + AND ${benchmarkProfiles.status} IN ('pending', 'running', 'ready') + ) + `), + { + __stmt: { + kind: 'event' as const, + guarded: true as const, + values: { + owner_type: values.owner_type, + owner_id: values.owner_id, + model: values.model, + variant: values.variant, + engine_identity: values.engine_identity, + repetitions: values.repetitions, + admitted_at: values.admitted_at, + }, + }, + } + ); +} + +export type RegisterProfilesInput = { + ownerType: ProfileOwnerType; + ownerId: string; + entries: readonly PoolEntry[]; + retryEntries?: readonly PoolEntry[]; + /** Injectable clock for tests. */ + now?: Date; +}; + +/** + * Atomically admit missing/stale/retried Benchmark profiles for an owner. + * All-or-nothing: either every charged admission + upsert commits in one + * `db.batch`, or a 429 quota result writes nothing. + */ +export async function registerProfiles( + db: D1Database, + config: Pick, + input: RegisterProfilesInput +): Promise { + assertUniqueEntries(input.entries); + const current = currentProfileContextFromConfig(config); + const now = input.now ?? new Date(); + const nowIso = now.toISOString(); + const windowStartIso = new Date(now.getTime() - PROFILE_ADMISSION_WINDOW_MS).toISOString(); + + const retryKeys = new Set((input.retryEntries ?? []).map(poolEntryKey)); + const existingRows = await loadRowsForEntries(db, input.entries); + const byEntry = groupRowsByEntry(existingRows); + + type PlanItem = { + entry: PoolEntry; + classification: AdmissionClass; + }; + const plan: PlanItem[] = input.entries.map(entry => ({ + entry, + classification: classifyProfileAdmission( + byEntry.get(poolEntryKey(entry)) ?? [], + current, + entry, + retryKeys.has(poolEntryKey(entry)) + ), + })); + + const chargedAdmissions = plan.filter( + p => p.classification.kind === 'admit' && p.classification.charged + ); + const freeAdmissions = plan.filter( + p => p.classification.kind === 'admit' && !p.classification.charged + ); + + const orm = drizzle(db); + const recentEvents = await orm + .select() + .from(profileRequestEvents) + .where( + and( + eq(profileRequestEvents.owner_type, input.ownerType), + eq(profileRequestEvents.owner_id, input.ownerId), + gte(profileRequestEvents.admitted_at, windowStartIso) + ) + ) + .orderBy(asc(profileRequestEvents.admitted_at)); + + if (recentEvents.length + chargedAdmissions.length > PROFILE_ADMISSION_LIMIT) { + const oldest = recentEvents[0]?.admitted_at ?? nowIso; + const retryAt = computeQuotaRetryAt(oldest, now.getTime()); + throw new ProfileQuotaExceededError({ + error: `Profile benchmark request limit reached. New benchmarks can be requested after ${retryAt}.`, + retryAt, + }); + } + + const admissions = [...chargedAdmissions, ...freeAdmissions]; + if (admissions.length > 0) { + const stmts: BatchItem<'sqlite'>[] = []; + + // Events first so the NOT EXISTS guard sees pre-batch profile state. + for (const item of chargedAdmissions) { + stmts.push( + chargedEventInsertStatement(orm, { + owner_type: input.ownerType, + owner_id: input.ownerId, + model: item.entry.model, + variant: variantToStorage(item.entry.variant), + engine_identity: current.engineIdentity, + repetitions: current.repetitions, + admitted_at: nowIso, + }) as BatchItem<'sqlite'> + ); + } + + for (const item of admissions) { + stmts.push( + pendingRegisterUpsertStatement( + orm, + buildPendingUpsertValues(item.entry, current, nowIso) + ) as BatchItem<'sqlite'> + ); + } + + // D1 batch requires a non-empty tuple; admissions.length > 0 guarantees it. + await orm.batch(stmts as [BatchItem<'sqlite'>, ...BatchItem<'sqlite'>[]]); + } + + // Derive reported statuses for admitted entries from post-batch rows so a + // guard-blocked concurrent admit reports the true current status. + const postRows = + admissions.length > 0 + ? await loadRowsForEntries( + db, + admissions.map(a => a.entry) + ) + : []; + const postByEntry = groupRowsByEntry(postRows); + + const statuses: BenchmarkProfileEntryStatus[] = plan.map(({ entry, classification }) => { + if (classification.kind === 'admit') { + const currentRow = (postByEntry.get(poolEntryKey(entry)) ?? []).find(row => + isCurrentBenchmarkProfile(row, current, entry) + ); + if (currentRow) return statusFromRow(entry, currentRow); + return pendingStatus(entry); + } + return { + entry, + status: classification.status, + failureReason: + classification.status === 'failed' + ? boundFailureReason(classification.failureReason) + : null, + }; + }); + + return { statuses }; +} + +export type LookupProfileStatusesInput = { + entries: readonly PoolEntry[]; + /** Injectable clock for tests. */ + now?: Date; +}; + +/** + * Status lookup for up to MAX_POOL_ENTRIES exact pairs. + * - current row → its status + * - stale only → free pending insert (no request event; on-conflict do nothing) + * + report the actual post-batch current status + * - missing → failed with missing message; never admits from a read path + */ +export async function lookupProfileStatuses( + db: D1Database, + config: Pick, + input: LookupProfileStatusesInput +): Promise { + assertUniqueEntries(input.entries); + const current = currentProfileContextFromConfig(config); + const nowIso = (input.now ?? new Date()).toISOString(); + + const existingRows = await loadRowsForEntries(db, input.entries); + const byEntry = groupRowsByEntry(existingRows); + + const staleToAdmit: PoolEntry[] = []; + /** Parallel to input.entries: pre-resolved status, or null when stale-admit pending. */ + const preStatuses: Array = []; + + for (const entry of input.entries) { + const rows = byEntry.get(poolEntryKey(entry)) ?? []; + const currentRow = rows.find(row => isCurrentBenchmarkProfile(row, current, entry)); + if (currentRow) { + preStatuses.push(statusFromRow(entry, currentRow)); + continue; + } + if (rows.length > 0) { + staleToAdmit.push(entry); + preStatuses.push(null); + continue; + } + preStatuses.push(missingStatus(entry)); + } + + if (staleToAdmit.length > 0) { + const orm = drizzle(db); + const stmts = staleToAdmit.map(entry => + pendingStatusInsertStatement(orm, buildPendingUpsertValues(entry, current, nowIso)) + ); + await orm.batch(stmts as unknown as [BatchItem<'sqlite'>, ...BatchItem<'sqlite'>[]]); + } + + const postRows = staleToAdmit.length > 0 ? await loadRowsForEntries(db, staleToAdmit) : []; + const postByEntry = groupRowsByEntry(postRows); + + const statuses: BenchmarkProfileEntryStatus[] = input.entries.map((entry, i) => { + const pre = preStatuses[i]; + if (pre) return pre; + const currentRow = (postByEntry.get(poolEntryKey(entry)) ?? []).find(row => + isCurrentBenchmarkProfile(row, current, entry) + ); + if (currentRow) return statusFromRow(entry, currentRow); + return pendingStatus(entry); + }); + + return { statuses }; +} diff --git a/services/auto-routing-benchmark/src/reasoning-effort.ts b/services/auto-routing-benchmark/src/reasoning-effort.ts index 05cee319e3..485e9c67dc 100644 --- a/services/auto-routing-benchmark/src/reasoning-effort.ts +++ b/services/auto-routing-benchmark/src/reasoning-effort.ts @@ -8,3 +8,22 @@ export function parsePersistedReasoningEffort(value: string | null): ReasoningEf const parsed = ReasoningEffortSchema.safeParse(value); return parsed.success ? parsed.data : null; } + +/** D1 stores '' for the null/default variant; application code uses null. */ +export function variantToStorage(variant: string | null | undefined): string { + return variant ?? ''; +} + +/** Convert a D1-stored variant ('' = null) back to the application null form. */ +export function variantFromStorage(stored: string | null | undefined): string | null { + if (stored == null || stored === '') return null; + return stored; +} + +/** + * Platform runs still select one reasoningEffort per model. Map that effort + * key to the canonical stored variant value (today's CLI --variant value). + */ +export function variantFromReasoningEffort(effort: string | null | undefined): string | null { + return effort ?? null; +} diff --git a/services/auto-routing-benchmark/src/routing-table-builder.test.ts b/services/auto-routing-benchmark/src/routing-table-builder.test.ts index 4bb0643e9e..0716ebece2 100644 --- a/services/auto-routing-benchmark/src/routing-table-builder.test.ts +++ b/services/auto-routing-benchmark/src/routing-table-builder.test.ts @@ -107,6 +107,126 @@ describe('buildRoutingTable', () => { expect(cheap?.reasoningEffort).toBeNull(); }); + it('platform table JSON shape has reasoningEffort and no variant key', () => { + const table = buildRoutingTable({ + runId: 'test-run-platform-shape', + generatedAt: '2026-01-01T00:00:00.000Z', + minAccuracy: 0.7, + switchCostFactor: 3, + bestAccuracySwitchThreshold: 0.05, + deciderModels: DECIDER_MODELS, + summaries: summariesForEveryRoute(), + }); + const cand = table.routes['implementation/code_generation']?.[0]; + expect(cand).toBeDefined(); + expect(cand).toHaveProperty('reasoningEffort'); + expect(Object.keys(cand ?? {})).not.toContain('variant'); + // Serialize/parse as published artifact would. + const json = JSON.parse(JSON.stringify(table)) as typeof table; + const jCand = json.routes['implementation/code_generation']?.[0]; + expect(jCand).toHaveProperty('reasoningEffort'); + expect(jCand && 'variant' in jCand).toBe(false); + }); + + it('two variants of one model appear as distinct candidates with matched efforts', () => { + const routeKey = 'implementation/code_generation' as const; + const table = buildRoutingTable({ + runId: 'test-run-two-variants', + generatedAt: '2026-01-01T00:00:00.000Z', + minAccuracy: 0.5, + switchCostFactor: 3, + bestAccuracySwitchThreshold: 0.05, + deciderModels: [ + { id: 'model/a', reasoningEffort: 'high' }, + { id: 'model/a', reasoningEffort: 'low' }, + ], + summaries: TAXONOMY_ROUTE_KEYS.flatMap(rk => { + if (rk !== routeKey) { + // Non-focus routes need both pairs so no route is empty and exact + // (model, variant) matching still binds each summary to its effort. + return [ + { ...summary('model/a', rk, 0.8, 0.001), variant: 'high' }, + { ...summary('model/a', rk, 0.75, 0.0015), variant: 'low' }, + ]; + } + return [ + { ...summary('model/a', rk, 0.9, 0.002), variant: 'high' }, + { ...summary('model/a', rk, 0.7, 0.001), variant: 'low' }, + ]; + }), + }); + const cands = table.routes[routeKey] ?? []; + expect(cands).toHaveLength(2); + expect(cands.map(c => c.model)).toEqual(['model/a', 'model/a']); + // Per-candidate pairing: accuracy 0.9 was measured at high, 0.7 at low. + // A swapping matcher that only checks the effort set would still fail here. + const highAcc = cands.find(c => c.accuracy === 0.9); + const lowAcc = cands.find(c => c.accuracy === 0.7); + expect(highAcc?.reasoningEffort).toBe('high'); + expect(lowAcc?.reasoningEffort).toBe('low'); + for (const c of cands) { + expect(Object.keys(c)).not.toContain('variant'); + } + }); + + it('binds reasoningEffort from an exact (model, variant) snapshot match', () => { + const routeKey = 'implementation/code_generation' as const; + const table = buildRoutingTable({ + runId: 'test-run-exact-pair', + generatedAt: '2026-01-01T00:00:00.000Z', + minAccuracy: 0.5, + switchCostFactor: 3, + bestAccuracySwitchThreshold: 0.05, + deciderModels: [ + { id: 'model/a', reasoningEffort: 'high' }, + { id: 'model/a', reasoningEffort: 'low' }, + { id: 'model/b', reasoningEffort: 'medium' }, + ], + summaries: TAXONOMY_ROUTE_KEYS.flatMap(rk => [ + { ...summary('model/a', rk, 0.9, 0.002), variant: 'high' }, + { ...summary('model/b', rk, 0.8, 0.003), variant: 'medium' }, + ]), + }); + const a = table.routes[routeKey]?.find(c => c.model === 'model/a'); + const b = table.routes[routeKey]?.find(c => c.model === 'model/b'); + expect(a?.reasoningEffort).toBe('high'); + expect(b?.reasoningEffort).toBe('medium'); + }); + + it('legacy single-row snapshot binds when summary omits variant', () => { + const table = buildRoutingTable({ + runId: 'test-run-legacy-single', + generatedAt: '2026-01-01T00:00:00.000Z', + minAccuracy: 0.7, + switchCostFactor: 3, + bestAccuracySwitchThreshold: 0.05, + deciderModels: DECIDER_MODELS, + summaries: summariesForEveryRoute(), + }); + const value = table.routes['implementation/code_generation']?.find( + c => c.model === 'model/value' + ); + expect(value?.reasoningEffort).toBe('medium'); + }); + + it('throws when multiple snapshot rows exist and none matches the summary pair', () => { + expect(() => + buildRoutingTable({ + runId: 'test-run-ambiguous-snapshot', + generatedAt: '2026-01-01T00:00:00.000Z', + minAccuracy: 0.5, + switchCostFactor: 3, + bestAccuracySwitchThreshold: 0.05, + deciderModels: [ + { id: 'model/a', reasoningEffort: 'high' }, + { id: 'model/a', reasoningEffort: 'low' }, + ], + // Summary has no variant and two snapshot rows → cannot bind safely. + summaries: TAXONOMY_ROUTE_KEYS.map(rk => summary('model/a', rk, 0.9, 0.002)), + }) + ).toThrow(/no snapshot row for model model\/a/); + }); + it('throws when any taxonomy route has no candidates', () => { expect(() => buildRoutingTable({ diff --git a/services/auto-routing-benchmark/src/routing-table-builder.ts b/services/auto-routing-benchmark/src/routing-table-builder.ts index 50df25ecbc..f2747a6194 100644 --- a/services/auto-routing-benchmark/src/routing-table-builder.ts +++ b/services/auto-routing-benchmark/src/routing-table-builder.ts @@ -1,12 +1,17 @@ import { + CustomRoutingTableSchema, + poolEntryKey, rankCandidates, RoutingTableSchema, TAXONOMY_ROUTE_KEYS, type BenchmarkDeciderModel, type BenchmarkModelSummary, + type CustomRoutingTable, + type PoolEntry, type RoutingTable, type TaxonomyRouteKey, } from '@kilocode/auto-routing-contracts'; +import type { BenchmarkModelSummaryWithRun } from './db'; // Builds the routing table from per-(model, taxonomy-route) decider summaries. Models // with zero graded cases in a route are excluded from that route, as are @@ -33,18 +38,51 @@ export function buildRoutingTable(params: { deciderModels, summaries, } = params; - const modelConfigById = new Map(deciderModels.map(m => [m.id, m] as const)); + // Prefer exact (model, variant) match so two variants of one model keep + // distinct snapshot rows. Legacy / platform: when the snapshot has exactly + // one row for the model (one effort per model), bind that row even if the + // summary omitted variant. Multiple snapshot rows without an exact match is + // corrupt — throw so buildRoutingTable fails and the caller keeps the + // previous published table. Emit reasoningEffort only — never variant — so + // the published PLATFORM artifact shape stays unchanged for rolling deploys. + const snapshotVariant = (m: BenchmarkDeciderModel): string | null => + m.variant !== undefined ? (m.variant ?? null) : (m.reasoningEffort ?? null); + + const findSnapshot = ( + model: string, + variant: string | null | undefined + ): BenchmarkDeciderModel => { + const appVariant = variant ?? null; + const exact = deciderModels.find(m => m.id === model && snapshotVariant(m) === appVariant); + if (exact) return exact; + const forModel = deciderModels.filter(m => m.id === model); + const sole = forModel.length === 1 ? forModel[0] : undefined; + // Exactly one snapshot row for this model → platform/legacy shape. + if (sole) return sole; + throw new Error( + `no snapshot row for model ${model} variant ${JSON.stringify(appVariant)} (${forModel.length} rows for model)` + ); + }; const routeCandidates = (routeKey: TaxonomyRouteKey) => rankCandidates( summaries .filter(s => s.routeKey === routeKey && s.cases > 0 && s.avgCostUsd !== null) - .map(s => ({ - model: s.model, - accuracy: s.accuracy, - avgCostUsd: s.avgCostUsd ?? 0, - reasoningEffort: modelConfigById.get(s.model)?.reasoningEffort ?? null, - })), + .map(s => { + const cfg = findSnapshot(s.model, s.variant); + // Platform artifact: emit reasoningEffort from the run snapshot's + // effort key, NOT variant (custom sparse tables are a later slice). + const effort = + cfg?.reasoningEffort !== undefined && cfg.reasoningEffort !== null + ? cfg.reasoningEffort + : null; + return { + model: s.model, + accuracy: s.accuracy, + avgCostUsd: s.avgCostUsd ?? 0, + reasoningEffort: effort, + }; + }), minAccuracy ); @@ -67,3 +105,110 @@ export function buildRoutingTable(params: { // live table intact. return RoutingTableSchema.parse(table); } + +function fnv1aHex(input: string): string { + let hash = 0x811c9dc5; + for (let i = 0; i < input.length; i++) { + hash ^= input.charCodeAt(i); + hash = Math.imul(hash, 0x01000193); + } + return (hash >>> 0).toString(16).padStart(8, '0'); +} + +/** + * Deterministic version id for a sparse custom table: hash of contributing + * (runId, model, variant) triples so identical assembly inputs cache-hit. + */ +export function computeCustomRoutingTableVersion( + contributors: readonly { runId: string; model: string; variant: string | null }[] +): string { + const parts = [...contributors].map(c => `${c.runId}\0${c.model}\0${c.variant ?? ''}`).sort(); + return `custom-${fnv1aHex(JSON.stringify(parts))}`; +} + +/** + * Assemble a SPARSE custom routing table for the requested ready/current Pool + * entries only. Uses per-route summaries from each profile's provenance run + * only — exact (model, variant) AND run_id must match; no cross-run leakage. + * Omits route keys with no graded candidates; returns null when no requested + * entry contributes. Candidates carry exact `variant` (never reasoningEffort). + */ +export function buildCustomRoutingTable(params: { + generatedAt: string; + minAccuracy: number; + switchCostFactor: number; + bestAccuracySwitchThreshold: number; + /** Ready current profiles with provenance run ids. */ + readyEntries: readonly { + entry: PoolEntry; + runId: string; + }[]; + /** + * Summaries from provenance runs, each carrying its measuring `runId`. + * Assembly selects only rows whose pair AND runId match the ready entry. + */ + summaries: readonly BenchmarkModelSummaryWithRun[]; +}): CustomRoutingTable | null { + const { + generatedAt, + minAccuracy, + switchCostFactor, + bestAccuracySwitchThreshold, + readyEntries, + summaries, + } = params; + + if (readyEntries.length === 0) return null; + + // Per ready entry: only summaries from that entry's provenance run + exact pair. + const provenanceByPair = new Map( + readyEntries.map(r => [poolEntryKey(r.entry), r.runId] as const) + ); + + const filtered = summaries.filter(s => { + const pairKey = poolEntryKey({ model: s.model, variant: s.variant ?? null }); + const provenanceRunId = provenanceByPair.get(pairKey); + return provenanceRunId !== undefined && s.runId === provenanceRunId; + }); + + const routes: CustomRoutingTable['routes'] = {}; + for (const routeKey of TAXONOMY_ROUTE_KEYS) { + const candidates = rankCandidates( + filtered + .filter(s => s.routeKey === routeKey && s.cases > 0 && s.avgCostUsd !== null) + .map(s => ({ + model: s.model, + accuracy: s.accuracy, + avgCostUsd: s.avgCostUsd ?? 0, + // Exact-pair identity only — never emit reasoningEffort on custom tables. + variant: s.variant ?? null, + })), + minAccuracy + ); + if (candidates.length > 0) { + routes[routeKey] = candidates; + } + } + + if (Object.keys(routes).length === 0) return null; + + const version = computeCustomRoutingTableVersion( + readyEntries.map(r => ({ + runId: r.runId, + model: r.entry.model, + variant: r.entry.variant, + })) + ); + + const table: CustomRoutingTable = { + version, + generatedAt, + minAccuracy, + switchCostFactor, + bestAccuracySwitchThreshold, + source: 'benchmark', + routes, + }; + + return CustomRoutingTableSchema.parse(table); +} diff --git a/services/auto-routing-benchmark/src/run-process-job.test.ts b/services/auto-routing-benchmark/src/run-process-job.test.ts index 8e9ecd00ed..7fc2df3050 100644 --- a/services/auto-routing-benchmark/src/run-process-job.test.ts +++ b/services/auto-routing-benchmark/src/run-process-job.test.ts @@ -17,6 +17,14 @@ vi.mock('./db', async importOriginal => { replaceModelSummaries: vi.fn(), saveRoutingTable: vi.fn(), upsertCaseResult: vi.fn(), + listPendingCurrentProfiles: vi.fn(), + listStaleRunningDeciderRunIds: vi.fn(), + markProfilesFailedForRun: vi.fn(), + markProfilesReadyForRun: vi.fn(), + markProfilesRunningForRun: vi.fn(), + markStaleRunsFailed: vi.fn(), + getRunningRun: vi.fn(), + getLatestSummariesByModel: vi.fn(), }; }); @@ -69,7 +77,14 @@ const env = { AUTO_ROUTING_CONFIG: { delete: vi.fn() }, } as unknown as Env; -function mockRunSnapshot(): void { +function mockRunSnapshot( + models: Array<{ + model: string; + variant: string; + enqueued?: boolean; + reasoning_effort?: string | null; + }> = [{ model, variant: '', enqueued: true, reasoning_effort: null }] +): void { vi.mocked(getRunWithModels).mockResolvedValue({ run: { max_concurrency: 4, @@ -81,19 +96,25 @@ function mockRunSnapshot(): void { repetitions: 1, classifier_max_p95_latency_ms: null, started_at: '2026-06-16T00:00:00.000Z', + purpose: 'platform', }, - models: [{ model, enqueued: true, reasoning_effort: null }], + models: models.map(m => ({ + enqueued: true, + reasoning_effort: null, + ...m, + })), } as never); } -function deciderMessage() { +function deciderMessage(overrides: { variant?: string | null } = {}) { return { runId, - kind: 'decider', + kind: 'decider' as const, model, caseIds: [benchCase.id], chunk: 0, rep: 0, + ...overrides, }; } @@ -119,6 +140,98 @@ afterEach(() => { vi.unstubAllGlobals(); }); +describe('processJob — decider exact-pair snapshot resolution', () => { + it('throws when an explicit-variant message has no exact snapshot row', async () => { + mockRunSnapshot([ + { model, variant: 'high', reasoning_effort: 'high' }, + { model, variant: 'low', reasoning_effort: 'low' }, + ]); + + await expect(processJob(env, deciderMessage({ variant: 'medium' }))).rejects.toThrow( + /no snapshot row for model/ + ); + + expect(runDeciderCaseViaCli).not.toHaveBeenCalled(); + expect(upsertCaseResult).not.toHaveBeenCalled(); + }); + + it('resolves the exact snapshot row for an explicit (model, variant) pair', async () => { + mockRunSnapshot([ + { model, variant: 'high', reasoning_effort: 'high' }, + { model, variant: 'low', reasoning_effort: 'low' }, + ]); + + await processJob(env, deciderMessage({ variant: 'low' })); + + expect(runDeciderCaseViaCli).toHaveBeenCalledWith( + env, + expect.objectContaining({ + model, + variant: 'low', + instanceName: `${runId}:${model}:low:0:0`, + }) + ); + expect(upsertCaseResult).toHaveBeenCalledWith( + env.BENCH_DB, + expect.objectContaining({ + model, + variant: 'low', + }) + ); + }); + + it('legacy no-variant message resolves the unique snapshot row for the model', async () => { + mockRunSnapshot([{ model, variant: 'high', reasoning_effort: 'high' }]); + + // Omit variant entirely (legacy pre-deploy message shape). + await processJob(env, { + runId, + kind: 'decider', + model, + caseIds: [benchCase.id], + chunk: 0, + rep: 0, + }); + + expect(runDeciderCaseViaCli).toHaveBeenCalledWith( + env, + expect.objectContaining({ + model, + variant: 'high', + instanceName: `${runId}:${model}:high:0:0`, + }) + ); + expect(upsertCaseResult).toHaveBeenCalledWith( + env.BENCH_DB, + expect.objectContaining({ + model, + variant: 'high', + }) + ); + }); + + it('throws when a legacy no-variant message has multiple snapshot rows for the model', async () => { + mockRunSnapshot([ + { model, variant: 'high', reasoning_effort: 'high' }, + { model, variant: 'low', reasoning_effort: 'low' }, + ]); + + await expect( + processJob(env, { + runId, + kind: 'decider', + model, + caseIds: [benchCase.id], + chunk: 0, + rep: 0, + }) + ).rejects.toThrow(/requires exactly one snapshot row/); + + expect(runDeciderCaseViaCli).not.toHaveBeenCalled(); + expect(upsertCaseResult).not.toHaveBeenCalled(); + }); +}); + describe('processJob — decider container availability failures', () => { it.each([ 'container /run failed: HTTP 503 There is no Container instance available at this time. This is likely because you have reached your max concurrent instance count.', @@ -155,10 +268,12 @@ describe('processJob — decider chunk chaining', () => { await processJob(env, message); + // Instance names include the stored variant segment ('' → empty between colons). + const instanceName = `${runId}:${model}::0:0`; expect(warmUpCliContainer).toHaveBeenCalledWith( env, expect.objectContaining({ - instanceName: `${runId}:${model}:0:0`, + instanceName, kiloApiUrl: 'http://host.docker.internal:3000', orgId: 'benchmark-org', }) @@ -166,7 +281,7 @@ describe('processJob — decider chunk chaining', () => { expect(runDeciderCaseViaCli).toHaveBeenCalledWith( env, expect.objectContaining({ - instanceName: `${runId}:${model}:0:0`, + instanceName, kiloApiUrl: 'http://host.docker.internal:3000', orgId: 'benchmark-org', }) @@ -177,6 +292,7 @@ describe('processJob — decider chunk chaining', () => { runId, kind: 'decider', model, + variant: null, chunk: 1, shard: 0, shardCount: 1, @@ -206,7 +322,7 @@ describe('processJob — decider chunk chaining', () => { expect(warmUpCliContainer).toHaveBeenCalledWith( env, - expect.objectContaining({ instanceName: `${runId}:${model}:0:2` }) + expect.objectContaining({ instanceName: `${runId}:${model}::0:2` }) ); expect(queueSendBatch).toHaveBeenCalledWith([ { @@ -214,6 +330,7 @@ describe('processJob — decider chunk chaining', () => { runId, kind: 'decider', model, + variant: null, chunk: nextChunk, shard, shardCount, @@ -258,6 +375,7 @@ describe('processJob — decider chunk chaining', () => { runId, kind: 'decider', model, + variant: null, chunk: 1, shard: 0, shardCount: 1, @@ -282,7 +400,7 @@ describe('processJob — decider chunk chaining', () => { expect(queueSendBatch).not.toHaveBeenCalled(); expect(destroyDeciderCliContainer).toHaveBeenCalledWith(env, { - instanceName: `${runId}:${model}:0:3`, + instanceName: `${runId}:${model}::0:3`, }); expect(countCaseResults).toHaveBeenCalled(); }); @@ -302,7 +420,7 @@ describe('processJob — decider chunk chaining', () => { }); expect(destroyDeciderCliContainer).toHaveBeenCalledWith(env, { - instanceName: `${runId}:${model}:0:3`, + instanceName: `${runId}:${model}::0:3`, }); expect(warn).toHaveBeenCalledWith( expect.stringContaining('benchmark_container_destroy_failed') diff --git a/services/auto-routing-benchmark/src/run.test.ts b/services/auto-routing-benchmark/src/run.test.ts index 4c38613658..4fbefca11e 100644 --- a/services/auto-routing-benchmark/src/run.test.ts +++ b/services/auto-routing-benchmark/src/run.test.ts @@ -17,6 +17,7 @@ function makeRow(overrides: Partial = {}): CaseResultRow { return { run_id: 'run-1', model: 'model/a', + variant: '', case_id: 'case-1', route_key: null, score: 1, @@ -124,6 +125,41 @@ describe('summarize — classifier kind', () => { }); }); +describe('summarize — exact-pair identity', () => { + it('groups two variants of one model separately', () => { + const rows: CaseResultRow[] = [ + makeRow({ + model: 'model/a', + variant: 'xhigh', + case_id: 'c1', + route_key: 'implementation/code_generation', + score: 1, + }), + makeRow({ + model: 'model/a', + variant: 'max', + case_id: 'c1', + route_key: 'implementation/code_generation', + score: 0, + }), + ]; + const summaries = summarize(rows, 'decider'); + expect(summaries).toHaveLength(2); + const xhigh = summaries.find(s => s.variant === 'xhigh'); + const max = summaries.find(s => s.variant === 'max'); + expect(xhigh?.accuracy).toBe(1); + expect(max?.accuracy).toBe(0); + }); + + it('emits null variant when stored as empty string', () => { + const [s] = summarize( + [makeRow({ variant: '', route_key: 'implementation/code_generation' })], + 'decider' + ); + expect(s.variant).toBeNull(); + }); +}); + describe('summarize — decider kind', () => { it('groups by taxonomy route key', () => { const rows: CaseResultRow[] = [ @@ -540,15 +576,105 @@ describe('decider message fan-out', () => { }); it('getDeciderContainerInstanceName reuses one container per model repetition shard', () => { - const base = { runId: 'run-test', kind: 'decider' as const, model: 'model/a', rep: 2 }; + const base = { + runId: 'run-test', + kind: 'decider' as const, + model: 'model/a', + variant: null as string | null, + rep: 2, + }; expect(getDeciderContainerInstanceName({ ...base, chunk: 0, shard: 0 })).toBe( - 'run-test:model/a:2:0' + 'run-test:model/a::2:0' ); expect(getDeciderContainerInstanceName({ ...base, chunk: 16, shard: 0 })).toBe( - 'run-test:model/a:2:0' + 'run-test:model/a::2:0' ); expect(getDeciderContainerInstanceName({ ...base, chunk: 1, shard: 1 })).toBe( - 'run-test:model/a:2:1' + 'run-test:model/a::2:1' + ); + }); + + it('two variants of one model get distinct queue messages and container lanes', () => { + const chunks = chunkArray( + Array.from({ length: 10 }, (_, i) => ({ id: `c${i}` })), + 5 + ); + const entries = [ + { model: 'model/a', variant: 'xhigh' as string | null }, + { model: 'model/a', variant: 'max' as string | null }, + ]; + const messages = buildDeciderMessages('run-v', 'decider', entries, 1, chunks, 100); + const models = messages.map(m => m.body.model); + const variants = messages.map(m => m.body.variant); + expect(new Set(models)).toEqual(new Set(['model/a'])); + expect(new Set(variants)).toEqual(new Set(['xhigh', 'max'])); + expect( + getDeciderContainerInstanceName({ + runId: 'run-v', + model: 'model/a', + variant: 'xhigh', + rep: 0, + shard: 0, + }) + ).not.toBe( + getDeciderContainerInstanceName({ + runId: 'run-v', + model: 'model/a', + variant: 'max', + rep: 0, + shard: 0, + }) + ); + }); + + it('message schema accepts optional nullable variant', () => { + expect( + BenchmarkJobMessageSchema.parse({ + runId: 'r1', + kind: 'decider', + model: 'm1', + variant: 'xhigh', + }).variant + ).toBe('xhigh'); + expect( + BenchmarkJobMessageSchema.parse({ + runId: 'r1', + kind: 'decider', + model: 'm1', + variant: null, + }).variant + ).toBeNull(); + }); +}); + +describe('migration 0005 exact-pair backfill', () => { + it('rebuild SQL copies rows and backfills variant from reasoning_effort', async () => { + const fs = await import('node:fs/promises'); + const path = await import('node:path'); + const { fileURLToPath } = await import('node:url'); + const migrationPath = path.join( + path.dirname(fileURLToPath(import.meta.url)), + '../migrations/0005_fuzzy_senator_kelly.sql' + ); + const sql = await fs.readFile(migrationPath, 'utf8'); + + // Table rebuilds preserve existing columns via INSERT…SELECT. + expect(sql).toContain('CREATE TABLE `__new_run_models`'); + expect(sql).toContain('CREATE TABLE `__new_model_summaries`'); + expect(sql).toContain('CREATE TABLE `__new_case_results`'); + expect(sql).toMatch(/INSERT INTO `__new_run_models`[\s\S]*COALESCE\("reasoning_effort", ''\)/); + expect(sql).toMatch( + /INSERT INTO `__new_model_summaries`[\s\S]*COALESCE\(rm\."reasoning_effort", ''\)/ + ); + expect(sql).toMatch( + /INSERT INTO `__new_case_results`[\s\S]*COALESCE\(rm\."reasoning_effort", ''\)/ + ); + // routing_table_candidates: additive column + backfill from effort. + expect(sql).toContain( + "ALTER TABLE `routing_table_candidates` ADD `variant` text DEFAULT '' NOT NULL" + ); + expect(sql).toContain( + "UPDATE `routing_table_candidates` SET `variant` = COALESCE(`reasoning_effort`, '')" ); }); }); diff --git a/services/auto-routing-benchmark/src/run.ts b/services/auto-routing-benchmark/src/run.ts index 38be8136d5..daaef54caa 100644 --- a/services/auto-routing-benchmark/src/run.ts +++ b/services/auto-routing-benchmark/src/run.ts @@ -16,6 +16,7 @@ import { DECIDER_CASES } from './datasets/decider-cases'; import type { RunModelRow } from './db'; import { countCaseResults, + exactPairKey, existsNewerCompletedRun, getCaseResults, getExistingCaseResultIds, @@ -24,12 +25,18 @@ import { getRunWithModels, getSummaries, insertRun, + listPendingCurrentProfiles, + listStaleRunningDeciderRunIds, + markProfilesFailedForRun, + markProfilesReadyForRun, + markProfilesRunningForRun, markRunCompleted, markRunFailed, markStaleRunsFailed, replaceModelSummaries, saveRoutingTable, upsertCaseResult, + type BenchmarkRunPurpose, type CaseResultRow, type PriorModelResult, } from './db'; @@ -42,13 +49,27 @@ import { runDeciderCaseViaCli, warmUpCliContainer, } from './cli-runner'; -import { parsePersistedReasoningEffort } from './reasoning-effort'; +import { + parsePersistedReasoningEffort, + variantFromReasoningEffort, + variantFromStorage, + variantToStorage, +} from './reasoning-effort'; import { pickClassifierWinner } from './winner'; +/** One exact Pool entry identity used throughout a decider/classifier run. */ +export type RunModelEntry = { + model: string; + /** Canonical variant; null = default/no variant. */ + variant: string | null; +}; + export type BenchmarkJobMessage = { runId: string; kind: BenchmarkKind; model: string; + /** Exact-pair identity. Optional for backward-compatible old queue messages. */ + variant?: string | null; // The case ids this message is responsible for, plus the chunk index. Decider // chunks are split across shard lanes; each lane has one stable container. caseIds?: string[]; @@ -63,6 +84,7 @@ export const BenchmarkJobMessageSchema = z.object({ runId: z.string().min(1), kind: z.enum(['classifier', 'decider']), model: z.string().min(1), + variant: z.string().trim().min(1).nullable().optional(), caseIds: z.array(z.string().min(1)).optional(), chunk: z.number().int().min(0).optional(), shard: z.number().int().min(0).optional(), @@ -128,24 +150,87 @@ async function enqueueRunMessages( try { await env.BENCH_QUEUE.sendBatch(messages.slice(i, i + QUEUE_SEND_BATCH_LIMIT)); } catch (error) { - await markRunFailed( - env.BENCH_DB, - runId, - `enqueue failed after ${i} of ${messages.length} messages: ${formatError(error).error}` - ).catch(() => {}); + const reason = `enqueue failed after ${i} of ${messages.length} messages: ${formatError(error).error}`; + await failRunAndDrain(env, runId, reason).catch(() => {}); throw error; } } } +/** + * Mark a running run failed, transition any profile rows it claimed, and attempt + * to drain the oldest pending profile batch into the freed single decider slot. + */ +export async function failRunAndDrain(env: Env, runId: string, error: string): Promise { + await markRunFailed(env.BENCH_DB, runId, error); + try { + await markProfilesFailedForRun(env.BENCH_DB, runId, error); + } catch (profileError) { + console.warn( + JSON.stringify({ + event: 'benchmark_profile_fail_transition_error', + runId, + ...formatError(profileError), + }) + ); + } + try { + await drainPendingProfileBatch(env); + } catch (drainError) { + console.warn( + JSON.stringify({ + event: 'benchmark_profile_drain_error', + afterRunId: runId, + ...formatError(drainError), + }) + ); + } +} + const STALE_RUN_MAX_AGE_MS = 6 * 3600_000; // Fails any run still 'running' past the stale threshold (queue retries // exhausted / dead-lettered). Called both before starting a run and when // listing runs, so a wedged run is recovered without depending on a new run // being started (the UI disables Start while a run shows 'running'). -export async function sweepStaleRuns(db: D1Database): Promise { - await markStaleRunsFailed(db, new Date(Date.now() - STALE_RUN_MAX_AGE_MS).toISOString()); +// Also fails Benchmark-profile rows claimed by those stale runs so they do +// not stay stuck in `running`. Drain is separate (needs Env for startRun). +export async function sweepStaleRuns(db: D1Database): Promise { + const olderThanIso = new Date(Date.now() - STALE_RUN_MAX_AGE_MS).toISOString(); + // Capture ids before the bulk status flip so profile claims can be failed. + const staleIds = await listStaleRunningDeciderRunIds(db, olderThanIso).catch( + () => [] as string[] + ); + await markStaleRunsFailed(db, olderThanIso); + for (const runId of staleIds) { + await markProfilesFailedForRun(db, runId, 'timed out').catch(() => {}); + } + return staleIds; +} + +/** + * Sweep stale runs and attempt the pending-profile drain. Used by the + * scheduled handler so a failed final queue message or platform run cannot + * strand pending profile work. Never throws for drain/slot contention. + */ +export async function sweepStaleRunsAndDrain(env: Env): Promise<{ + staleRunIds: string[]; + drained: { runId: string; entryCount: number } | null; +}> { + const staleRunIds = await sweepStaleRuns(env.BENCH_DB); + let drained: { runId: string; entryCount: number } | null = null; + try { + drained = await drainPendingProfileBatch(env); + } catch (error) { + console.warn( + JSON.stringify({ + event: 'benchmark_profile_drain_error', + afterSweep: true, + ...formatError(error), + }) + ); + } + return { staleRunIds, drained }; } // Bump when grading logic, the CLI invocation/variant handling, the container @@ -182,26 +267,32 @@ export function computeEngineIdentity(kind: BenchmarkKind): string { return `v${BENCHMARK_ENGINE_VERSION}:${fnv1aHex(JSON.stringify(datasetSignature))}`; } +function normalizeRunModelEntries(models: readonly (string | RunModelEntry)[]): RunModelEntry[] { + return models.map(entry => (typeof entry === 'string' ? { model: entry, variant: null } : entry)); +} + /** Pure helper: produces the initial sendBatch bodies for a decider run. - * Extracted for unit-testability; the shape is models × reps messages. Later + * Extracted for unit-testability; the shape is entries × reps messages. Later * chunks are chained by processDeciderJob after the previous chunk completes. + * Accepts model id strings or exact `{ model, variant }` entries. */ export function buildDeciderMessages( runId: string, kind: BenchmarkKind, - modelIds: string[], + models: readonly (string | RunModelEntry)[], repetitions: number, chunks: readonly (readonly { id: string }[])[], maxLiveContainers: number = DECIDER_CONTAINER_INSTANCE_CAP ): { body: BenchmarkJobMessage }[] { + const entries = normalizeRunModelEntries(models); const shardCount = computeDeciderShardCount({ - modelCount: modelIds.length, + modelCount: entries.length, repetitions, chunkCount: chunks.length, maxLiveContainers, }); if (shardCount === 0) return []; - return modelIds.flatMap(model => + return entries.flatMap(entry => Array.from({ length: repetitions }, (_, rep) => Array.from({ length: shardCount }, (_, shard) => { const chunkCases = chunks[shard]; @@ -211,7 +302,8 @@ export function buildDeciderMessages( body: { runId, kind, - model, + model: entry.model, + variant: entry.variant, chunk: shard, shard, shardCount, @@ -226,24 +318,29 @@ export function buildDeciderMessages( } export function getDeciderContainerInstanceName( - message: Pick + message: Pick ): string { - return `${message.runId}:${message.model}:${message.rep ?? 0}:${message.shard ?? 0}`; + // Include variant so two variants of one model in one run get distinct lanes. + // Empty string for null keeps names stable and filesystem/DO-name safe. + const variantPart = variantToStorage(message.variant); + return `${message.runId}:${message.model}:${variantPart}:${message.rep ?? 0}:${message.shard ?? 0}`; } export function buildClassifierMessages( runId: string, - modelIds: string[], + models: readonly (string | RunModelEntry)[], repetitions: number, chunks: readonly (readonly { id: string }[])[] ): { body: BenchmarkJobMessage }[] { - return modelIds.flatMap(model => + const entries = normalizeRunModelEntries(models); + return entries.flatMap(entry => Array.from({ length: repetitions }, (_, rep) => chunks.map((chunkCases, chunk) => ({ body: { runId, kind: 'classifier', - model, + model: entry.model, + variant: entry.variant, chunk, rep, caseIds: chunkCases.map(c => c.id), @@ -275,7 +372,7 @@ export class BenchmarkRunConfigError extends Error { } } -function validateDeciderContainerBudget({ +export function validateDeciderContainerBudget({ modelCount, repetitions, maxLiveContainers, @@ -292,11 +389,34 @@ function validateDeciderContainerBudget({ ); } +export type StartRunOptions = { + force?: boolean; + /** + * 'platform' (default): publish the default routing table / classifier winner. + * 'profile': measure an explicit Pool-entry snapshot for the global registry; + * never publishes platform artifacts. + */ + purpose?: BenchmarkRunPurpose; + /** + * Explicit decider entry snapshot for profile runs. Required when + * purpose === 'profile'. Ignored for platform runs (entries come from config). + */ + entries?: readonly RunModelEntry[]; +}; + export async function startRun( env: Env, kind: BenchmarkKind, - options: { force?: boolean } = {} + options: StartRunOptions = {} ): Promise<{ runId: string; enqueuedModels: number; skippedModels: string[] }> { + const purpose: BenchmarkRunPurpose = options.purpose ?? 'platform'; + if (purpose === 'profile' && kind !== 'decider') { + throw new BenchmarkRunConfigError('profile runs must be decider kind'); + } + if (purpose === 'profile' && (!options.entries || options.entries.length === 0)) { + throw new BenchmarkRunConfigError('profile runs require a non-empty entries snapshot'); + } + // Stale-run sweeper: fail dead 'running' runs first so a wedged run can't // block new ones and the admin panel shows the truth. await sweepStaleRuns(env.BENCH_DB); @@ -309,63 +429,88 @@ export async function startRun( // One active run per kind. The unique partial index is the atomic backstop; // this pre-check turns the common case (a run already going) into a clean // RunAlreadyActiveError instead of an insert-constraint failure. + // Platform and profile share this single decider slot. const activeRun = await getRunningRun(env.BENCH_DB, kind); if (activeRun) { throw new RunAlreadyActiveError(kind, activeRun.id); } const repetitions = kind === 'classifier' ? config.classifierRepetitions : config.deciderRepetitions; - const models = - kind === 'classifier' ? config.classifierModels : config.deciderModels.map(m => m.id); + + // Platform: config maps reasoningEffort → stored variant. Profile: explicit + // snapshot (may include two variants of one model). + const profileEntries = options.entries; + const modelEntries: RunModelEntry[] = + purpose === 'profile' && profileEntries + ? normalizeRunModelEntries(profileEntries) + : kind === 'classifier' + ? config.classifierModels.map(model => ({ model, variant: null })) + : config.deciderModels.map(m => ({ + model: m.id, + variant: + m.variant !== undefined + ? (m.variant ?? null) + : variantFromReasoningEffort(m.reasoningEffort ?? null), + })); const engineIdentity = computeEngineIdentity(kind); - const reasoningEffortFor = (modelId: string): string | null => - kind === 'classifier' - ? null - : (config.deciderModels.find(m => m.id === modelId)?.reasoningEffort ?? null); - - // Models with prior results are skipped (their latest summaries are carried - // into this run's aggregate) unless the admin forces a full re-run. A prior - // result is only carried when it was measured under the SAME benchmark - // identity — engine identity (dataset + grading/CLI version), repetitions, - // and the model's reasoning_effort — so a config/dataset change re-benchmarks - // the model instead of pairing current serving config with stale numbers. - const priorByModel = options.force - ? new Map() - : await getLatestSummariesByModel(env.BENCH_DB, kind); - const isCarryable = (modelId: string): boolean => { - const prior = priorByModel.get(modelId); + + // Exact entries with prior results are skipped (their latest summaries are + // carried into this run's aggregate) unless the admin forces a full re-run. + // A prior result is only carried when it was measured under the SAME + // benchmark identity — engine identity (dataset + grading/CLI version), + // repetitions, AND exact variant — so a config/dataset/variant change + // re-benchmarks instead of pairing current serving config with stale numbers. + // Profile runs always re-measure the snapshot (no carry) so registry + // provenance points at a run that actually graded these entries. + const priorByPair = + options.force || purpose === 'profile' + ? new Map() + : await getLatestSummariesByModel(env.BENCH_DB, kind); + const isCarryable = (entry: RunModelEntry): boolean => { + const prior = priorByPair.get(exactPairKey(entry.model, entry.variant)); return ( prior !== undefined && prior.engineIdentity === engineIdentity && prior.repetitions === repetitions && - (prior.reasoningEffort ?? null) === reasoningEffortFor(modelId) + (prior.variant ?? null) === (entry.variant ?? null) ); }; - const enqueuedModelIds = models.filter(m => !isCarryable(m)); - const skippedModels = models.filter(m => isCarryable(m)); - const carriedSummaries = skippedModels.flatMap(m => priorByModel.get(m)?.summaries ?? []); + const enqueuedEntries = modelEntries.filter(e => !isCarryable(e)); + const skippedEntries = modelEntries.filter(e => isCarryable(e)); + const skippedModels = skippedEntries.map(e => e.model); + const carriedSummaries = skippedEntries.flatMap( + e => priorByPair.get(exactPairKey(e.model, e.variant))?.summaries ?? [] + ); const benchmarkIdentity = resolveBenchmarkIdentity(config); const maxLiveDeciderContainers = Math.min(config.maxConcurrency, DECIDER_CONTAINER_INSTANCE_CAP); if (kind === 'decider') { validateDeciderContainerBudget({ - modelCount: enqueuedModelIds.length, + modelCount: enqueuedEntries.length, repetitions, maxLiveContainers: maxLiveDeciderContainers, }); } const startedAt = new Date().toISOString(); - const runId = `${kind}-${startedAt.replace(/[:.]/g, '-')}`; - - // Build run_models rows for ALL models of this run's kind. - const runModelRows: RunModelRow[] = models.map(modelId => ({ + // Distinguish profile run ids so ops logs and D1 rows are easy to filter. + const runId = + purpose === 'profile' + ? `profile-${startedAt.replace(/[:.]/g, '-')}` + : `${kind}-${startedAt.replace(/[:.]/g, '-')}`; + + // Build run_models rows for ALL exact entries of this run. + // Platform runs write variant AND keep reasoning_effort as a legacy mirror. + // Profile runs write variant only (reasoning_effort null). + const enqueuedPairKeys = new Set(enqueuedEntries.map(e => exactPairKey(e.model, e.variant))); + const runModelRows: RunModelRow[] = modelEntries.map(entry => ({ run_id: runId, - model: modelId, - enqueued: enqueuedModelIds.includes(modelId), - reasoning_effort: reasoningEffortFor(modelId), + model: entry.model, + variant: variantToStorage(entry.variant), + enqueued: enqueuedPairKeys.has(exactPairKey(entry.model, entry.variant)), + reasoning_effort: purpose === 'profile' ? null : entry.variant, })); try { @@ -385,6 +530,7 @@ export async function startRun( classifier_max_p95_latency_ms: kind === 'classifier' ? config.classifierMaxP95LatencyMs : null, engine_identity: engineIdentity, + purpose, }, runModelRows, carriedSummaries @@ -400,17 +546,37 @@ export async function startRun( throw error; } + if (purpose === 'profile') { + try { + await markProfilesRunningForRun(env.BENCH_DB, runId, modelEntries, { + engineIdentity, + repetitions, + }); + } catch (error) { + console.error( + JSON.stringify({ + event: 'benchmark_profile_running_transition_error', + runId, + ...formatError(error), + }) + ); + await failRunAndDrain(env, runId, formatError(error).error); + throw error; + } + } + console.log( JSON.stringify({ event: 'benchmark_run_started', runId, kind, - enqueuedModels: enqueuedModelIds, + purpose, + enqueuedModels: enqueuedEntries.map(e => e.model), skippedModels, }) ); - if (enqueuedModelIds.length === 0) { + if (enqueuedEntries.length === 0) { // Everything already has results: complete immediately and republish the // aggregate so config-only changes (model removed, threshold tweaked) // take effect without re-running any model. The state mirrors the rows @@ -426,15 +592,16 @@ export async function startRun( repetitions, classifierMaxP95LatencyMs: kind === 'classifier' ? config.classifierMaxP95LatencyMs : null, startedAt, + purpose, }); return { runId, enqueuedModels: 0, skippedModels }; } if (kind === 'classifier') { const chunks = chunkArray(CLASSIFIER_CASES, CLASSIFIER_CHUNK_SIZE); - const messages = buildClassifierMessages(runId, enqueuedModelIds, repetitions, chunks); + const messages = buildClassifierMessages(runId, enqueuedEntries, repetitions, chunks); await enqueueRunMessages(env, runId, messages); - return { runId, enqueuedModels: enqueuedModelIds.length, skippedModels }; + return { runId, enqueuedModels: enqueuedEntries.length, skippedModels }; } // Decider: seed as many shard lanes as fit under the live-container cap. Each @@ -444,13 +611,100 @@ export async function startRun( const messages = buildDeciderMessages( runId, kind, - enqueuedModelIds, + enqueuedEntries, repetitions, chunks, maxLiveDeciderContainers ); await enqueueRunMessages(env, runId, messages); - return { runId, enqueuedModels: enqueuedModelIds.length, skippedModels }; + return { runId, enqueuedModels: enqueuedEntries.length, skippedModels }; +} + +/** + * After any decider run reaches a terminal state, claim the oldest pending + * profile batch that fits the container budget and start a profile run. + * No-ops when the decider slot is occupied or there is no pending work. + * Returns the started run id, or null. + */ +export async function drainPendingProfileBatch( + env: Env +): Promise<{ runId: string; entryCount: number } | null> { + const config = await getBenchmarkConfig(env.BENCH_DB); + if (!config) return null; + + const active = await getRunningRun(env.BENCH_DB, 'decider'); + if (active) { + console.log( + JSON.stringify({ + event: 'benchmark_profile_drain_skipped_slot_occupied', + activeRunId: active.id, + }) + ); + return null; + } + + const engineIdentity = computeEngineIdentity('decider'); + const repetitions = config.deciderRepetitions; + const pending = await listPendingCurrentProfiles(env.BENCH_DB, { + engineIdentity, + repetitions, + }); + if (pending.length === 0) return null; + + // Oldest-first (query already ordered by requested_at). Cap by container budget: + // modelCount * repetitions <= maxLiveContainers. + const maxLive = Math.min(config.maxConcurrency, DECIDER_CONTAINER_INSTANCE_CAP); + const maxEntries = Math.max(1, Math.floor(maxLive / Math.max(1, repetitions))); + const batch = pending.slice(0, maxEntries); + const entries: RunModelEntry[] = batch.map(row => ({ + model: row.model, + variant: variantFromStorage(row.variant), + })); + + try { + validateDeciderContainerBudget({ + modelCount: entries.length, + repetitions, + maxLiveContainers: maxLive, + }); + } catch (error) { + // Even a single entry can exceed the budget when repetitions are huge; + // leave pending and log — scheduled sweep will retry after config fix. + console.warn( + JSON.stringify({ + event: 'benchmark_profile_drain_budget_exceeded', + entryCount: entries.length, + repetitions, + maxLive, + ...formatError(error), + }) + ); + return null; + } + + try { + const result = await startRun(env, 'decider', { purpose: 'profile', entries }); + console.log( + JSON.stringify({ + event: 'benchmark_profile_drain_started', + runId: result.runId, + entryCount: entries.length, + models: entries.map(e => e.model), + }) + ); + return { runId: result.runId, entryCount: entries.length }; + } catch (error) { + if (error instanceof RunAlreadyActiveError) { + console.log( + JSON.stringify({ + event: 'benchmark_profile_drain_skipped_slot_occupied', + activeRunId: error.activeRunId, + }) + ); + return null; + } + throw error; + } } export async function processJob(env: Env, rawMessage: unknown): Promise { @@ -491,6 +745,7 @@ export async function processJob(env: Env, rawMessage: unknown): Promise { const expandedItems = CLASSIFIER_CASES.filter(benchCase => caseIds.has(benchCase.id)).map( benchCase => ({ benchCase, rep }) ); + const classifierVariant = variantToStorage(message.variant); await runCasesWithConcurrency( expandedItems, state.maxConcurrency, @@ -504,6 +759,7 @@ export async function processJob(env: Env, rawMessage: unknown): Promise { await upsertCaseResult(env.BENCH_DB, { run_id: message.runId, model: message.model, + variant: classifierVariant, case_id: benchCase.id, route_key: null, score, @@ -548,6 +804,7 @@ type RunState = { repetitions: number; classifierMaxP95LatencyMs: number | null; startedAt: string; + purpose: BenchmarkRunPurpose; }; async function getRunState(env: Env, runId: string): Promise { @@ -566,6 +823,7 @@ async function getRunState(env: Env, runId: string): Promise { repetitions: run.repetitions, classifierMaxP95LatencyMs: run.classifier_max_p95_latency_ms, startedAt: run.started_at, + purpose: run.purpose === 'profile' ? 'profile' : 'platform', }; } @@ -604,20 +862,55 @@ async function processDeciderJob( const chunk = message.chunk ?? 0; const shard = message.shard ?? 0; const shardCount = message.shardCount ?? 1; - const instanceName = getDeciderContainerInstanceName(message); + + // Resolve the run-snapshot row for this message's exact pair. + // - Explicit message.variant (including null): require exact (model, variant). + // A missing pair is corrupt — throw so the queue retries then dead-letters; + // never grade with a CLI variant taken from a non-matching row. + // - Legacy pre-deploy messages (variant undefined): model-only fallback only + // when the snapshot has exactly one row for that model; otherwise throw. + const modelRowsForModel = state.models.filter(m => m.model === message.model); + let modelRow: RunModelRow; + if (message.variant !== undefined) { + const messageVariant = message.variant ?? null; + const exact = modelRowsForModel.find(m => variantFromStorage(m.variant) === messageVariant); + if (!exact) { + throw new Error( + `run ${message.runId} has no snapshot row for model ${message.model} variant ${JSON.stringify(messageVariant)}` + ); + } + modelRow = exact; + } else { + const sole = modelRowsForModel.length === 1 ? modelRowsForModel[0] : undefined; + if (!sole) { + throw new Error( + `run ${message.runId}: legacy decider message for model ${message.model} requires exactly one snapshot row, found ${modelRowsForModel.length}` + ); + } + modelRow = sole; + } + const entryVariant = + message.variant !== undefined + ? (message.variant ?? null) + : variantFromStorage(modelRow.variant); + const storedVariant = variantToStorage(entryVariant); + // Canonical variant passed to the CLI — derived only from the resolved pair. + const cliVariant = entryVariant ?? modelRow.reasoning_effort ?? null; + + const instanceName = getDeciderContainerInstanceName({ + ...message, + variant: entryVariant, + }); const existingCaseIds = await getExistingCaseResultIds(env.BENCH_DB, { runId: message.runId, model: message.model, + variant: entryVariant, rep, caseIds: cases.map(c => c.id), }); const casesToRun = cases.filter(c => !existingCaseIds.has(c.id)); - // Reasoning effort comes from the run snapshot (run_models row), not live config. - const modelRow = state.models.find(m => m.model === message.model); - const reasoningEffort = modelRow?.reasoning_effort ?? null; - if (casesToRun.length > 0) { // Fetch a short-lived user token ONCE per queue message. Non-OK throws so the // queue retries the message. The token is never logged. @@ -655,7 +948,7 @@ async function processDeciderJob( kiloToken, kiloApiUrl: env.KILO_CLI_API_URL, orgId: state.benchmarkOrgId, - reasoningEffort, + variant: cliVariant, }); // The CLI occasionally ends a session with no assistant text at all // (transient empty completion: a lone step_finish with cost 0). Mirror @@ -670,7 +963,7 @@ async function processDeciderJob( kiloToken, kiloApiUrl: env.KILO_CLI_API_URL, orgId: state.benchmarkOrgId, - reasoningEffort, + variant: cliVariant, }); retry.costUsd = retry.costUsd === null && result.costUsd === null @@ -685,6 +978,7 @@ async function processDeciderJob( await upsertCaseResult(env.BENCH_DB, { run_id: message.runId, model: message.model, + variant: storedVariant, case_id: benchCase.id, route_key: taxonomyRouteKey(benchCase), score: succeeded ? 1 : 0, @@ -704,7 +998,14 @@ async function processDeciderJob( if (isRetryableContainerAvailabilityError(error)) throw error; await upsertCaseResult( env.BENCH_DB, - failedRow(message, benchCase.id, taxonomyRouteKey(benchCase), startedAt, error, rep) + failedRow( + { ...message, variant: entryVariant }, + benchCase.id, + taxonomyRouteKey(benchCase), + startedAt, + error, + rep + ) ); } }); @@ -712,7 +1013,7 @@ async function processDeciderJob( const hasNextChunk = await enqueueNextDeciderChunkIfNeeded( env, - message, + { ...message, variant: entryVariant }, rep, chunk, shard, @@ -749,6 +1050,7 @@ async function enqueueNextDeciderChunkIfNeeded( const existingNextCaseIds = await getExistingCaseResultIds(env.BENCH_DB, { runId: message.runId, model: message.model, + variant: message.variant ?? null, rep, caseIds: nextCaseIds, }); @@ -760,6 +1062,7 @@ async function enqueueNextDeciderChunkIfNeeded( runId: message.runId, kind: 'decider', model: message.model, + variant: message.variant ?? null, chunk: nextChunkIndex, shard, shardCount, @@ -814,6 +1117,7 @@ function failedRow( return { run_id: message.runId, model: message.model, + variant: variantToStorage(message.variant), case_id: caseId, route_key: routeKey, score: 0, @@ -869,6 +1173,42 @@ async function finalizeRunIfComplete( await replaceModelSummaries(env.BENCH_DB, runId, freshSummaries); await markRunCompleted(env.BENCH_DB, runId); + // Profile runs update the registry only — never publish platform artifacts, + // never touch the classifier winner, never clear platform KV keys. + if (state.purpose === 'profile') { + try { + await markProfilesReadyForRun(env.BENCH_DB, runId); + } catch (error) { + console.warn( + JSON.stringify({ + event: 'benchmark_profile_ready_transition_error', + runId, + ...formatError(error), + }) + ); + } + console.log( + JSON.stringify({ + event: 'benchmark_profile_run_completed', + runId, + kind, + purpose: state.purpose, + }) + ); + try { + await drainPendingProfileBatch(env); + } catch (drainError) { + console.warn( + JSON.stringify({ + event: 'benchmark_profile_drain_error', + afterRunId: runId, + ...formatError(drainError), + }) + ); + } + return; + } + // Read back all summaries (fresh + carried) for publishing. const allSummaries = await getSummaries(env.BENCH_DB, runId); @@ -907,10 +1247,15 @@ async function finalizeRunIfComplete( const generatedAt = new Date().toISOString(); try { // Built from the run's own model snapshot, not live config, so a mid-run - // admin edit can't skew the published table. + // admin edit can't skew the published table. Platform tables still emit + // reasoningEffort (sourced from the run snapshot's effort key) — not + // variant — so the published artifact shape stays unchanged for rolling + // deploys of old auto-routing workers. const deciderModels: BenchmarkDeciderModel[] = state.models.map(m => ({ id: m.model, - reasoningEffort: parsePersistedReasoningEffort(m.reasoning_effort), + reasoningEffort: parsePersistedReasoningEffort( + m.reasoning_effort ?? variantFromStorage(m.variant) + ), })); const table = buildRoutingTable({ runId, @@ -943,18 +1288,35 @@ async function finalizeRunIfComplete( event: 'benchmark_run_completed', runId, kind, + purpose: state.purpose, summaries: allSummaries, }) ); + + // Free the single decider slot into the oldest pending profile batch. + if (kind === 'decider') { + try { + await drainPendingProfileBatch(env); + } catch (drainError) { + console.warn( + JSON.stringify({ + event: 'benchmark_profile_drain_error', + afterRunId: runId, + ...formatError(drainError), + }) + ); + } + } } export function summarize(rows: CaseResultRow[], kind: BenchmarkKind): BenchmarkModelSummary[] { - // Group by "model route-key" using a plain reduce so this works in all runtimes. - // Classifier rows use '*' because classification has no decider taxonomy route. + // Group by (model, variant, routeKey). Classifier rows use '*' because + // classification has no decider taxonomy route. variant '' → null on emit. const groups = new Map(); for (const row of rows) { const routeKey = kind === 'classifier' ? '*' : (row.route_key ?? '*'); - const key = `${row.model}\0${routeKey}`; + const storedVariant = row.variant ?? ''; + const key = `${row.model}\0${storedVariant}\0${routeKey}`; const existing = groups.get(key); if (existing) { existing.push(row); @@ -964,7 +1326,7 @@ export function summarize(rows: CaseResultRow[], kind: BenchmarkKind): Benchmark } return [...groups.entries()].map(([key, group]) => { - const [model, routeKey] = key.split('\0'); + const [model, storedVariant, routeKey] = key.split('\0'); const latencies = group.map(r => r.latency_ms).toSorted((a, b) => a - b); const costs = group.filter(r => r.cost_usd !== null); const p95LatencyMs = @@ -974,6 +1336,7 @@ export function summarize(rows: CaseResultRow[], kind: BenchmarkKind): Benchmark : null; return { model, + variant: variantFromStorage(storedVariant), routeKey: routeKey as BenchmarkModelSummary['routeKey'], accuracy: Number((group.reduce((a, r) => a + r.score, 0) / group.length).toFixed(4)), avgCostUsd: costs.length diff --git a/services/auto-routing/src/admin-routing-settings.ts b/services/auto-routing/src/admin-routing-settings.ts new file mode 100644 index 0000000000..c8d4c3c577 --- /dev/null +++ b/services/auto-routing/src/admin-routing-settings.ts @@ -0,0 +1,127 @@ +import { + AutoRoutingModeOwnerQuerySchema, + AutoRoutingSettingsResponseSchema, + BenchmarkProfileQuotaErrorSchema, + DEFAULT_AUTO_ROUTING_MODE, + UpdateAutoRoutingSettingsRequestSchema, + type AutoRoutingMode, + type AutoRoutingModeOwnerType, + type AutoRoutingSettingsResponse, + type BenchmarkProfileEntryStatus, + type EfficientModelPool, +} from '@kilocode/auto-routing-contracts'; +import type { Handler } from 'hono'; +import { + BenchmarkProfileQuotaError, + fetchBenchmarkProfileStatuses, + registerBenchmarkProfiles, +} from './benchmark-origin'; +import type { HonoEnv } from './hono-env'; +import { getConfiguredAutoRoutingSettings, setAutoRoutingSettings } from './routing-mode'; + +function settingsResponse(params: { + ownerType: AutoRoutingModeOwnerType; + ownerId: string; + configuredMode: AutoRoutingMode | null; + configuredPool: EfficientModelPool | null; + poolStatuses: BenchmarkProfileEntryStatus[]; +}): AutoRoutingSettingsResponse { + return AutoRoutingSettingsResponseSchema.parse({ + ownerType: params.ownerType, + ownerId: params.ownerId, + mode: params.configuredMode ?? DEFAULT_AUTO_ROUTING_MODE, + configuredMode: params.configuredMode, + defaultMode: DEFAULT_AUTO_ROUTING_MODE, + configuredPool: params.configuredPool, + poolStatuses: params.poolStatuses, + }); +} + +export const getRoutingSettingsHandler: Handler = async c => { + const parsed = AutoRoutingModeOwnerQuerySchema.safeParse({ + ownerType: c.req.query('ownerType'), + ownerId: c.req.query('ownerId'), + }); + if (!parsed.success) { + return c.json({ error: 'Invalid routing settings owner' }, 400); + } + + const configured = await getConfiguredAutoRoutingSettings(c.env, parsed.data); + if (configured.pool === null) { + return c.json( + settingsResponse({ + ...parsed.data, + configuredMode: configured.mode, + configuredPool: null, + poolStatuses: [], + }) + ); + } + + try { + const statuses = await fetchBenchmarkProfileStatuses(c.env, configured.pool); + return c.json( + settingsResponse({ + ...parsed.data, + configuredMode: configured.mode, + configuredPool: configured.pool, + poolStatuses: statuses.statuses, + }) + ); + } catch { + return c.json({ error: 'Failed to load pool profile statuses' }, 502); + } +}; + +export const putRoutingSettingsHandler: Handler = async c => { + let rawBody: unknown; + try { + rawBody = await c.req.json(); + } catch { + return c.json({ error: 'Invalid JSON body' }, 400); + } + + const parsed = UpdateAutoRoutingSettingsRequestSchema.safeParse(rawBody); + if (!parsed.success) { + return c.json({ error: 'Invalid routing settings' }, 400); + } + + const { ownerType, ownerId, mode, pool, retryEntries } = parsed.data; + + let poolStatuses: BenchmarkProfileEntryStatus[] = []; + if (pool !== null) { + try { + const registered = await registerBenchmarkProfiles(c.env, { + ownerType, + ownerId, + entries: pool, + ...(retryEntries !== undefined ? { retryEntries } : {}), + }); + poolStatuses = registered.statuses; + } catch (error) { + if (error instanceof BenchmarkProfileQuotaError) { + return c.json( + BenchmarkProfileQuotaErrorSchema.parse({ + error: error.message, + retryAt: error.retryAt, + }), + 429 + ); + } + return c.json({ error: 'Failed to register benchmark profiles' }, 502); + } + } + + // Persist only after successful admission (or when clearing the pool). + await setAutoRoutingSettings(c.env, { ownerType, ownerId }, { mode, pool }); + + return c.json( + settingsResponse({ + ownerType, + ownerId, + configuredMode: mode, + configuredPool: pool, + poolStatuses, + }) + ); +}; diff --git a/services/auto-routing/src/benchmark-origin.ts b/services/auto-routing/src/benchmark-origin.ts index ecbadc2fd6..92bafab73d 100644 --- a/services/auto-routing/src/benchmark-origin.ts +++ b/services/auto-routing/src/benchmark-origin.ts @@ -1,17 +1,49 @@ import { + BenchmarkProfileQuotaErrorSchema, + BenchmarkProfileStatusesRequestSchema, + BenchmarkProfileStatusesResponseSchema, BenchmarkRoutingTableResponseSchema, ClassifierWinnerResponseSchema, + CustomRoutingTableRequestSchema, + CustomRoutingTableResponseSchema, + RegisterBenchmarkProfilesRequestSchema, + type BenchmarkProfileStatusesResponse, type ClassifierWinner, + type CustomRoutingTable, + type EfficientModelPool, + type RegisterBenchmarkProfilesRequest, type RoutingTable, } from '@kilocode/auto-routing-contracts'; type BenchmarkEnv = Pick; -async function fetchBenchmark(env: BenchmarkEnv, path: string): Promise { +export class BenchmarkProfileQuotaError extends Error { + readonly retryAt: string; + + constructor(message: string, retryAt: string) { + super(message); + this.name = 'BenchmarkProfileQuotaError'; + this.retryAt = retryAt; + } +} + +async function fetchBenchmark( + env: BenchmarkEnv, + path: string, + init?: RequestInit +): Promise { const secret = await env.INTERNAL_API_SECRET_PROD.get(); - const res = await env.BENCHMARK_SERVICE.fetch(`https://auto-routing-benchmark${path}`, { - headers: { authorization: `Bearer ${secret}` }, + return env.BENCHMARK_SERVICE.fetch(`https://auto-routing-benchmark${path}`, { + ...init, + headers: { + authorization: `Bearer ${secret}`, + ...(init?.headers ?? {}), + }, }); +} + +async function fetchBenchmarkJson(env: BenchmarkEnv, path: string): Promise { + const res = await fetchBenchmark(env, path); if (!res.ok) { const detail = (await res.text().catch(() => '')).slice(0, 200); throw new Error(`benchmark origin ${path} responded ${res.status} ${detail}`); @@ -20,7 +52,7 @@ async function fetchBenchmark(env: BenchmarkEnv, path: string): Promise } export async function fetchRoutingTableFromOrigin(env: BenchmarkEnv): Promise { - const body = await fetchBenchmark(env, '/admin/routing-table'); + const body = await fetchBenchmarkJson(env, '/admin/routing-table'); const parsed = BenchmarkRoutingTableResponseSchema.safeParse(body); if (!parsed.success) { throw new Error( @@ -33,7 +65,7 @@ export async function fetchRoutingTableFromOrigin(env: BenchmarkEnv): Promise { - const body = await fetchBenchmark(env, '/admin/classifier-winner'); + const body = await fetchBenchmarkJson(env, '/admin/classifier-winner'); const parsed = ClassifierWinnerResponseSchema.safeParse(body); if (!parsed.success) { throw new Error( @@ -42,3 +74,88 @@ export async function fetchClassifierWinnerFromOrigin( } return parsed.data.winner; } + +export async function registerBenchmarkProfiles( + env: BenchmarkEnv, + request: RegisterBenchmarkProfilesRequest +): Promise { + const body = RegisterBenchmarkProfilesRequestSchema.parse(request); + const res = await fetchBenchmark(env, '/admin/profiles/register', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify(body), + }); + + if (res.status === 429) { + const raw: unknown = await res.json().catch(() => null); + const quota = BenchmarkProfileQuotaErrorSchema.safeParse(raw); + if (quota.success) { + throw new BenchmarkProfileQuotaError(quota.data.error, quota.data.retryAt); + } + throw new Error('benchmark profile register responded 429 with invalid body'); + } + + if (!res.ok) { + const detail = (await res.text().catch(() => '')).slice(0, 200); + throw new Error(`benchmark origin /admin/profiles/register responded ${res.status} ${detail}`); + } + + const raw: unknown = await res.json(); + const parsed = BenchmarkProfileStatusesResponseSchema.safeParse(raw); + if (!parsed.success) { + throw new Error( + `benchmark profile register response invalid: ${parsed.error.issues[0]?.message ?? 'unknown'}` + ); + } + return parsed.data; +} + +export async function fetchBenchmarkProfileStatuses( + env: BenchmarkEnv, + entries: EfficientModelPool +): Promise { + const body = BenchmarkProfileStatusesRequestSchema.parse({ entries }); + const res = await fetchBenchmark(env, '/admin/profiles/status', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify(body), + }); + if (!res.ok) { + const detail = (await res.text().catch(() => '')).slice(0, 200); + throw new Error(`benchmark origin /admin/profiles/status responded ${res.status} ${detail}`); + } + const raw: unknown = await res.json(); + const parsed = BenchmarkProfileStatusesResponseSchema.safeParse(raw); + if (!parsed.success) { + throw new Error( + `benchmark profile statuses response invalid: ${parsed.error.issues[0]?.message ?? 'unknown'}` + ); + } + return parsed.data; +} + +export async function fetchCustomRoutingTable( + env: BenchmarkEnv, + entries: EfficientModelPool +): Promise { + const body = CustomRoutingTableRequestSchema.parse({ entries }); + const res = await fetchBenchmark(env, '/admin/custom-routing-table', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify(body), + }); + if (!res.ok) { + const detail = (await res.text().catch(() => '')).slice(0, 200); + throw new Error( + `benchmark origin /admin/custom-routing-table responded ${res.status} ${detail}` + ); + } + const raw: unknown = await res.json(); + const parsed = CustomRoutingTableResponseSchema.safeParse(raw); + if (!parsed.success) { + throw new Error( + `benchmark custom routing-table response invalid: ${parsed.error.issues[0]?.message ?? 'unknown'}` + ); + } + return parsed.data.table; +} diff --git a/services/auto-routing/src/decide.ts b/services/auto-routing/src/decide.ts index a01f9e041c..1215825ca9 100644 --- a/services/auto-routing/src/decide.ts +++ b/services/auto-routing/src/decide.ts @@ -1,12 +1,16 @@ -import { MirrorPayloadSchema, taxonomyRouteKey } from '@kilocode/auto-routing-contracts'; -import type { - AutoRoutingDecision, - AutoRoutingDecisionResponse, - MirrorPayload, - NormalizedClassifierInput, - RoutingConstraints, +import { + MirrorPayloadSchema, + poolEntryKey, + taxonomyRouteKey, + type AutoRoutingDecision, + type AutoRoutingDecisionResponse, + type EfficientModelPool, + type MirrorPayload, + type NormalizedClassifierInput, + type RoutingConstraints, + type RoutingTable, } from '@kilocode/auto-routing-contracts'; -import { formatError } from '@kilocode/worker-utils'; +import { formatError, ttlCached } from '@kilocode/worker-utils'; import type { Handler } from 'hono'; import { writeClassifierMetricsDataPoint } from './classifier-analytics'; import type { ClassifierAnalyticsStatus } from './classifier-analytics'; @@ -26,16 +30,77 @@ import { putStickyDecision, } from './decision-cache'; import type { StickyDecision } from './decision-cache'; -import { computeDecision, modelSatisfiesConstraints } from './decision-engine'; +import { + computeDecision, + modelSatisfiesConstraints, + type DecisionIncumbent, +} from './decision-engine'; import { ClassifierRunError, classifyNormalizedInput } from './model-classifier'; import type { ClassifierRunResult } from './model-classifier'; +import { fetchCustomRoutingTable } from './benchmark-origin'; import { getRoutingTable } from './routing-table'; -import { getAutoRoutingMode } from './routing-mode'; +import { getEffectiveAutoRoutingSettings } from './routing-mode'; import type { HonoEnv } from './hono-env'; import { codingPlanDefaultDecision, getCodingPlanPreference } from './coding-plan-preference'; import { getModelCapabilities } from './model-capabilities'; import type { ModelCapabilitiesMap } from './model-capabilities'; +const CUSTOM_ROUTING_TABLE_CACHE_TTL_MS = 60_000; + +/** Ordered canonical key for an Efficient model pool (sort by poolEntryKey). */ +export function orderedPoolCacheKey(pool: EfficientModelPool): string { + return [...pool] + .map(entry => poolEntryKey(entry)) + .sort() + .join('\n'); +} + +// Per-pool-key isolate caches that close over the pool entries (~60s TTL). +const customTableLoaders = new Map< + string, + ReturnType> +>(); + +export function clearCustomRoutingTableCache(): void { + for (const cache of customTableLoaders.values()) { + cache.clear(); + } + customTableLoaders.clear(); +} + +async function loadCustomRoutingTable( + env: Env, + pool: EfficientModelPool +): Promise { + const key = orderedPoolCacheKey(pool); + let cache = customTableLoaders.get(key); + if (!cache) { + const poolSnapshot = pool; + cache = ttlCached(CUSTOM_ROUTING_TABLE_CACHE_TTL_MS, async (e: Env) => { + return fetchCustomRoutingTable(e, poolSnapshot); + }); + customTableLoaders.set(key, cache); + } + return cache.get(env).catch((error: unknown) => { + console.warn( + JSON.stringify({ event: 'auto_routing_custom_table_read_failed', ...formatError(error) }) + ); + return null; + }); +} + +function stickyToIncumbent(sticky: StickyDecision | null): DecisionIncumbent | null { + if (!sticky) return null; + return { model: sticky.model, variant: sticky.variant ?? null }; +} + +function decisionVariant(decision: AutoRoutingDecision): string | null { + if (decision.source !== 'benchmark') return null; + if (decision.variant !== undefined && decision.variant !== null) return decision.variant; + // Legacy effort-only decisions: sticky stores null variant (model-only). + return null; +} + // Isolate-scoped request counter, used to correlate latency with isolate // warm-up in logs. let isolateRequestSeq = 0; @@ -228,8 +293,13 @@ function recordDecision( }); const incumbentModel = incumbent?.model ?? null; + const decidedVariant = + decision !== null && decision.source === 'benchmark' ? (decision.variant ?? null) : null; const switched = - decision !== null && incumbentModel !== null && decision.model !== incumbentModel; + decision !== null && + incumbentModel !== null && + (decision.model !== incumbentModel || + (incumbent?.variant != null && decidedVariant !== incumbent.variant)); const routeKey = summary.classification ? taxonomyRouteKey(summary.classification) : null; // Null when there is no incumbent route to compare against (no incumbent, // a pre-routeKey cache entry, or no classification this request). @@ -275,6 +345,7 @@ function recordDecision( autoRoutingMode, uaPrefix: ctx.payload.userAgent?.slice(0, 40) ?? null, decidedModel: decision?.model ?? null, + decidedVariant, decidedTaskType: decision?.taskType ?? null, decidedSubtaskType: decision?.subtaskType ?? null, decisionSource: decision?.source ?? null, @@ -319,27 +390,51 @@ export const decideHandler: Handler = async c => { 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. + // Constraints-absent coding-plan short-circuit stays byte-identical: no + // settings hop, capability fetch, routing-table fetch, or benchmark hop. + if (codingPlanActive && !hasConstraints) { + 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 }); + } + + // Resolve settings before the (single) capability load so custom-pool + // model ids are included when a pool is configured. Settings do not + // depend on capabilities. Null pool keeps platform-table semantics and + // does not add a custom benchmark hop beyond the table loader below. + const effectiveSettings = await getEffectiveAutoRoutingSettings(c.env, { + userId: payload.userId, + organizationId: payload.organizationId, + }); + const configuredPool = effectiveSettings.pool; + const routingMode = effectiveSettings.mode; + const failClosedOnInactive = configuredPool !== null; + + // One capability load on the decide path. Include pool model ids whenever + // a pool is configured so constrained + custom-pool traffic does not pay + // two sequential 500ms capability budgets. let capabilities: ModelCapabilitiesMap = new Map(); - if (hasConstraints && constraints) { + if (hasConstraints || configuredPool !== null) { capabilities = await getModelCapabilities(c.env, { codingPlanModelId: codingPlanActive ? codingPlanPreference.modelId : null, + ...(configuredPool !== null + ? { additionalModelIds: configuredPool.map(entry => entry.model) } + : {}), }); } - if (codingPlanActive) { - const canTakeShortCircuit = - hasConstraints && constraints - ? modelSatisfiesConstraints(capabilities.get(codingPlanPreference.modelId), constraints) - : true; + if (codingPlanActive && hasConstraints && constraints) { + const canTakeShortCircuit = modelSatisfiesConstraints( + capabilities.get(codingPlanPreference.modelId), + constraints + ); if (canTakeShortCircuit) { const decision = codingPlanDefaultDecision(codingPlanPreference); writeClassifierMetricsDataPoint(c.env, { @@ -357,18 +452,19 @@ export const decideHandler: Handler = async c => { // from subscription-billed to credit-billed benchmark routing. } - const [hashes, userIdHash, classifierModel, successSampleRate, routingTable, routingMode] = - await Promise.all([ - computeContentHashes(payload.input), - hashIdentifierForTelemetry(payload.userId), - getClassifierModel(c.env), - getDecisionLogSampleRate(c.env), - getRoutingTable(c.env), - getAutoRoutingMode(c.env, { - userId: payload.userId, - organizationId: payload.organizationId, - }), - ]); + // Null pool uses the platform cached table (no custom benchmark hop). + // Configured pool loads the sparse custom table for exactly those entries; + // on lookup failure the table is null → null decision → gateway balanced fallback. + const [hashes, userIdHash, classifierModel, successSampleRate, routingTable] = await Promise.all([ + computeContentHashes(payload.input), + hashIdentifierForTelemetry(payload.userId), + getClassifierModel(c.env), + getDecisionLogSampleRate(c.env), + configuredPool === null + ? getRoutingTable(c.env) + : loadCustomRoutingTable(c.env, configuredPool), + ]); + const ctx: DecisionContext = { payload, hashes, @@ -384,21 +480,30 @@ export const decideHandler: Handler = async c => { getCachedClassification(c.env, ctx.conversationKey, hashes.exact, classifierModel), getStickyDecision(c.env, ctx.conversationKey), ]); + const incumbent = stickyToIncumbent(sticky); + const decisionOptions = { + constraints: payload.constraints, + capabilityMap: hasConstraints || configuredPool !== null ? capabilities : undefined, + failClosedOnInactive, + }; if (cached) { const decision = computeDecision( cached, routingTable, - sticky?.model ?? null, + incumbent, deniedModelIds, routingMode, - { - constraints: payload.constraints, - capabilityMap: hasConstraints ? capabilities : undefined, - } + decisionOptions ); if (decision) { c.executionCtx.waitUntil( - putStickyDecision(c.env, ctx.conversationKey, decision.model, taxonomyRouteKey(cached)) + putStickyDecision( + c.env, + ctx.conversationKey, + decision.model, + decisionVariant(decision), + taxonomyRouteKey(cached) + ) ); } recordDecision( @@ -431,13 +536,10 @@ export const decideHandler: Handler = async c => { const decision = computeDecision( classifier.classification, routingTable, - sticky?.model ?? null, + incumbent, deniedModelIds, routingMode, - { - constraints: payload.constraints, - capabilityMap: hasConstraints ? capabilities : undefined, - } + decisionOptions ); // Like the classification cache, sticky state only trusts real classifier // output: a heuristic fallback must not re-anchor the session's model. @@ -447,6 +549,7 @@ export const decideHandler: Handler = async c => { c.env, ctx.conversationKey, decision.model, + decisionVariant(decision), taxonomyRouteKey(classifier.classification) ) ); diff --git a/services/auto-routing/src/decision-cache.test.ts b/services/auto-routing/src/decision-cache.test.ts index c4e529cf7d..d9aa399038 100644 --- a/services/auto-routing/src/decision-cache.test.ts +++ b/services/auto-routing/src/decision-cache.test.ts @@ -126,20 +126,49 @@ describe('sticky decision storage', () => { return { env, cacheDO, storage }; } - it('round-trips the sticky model and route for a conversation', async () => { + it('round-trips the sticky model, variant, and route for a conversation', async () => { const { env } = createStickyEnv(); await expect(getStickyDecision(env, 'conversation-1')).resolves.toBeNull(); - await putStickyDecision(env, 'conversation-1', 'mid/chat', 'implementation/code_generation'); + await putStickyDecision( + env, + 'conversation-1', + 'mid/chat', + 'thinking', + 'implementation/code_generation' + ); await expect(getStickyDecision(env, 'conversation-1')).resolves.toEqual({ model: 'mid/chat', + variant: 'thinking', + routeKey: 'implementation/code_generation', + }); + }); + + it('round-trips a null sticky variant', async () => { + const { env } = createStickyEnv(); + await putStickyDecision( + env, + 'conversation-1', + 'mid/chat', + null, + 'implementation/code_generation' + ); + await expect(getStickyDecision(env, 'conversation-1')).resolves.toEqual({ + model: 'mid/chat', + variant: null, routeKey: 'implementation/code_generation', }); }); it('expires sticky entries after the TTL', async () => { const { env } = createStickyEnv(); - await putStickyDecision(env, 'conversation-1', 'mid/chat', 'implementation/code_generation'); + await putStickyDecision( + env, + 'conversation-1', + 'mid/chat', + null, + 'implementation/code_generation' + ); vi.advanceTimersByTime(31 * 60 * 1000); await expect(getStickyDecision(env, 'conversation-1')).resolves.toBeNull(); @@ -152,13 +181,14 @@ describe('sticky decision storage', () => { await expect(getStickyDecision(env, 'conversation-1')).resolves.toBeNull(); }); - it('serves entries written before the routeKey field existed', async () => { + it('serves entries written before the routeKey and variant fields existed', async () => { const { env, cacheDO } = createStickyEnv(); await cacheDO.putEntry('sticky', { model: 'mid/chat' } as unknown as ClassifierOutput); await expect(getStickyDecision(env, 'conversation-1')).resolves.toEqual({ model: 'mid/chat', routeKey: null, + variant: null, }); }); }); diff --git a/services/auto-routing/src/decision-cache.ts b/services/auto-routing/src/decision-cache.ts index 4e0fbf5770..cfa9b58841 100644 --- a/services/auto-routing/src/decision-cache.ts +++ b/services/auto-routing/src/decision-cache.ts @@ -91,6 +91,9 @@ const StickyDecisionSchema = z.object({ // Taxonomy route the model was decided on, for route-change telemetry. // Nullable/defaulted so entries written before the field existed parse. routeKey: z.string().min(1).nullish().default(null), + // Exact catalog variant the model was decided with. Null for legacy + // model-only sticky records and for candidates without a variant. + variant: z.string().min(1).nullish().default(null), }); export type StickyDecision = z.infer; @@ -134,11 +137,13 @@ export async function putStickyDecision( env: DecisionCacheEnv, conversationKey: string, model: string, + variant: string | null, routeKey: string ): Promise { try { await cacheStub(env, conversationKey).putEntry(STICKY_DECISION_KEY, { model, + variant, routeKey, } satisfies StickyDecision); } catch { diff --git a/services/auto-routing/src/decision-engine.test.ts b/services/auto-routing/src/decision-engine.test.ts index a1f1581f66..815fc4f997 100644 --- a/services/auto-routing/src/decision-engine.test.ts +++ b/services/auto-routing/src/decision-engine.test.ts @@ -1,21 +1,29 @@ import { describe, expect, it } from 'vitest'; import type { ClassifierOutput, RoutingTable } from '@kilocode/auto-routing-contracts'; -import { computeDecision } from './decision-engine'; +import { computeDecision, type DecisionIncumbent } from './decision-engine'; import type { ModelCapabilities, ModelCapabilitiesMap } from './model-capabilities'; function makeCaps( - rows: Record + rows: Record< + string, + { inputModalities?: string[]; contextLength?: number | null; isActive?: boolean | null } + > ): 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, + isActive: row.isActive ?? true, }); } return map; } +function inc(model: string | null, variant: string | null = null): DecisionIncumbent | null { + return model === null ? null : { model, variant }; +} + const classification: ClassifierOutput = { taskType: 'implementation', subtaskType: 'code_generation', @@ -113,7 +121,13 @@ describe('computeDecision', () => { }); }); it('does not keep a lower-accuracy incumbent in best accuracy mode', () => { - const decision = computeDecision(classification, table, 'mid/chat', new Set(), 'best_accuracy'); + const decision = computeDecision( + classification, + table, + inc('mid/chat'), + new Set(), + 'best_accuracy' + ); expect(decision).toMatchObject({ model: 'pricey/chat', sticky: false, switchReason: 'cost' }); }); it('keeps a best-accuracy incumbent when the fresh pick is less than five points better', () => { @@ -142,7 +156,7 @@ describe('computeDecision', () => { const decision = computeDecision( classification, nearTieTable, - 'incumbent/chat', + inc('incumbent/chat'), new Set(), 'best_accuracy' ); @@ -174,7 +188,7 @@ describe('computeDecision', () => { const decision = computeDecision( classification, betterTable, - 'incumbent/chat', + inc('incumbent/chat'), new Set(), 'best_accuracy' ); @@ -244,7 +258,7 @@ describe('computeDecision', () => { it('keeps the incumbent on route changes when it is within the switch-cost factor', () => { // Fresh pick cheap/chat at 0.002; mid/chat at 0.005 is not cheaper by // more than 3x (0.002 * 3 = 0.006 >= 0.005), so the session stays put. - const decision = computeDecision(classification, table, 'mid/chat'); + const decision = computeDecision(classification, table, inc('mid/chat')); expect(decision).toEqual({ model: 'mid/chat', taskType: 'implementation', @@ -278,16 +292,21 @@ describe('computeDecision', () => { ], }, }; - const decision = computeDecision(classification, boundaryTable, 'incumbent/chat'); + const decision = computeDecision(classification, boundaryTable, inc('incumbent/chat')); expect(decision).toMatchObject({ model: 'incumbent/chat', sticky: true }); }); it('switches when the fresh pick is cheaper by more than the factor', () => { // pricey/chat at 0.02 vs fresh 0.002 * 3 = 0.006: switch pays off. - const decision = computeDecision(classification, table, 'pricey/chat'); + const decision = computeDecision(classification, table, inc('pricey/chat')); expect(decision).toMatchObject({ model: 'cheap/chat', sticky: false, switchReason: 'cost' }); }); it('does not keep a denied incumbent', () => { - const decision = computeDecision(classification, table, 'mid/chat', new Set(['mid/chat'])); + const decision = computeDecision( + classification, + table, + inc('mid/chat'), + new Set(['mid/chat']) + ); expect(decision).toMatchObject({ model: 'cheap/chat', sticky: false, @@ -295,7 +314,7 @@ describe('computeDecision', () => { }); }); it('switches when the incumbent no longer meets the route threshold', () => { - const decision = computeDecision(classification, table, 'weak/chat'); + const decision = computeDecision(classification, table, inc('weak/chat')); expect(decision).toMatchObject({ model: 'cheap/chat', sticky: false, @@ -303,7 +322,7 @@ describe('computeDecision', () => { }); }); it('serves the fresh pick when the incumbent is not in the route', () => { - const decision = computeDecision(classification, table, 'gone/model'); + const decision = computeDecision(classification, table, inc('gone/model')); expect(decision).toMatchObject({ model: 'cheap/chat', sticky: false, @@ -311,7 +330,7 @@ describe('computeDecision', () => { }); }); it('is not sticky when the incumbent is the fresh pick', () => { - const decision = computeDecision(classification, table, 'cheap/chat'); + const decision = computeDecision(classification, table, inc('cheap/chat')); expect(decision).toMatchObject({ model: 'cheap/chat', sticky: false, switchReason: null }); }); }); @@ -356,7 +375,7 @@ describe('computeDecision', () => { const decision = computeDecision( liveClassification('investigation', 'codebase_understanding'), zeroPasserTable, - 'moonshotai/kimi-k2.7-code' + inc('moonshotai/kimi-k2.7-code') ); expect(decision).toMatchObject({ model: 'moonshotai/kimi-k2.7-code', @@ -395,7 +414,7 @@ describe('computeDecision', () => { const decision = computeDecision( liveClassification('planning_design', 'technical_planning'), solePasserTable, - 'thinkingmachines/inkling' + inc('thinkingmachines/inkling') ); expect(decision).toMatchObject({ model: 'thinkingmachines/inkling', @@ -433,7 +452,7 @@ describe('computeDecision', () => { const decision = computeDecision( liveClassification('debugging', 'root_cause_analysis'), modalTable, - 'moonshotai/kimi-k2.7-code' + inc('moonshotai/kimi-k2.7-code') ); expect(decision).toMatchObject({ model: 'moonshotai/kimi-k2.7-code', @@ -472,7 +491,7 @@ describe('computeDecision', () => { const decision = computeDecision( liveClassification('planning_design', 'architecture_design'), relabelTable, - 'anthropic/claude-sonnet-5' + inc('anthropic/claude-sonnet-5') ); expect(decision).toMatchObject({ model: 'thinkingmachines/inkling', @@ -510,7 +529,7 @@ describe('computeDecision', () => { const decision = computeDecision( liveClassification('planning_design', 'technical_planning'), solePasserTable, - 'thinkingmachines/inkling', + inc('thinkingmachines/inkling'), new Set(), 'best_accuracy' ); @@ -658,7 +677,7 @@ describe('computeDecision', () => { const decision = computeDecision( classification, visionTable, - 'text-only/chat', + inc('text-only/chat'), new Set(), 'cost_per_accuracy', { @@ -749,7 +768,7 @@ describe('computeDecision', () => { const decision = computeDecision( classification, sizedTable, - 'large/chat', + inc('large/chat'), new Set(), 'cost_per_accuracy', { @@ -847,7 +866,7 @@ describe('computeDecision', () => { const decision = computeDecision( classification, sizedTable, - 'b/chat', + inc('b/chat'), new Set(), 'cost_per_accuracy', { @@ -925,4 +944,205 @@ describe('computeDecision', () => { expect(emptyConstraints).toEqual(noConstraints); }); }); + + describe('exact-pair stickiness and variant emission', () => { + const variantTable: RoutingTable = { + ...table, + routes: { + 'implementation/code_generation': [ + { + model: 'shared/chat', + accuracy: 0.9, + avgCostUsd: 0.002, + meetsThreshold: true, + variant: 'instant', + }, + { + model: 'shared/chat', + accuracy: 0.85, + avgCostUsd: 0.005, + meetsThreshold: true, + variant: 'thinking', + }, + { + model: 'other/chat', + accuracy: 0.8, + avgCostUsd: 0.01, + meetsThreshold: true, + variant: null, + }, + ], + }, + }; + + it('keeps an exact-pair incumbent over a cheaper different variant of the same model', () => { + const decision = computeDecision( + classification, + variantTable, + inc('shared/chat', 'thinking') + ); + expect(decision).toMatchObject({ + model: 'shared/chat', + variant: 'thinking', + sticky: true, + }); + expect(decision).not.toHaveProperty('reasoningEffort'); + }); + + it('switches when the exact-pair incumbent was removed from the candidate set', () => { + const decision = computeDecision( + classification, + variantTable, + inc('shared/chat', 'gone-variant') + ); + expect(decision).toMatchObject({ + model: 'shared/chat', + variant: 'instant', + sticky: false, + switchReason: 'threshold', + }); + }); + + it('legacy model-only sticky matches a single candidate with that model', () => { + // Fresh pick is cheaper but not by more than switchCostFactor (3x), so + // the unique model-only sticky match is kept. + const singleVariantTable: RoutingTable = { + ...table, + routes: { + 'implementation/code_generation': [ + { + model: 'fresh/chat', + accuracy: 0.9, + avgCostUsd: 0.002, + meetsThreshold: true, + variant: null, + }, + { + model: 'only/chat', + accuracy: 0.85, + avgCostUsd: 0.005, + meetsThreshold: true, + variant: 'max', + }, + ], + }, + }; + const decision = computeDecision(classification, singleVariantTable, inc('only/chat', null)); + expect(decision).toMatchObject({ + model: 'only/chat', + variant: 'max', + sticky: true, + }); + }); + + it('legacy model-only sticky is ignored when two variants of the model are candidates', () => { + const decision = computeDecision(classification, variantTable, inc('shared/chat', null)); + // No unique model match → no incumbent → fresh pick + expect(decision).toMatchObject({ + model: 'shared/chat', + variant: 'instant', + sticky: false, + switchReason: null, + }); + }); + + it('emits variant without reasoningEffort for new-style candidates', () => { + const decision = computeDecision(classification, variantTable, null); + expect(decision).toEqual({ + model: 'shared/chat', + taskType: 'implementation', + subtaskType: 'code_generation', + source: 'benchmark', + tableVersion: 'run-1', + variant: 'instant', + sticky: false, + switchReason: null, + }); + }); + + it('emits reasoningEffort without variant for legacy candidates', () => { + const decision = computeDecision(classification, table, null, new Set(['cheap/chat'])); + expect(decision).toEqual({ + model: 'mid/chat', + taskType: 'implementation', + subtaskType: 'code_generation', + source: 'benchmark', + tableVersion: 'run-1', + reasoningEffort: 'medium', + sticky: false, + switchReason: null, + }); + }); + }); + + describe('custom fail-closed on inactive', () => { + const activeTable: RoutingTable = { + ...table, + routes: { + 'implementation/code_generation': [ + { + model: 'inactive/chat', + accuracy: 0.95, + avgCostUsd: 0.001, + meetsThreshold: true, + variant: 'max', + }, + { + model: 'active/chat', + accuracy: 0.85, + avgCostUsd: 0.002, + meetsThreshold: true, + variant: 'max', + }, + ], + }, + }; + + it('drops missing-row and isActive !== true candidates when failClosedOnInactive', () => { + const caps = makeCaps({ + 'inactive/chat': { inputModalities: [], isActive: false }, + 'active/chat': { inputModalities: [], isActive: true }, + }); + const decision = computeDecision( + classification, + activeTable, + null, + new Set(), + 'cost_per_accuracy', + { capabilityMap: caps, failClosedOnInactive: true } + ); + expect(decision).toMatchObject({ model: 'active/chat', variant: 'max' }); + }); + + it('returns null when every custom candidate is inactive or missing', () => { + const caps = makeCaps({ + 'inactive/chat': { inputModalities: [], isActive: false }, + }); + const decision = computeDecision( + classification, + activeTable, + null, + new Set(), + 'cost_per_accuracy', + { capabilityMap: caps, failClosedOnInactive: true } + ); + expect(decision).toBeNull(); + }); + + it('does not drop inactive candidates on the platform path', () => { + const caps = makeCaps({ + 'inactive/chat': { inputModalities: [], isActive: false }, + 'active/chat': { inputModalities: [], isActive: true }, + }); + const decision = computeDecision( + classification, + activeTable, + null, + new Set(), + 'cost_per_accuracy', + { capabilityMap: caps } + ); + expect(decision).toMatchObject({ model: 'inactive/chat' }); + }); + }); }); diff --git a/services/auto-routing/src/decision-engine.ts b/services/auto-routing/src/decision-engine.ts index 1cea4834a6..54ca907956 100644 --- a/services/auto-routing/src/decision-engine.ts +++ b/services/auto-routing/src/decision-engine.ts @@ -6,6 +6,7 @@ import { type AutoRoutingMode, type ClassifierOutput, type RankedCandidate, + type ReasoningEffort, type RoutingConstraints, type RoutingTable, } from '@kilocode/auto-routing-contracts'; @@ -21,6 +22,11 @@ import type { ModelCapabilities, ModelCapabilitiesMap } from './model-capabiliti // `model_stats.inputModalities` (`apps/web/src/lib/model-stats/sync-openrouter.ts:77,95,124`). export const ENFORCED_MODALITIES: ReadonlyArray = ['image', 'file']; +export type DecisionIncumbent = { + model: string; + variant: string | null; +}; + // Single source of the constraint policy shared by benchmark routing // (`applyCapabilityFilters`) and the coding-plan short-circuit in `decide.ts`. // Keeping these in one place stops the two paths from silently diverging when @@ -67,6 +73,58 @@ export function modelSatisfiesConstraints( ); } +function candidateVariant(candidate: RankedCandidate): string | null { + // Prefer canonical variant; legacy effort-only candidates use reasoningEffort + // as the variant key for exact-pair identity during rolling deploy. + if (candidate.variant !== undefined && candidate.variant !== null) { + return candidate.variant; + } + if (candidate.reasoningEffort !== undefined && candidate.reasoningEffort !== null) { + return candidate.reasoningEffort; + } + return null; +} + +function findIncumbentCandidate( + candidates: ReadonlyArray, + incumbent: DecisionIncumbent | null +): RankedCandidate | undefined { + if (incumbent === null) return undefined; + + // Exact-pair match when the sticky record carries a variant. + if (incumbent.variant !== null) { + return candidates.find( + c => c.model === incumbent.model && candidateVariant(c) === incumbent.variant + ); + } + + // Legacy model-only sticky: match only when exactly one current candidate + // has that model. Multiple variants of the same model → no incumbent. + const modelMatches = candidates.filter(c => c.model === incumbent.model); + if (modelMatches.length === 1) { + return modelMatches[0]; + } + return undefined; +} + +function decisionFieldsFromCandidate(candidate: RankedCandidate): { + variant?: string | null; + reasoningEffort?: ReasoningEffort | null; +} { + // New writers emit variant only. Legacy platform candidates keep + // reasoningEffort. Never both (contract rejects both-set). + if (candidate.variant !== undefined && candidate.variant !== null) { + return { variant: candidate.variant }; + } + if (candidate.variant === null) { + // Explicit null variant on a new-style candidate: emit variant null, + // no reasoningEffort. + return { variant: null }; + } + // Absent variant: legacy path. + return { reasoningEffort: candidate.reasoningEffort ?? null }; +} + function pickFreshCandidate( candidates: ReadonlyArray, mode: AutoRoutingMode @@ -150,6 +208,18 @@ function applyCapabilityFilters( return { filtered: maxContextFallback, reason: 'ok' }; } +function applyActiveFilter( + candidates: ReadonlyArray, + capabilityMap: ModelCapabilitiesMap | undefined, + failClosedOnInactive: boolean +): ReadonlyArray { + if (!failClosedOnInactive) return candidates; + return candidates.filter(c => { + const caps = capabilityMap?.get(c.model); + return caps != null && caps.isActive === true; + }); +} + // A route's accuracy is graded on 10 distinct benchmark cases // (datasets/decider-cases.ts, >=10 per taxonomy route; repetitions re-run the // same prompts and add no independent tasks). At n=10 the one-sided 95% Wilson @@ -167,15 +237,25 @@ function applyCapabilityFilters( // the single largest group, which sits at 27/30 and is unaffected). const STICKY_ACCURACY_TOLERANCE = 0.1; +function sameExactPair( + a: DecisionIncumbent, + b: { model: string; variant: string | null } +): boolean { + return a.model === b.model && a.variant === b.variant; +} + export function computeDecision( classification: ClassifierOutput, table: RoutingTable | null, - incumbentModel: string | null, + incumbent: DecisionIncumbent | null, deniedModelIds: ReadonlySet = new Set(), mode: AutoRoutingMode = DEFAULT_AUTO_ROUTING_MODE, options: { constraints?: RoutingConstraints | undefined; capabilityMap?: ModelCapabilitiesMap | undefined; + // True only for custom-pool tables: drop candidates whose capability row + // is missing or isActive !== true. Platform tables leave this unset. + failClosedOnInactive?: boolean; } = {} ): AutoRoutingDecision | null { if (!table) return null; @@ -185,8 +265,17 @@ export function computeDecision( ); if (!routeCandidates?.length) return null; - const { filtered: candidates, reason } = applyCapabilityFilters( + const activeFiltered = applyActiveFilter( routeCandidates, + options.capabilityMap, + options.failClosedOnInactive === true + ); + if (activeFiltered.length === 0) { + return null; + } + + const { filtered: candidates, reason } = applyCapabilityFilters( + activeFiltered, options.constraints, options.capabilityMap ); @@ -195,6 +284,7 @@ export function computeDecision( } const freshPick = pickFreshCandidate(candidates, mode); + const freshPair = { model: freshPick.model, variant: candidateVariant(freshPick) }; // Keep the session on its incumbent model when it is still good enough for // the current taxonomy route. A model switch discards the provider's prompt cache, @@ -205,8 +295,7 @@ export function computeDecision( // 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 incumbentCandidate = findIncumbentCandidate(candidates, incumbent); // Sticky eligibility, shared by the keep decision and the switchReason // telemetry below so the two can never disagree. best_accuracy keeps the // strict bar (that mode exists to buy accuracy); cost_per_accuracy keeps @@ -216,48 +305,64 @@ export function computeDecision( ? candidate.meetsThreshold : candidate.accuracy >= table.minAccuracy - STICKY_ACCURACY_TOLERANCE; const stickyIncumbent = - incumbent && - incumbentStickyEligible(incumbent) && - incumbent.model !== freshPick.model && + incumbentCandidate && + incumbentStickyEligible(incumbentCandidate) && + !sameExactPair( + { model: incumbentCandidate.model, variant: candidateVariant(incumbentCandidate) }, + freshPair + ) && ((mode === 'cost_per_accuracy' && - !(freshPick.avgCostUsd * table.switchCostFactor < incumbent.avgCostUsd)) || + !(freshPick.avgCostUsd * table.switchCostFactor < incumbentCandidate.avgCostUsd)) || (mode === 'best_accuracy' && - !(freshPick.accuracy - incumbent.accuracy > table.bestAccuracySwitchThreshold))); + !(freshPick.accuracy - incumbentCandidate.accuracy > table.bestAccuracySwitchThreshold))); - if (stickyIncumbent) { + if (stickyIncumbent && incumbentCandidate) { return { - model: incumbent.model, + model: incumbentCandidate.model, taskType: classification.taskType, subtaskType: classification.subtaskType, source: table.source, tableVersion: table.version, - reasoningEffort: incumbent.reasoningEffort ?? null, + ...decisionFieldsFromCandidate(incumbentCandidate), sticky: true, switchReason: null, }; } - const switched = incumbentModel !== null && incumbentModel !== freshPick.model; + // Legacy model-only sticky compares model only; exact-pair sticky compares + // both fields so two variants of the same model count as a switch. + const switched = + incumbent !== null && + (incumbent.variant === null + ? incumbent.model !== freshPick.model + : !sameExactPair(incumbent, freshPair)); + return { model: freshPick.model, taskType: classification.taskType, subtaskType: classification.subtaskType, source: table.source, tableVersion: table.version, - reasoningEffort: freshPick.reasoningEffort ?? null, + ...decisionFieldsFromCandidate(freshPick), sticky: false, // 'cost': the incumbent was eligible but the mode's switch condition // (cost factor / accuracy gap) made the fresh pick worth it; // 'capability': the modality/context filters ejected it from the route; // 'threshold': it is denied, off the route, or outside the accuracy band. - switchReason: !switched - ? null - : incumbent - ? incumbentStickyEligible(incumbent) - ? 'cost' - : 'threshold' - : routeCandidates.some(c => c.model === incumbentModel) - ? 'capability' - : 'threshold', + switchReason: (() => { + if (!switched || incumbent === null) return null; + if (incumbentCandidate) { + return incumbentStickyEligible(incumbentCandidate) ? 'cost' : 'threshold'; + } + // Exact-pair sticky: model still on the route with a different + // variant is a removed pair (threshold), not a capability eject. + // Model-only legacy sticky still uses model presence. + const stillOnRoute = routeCandidates.some( + c => + c.model === incumbent.model && + (incumbent.variant === null || candidateVariant(c) === incumbent.variant) + ); + return stillOnRoute ? 'capability' : 'threshold'; + })(), }; } diff --git a/services/auto-routing/src/index.test.ts b/services/auto-routing/src/index.test.ts index 2ef838dc67..aeb9d26d37 100644 --- a/services/auto-routing/src/index.test.ts +++ b/services/auto-routing/src/index.test.ts @@ -2,10 +2,13 @@ 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 type * as ModelCapabilitiesModule from './model-capabilities'; +import { clearCustomRoutingTableCache } from './decide'; import { app } from './index'; import { ClassifierRunError } from './model-classifier'; import type * as DbModule from '@kilocode/db'; import type * as ModelClassifierModule from './model-classifier'; +import type { AutoRoutingMode, EfficientModelPool } from '@kilocode/auto-routing-contracts'; const classifyNormalizedInput = vi.hoisted(() => vi.fn()); const getWorkerDb = vi.hoisted(() => vi.fn()); @@ -16,12 +19,23 @@ 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()); +const getModelCapabilitiesMock = vi.hoisted(() => vi.fn()); vi.mock('./model-classifier', async importOriginal => { const actual = await importOriginal(); return { ...actual, classifyNormalizedInput }; }); +vi.mock('./model-capabilities', async importOriginal => { + const actual = await importOriginal(); + getModelCapabilitiesMock.mockImplementation(actual.getModelCapabilities); + return { + ...actual, + getModelCapabilities: (...args: Parameters) => + getModelCapabilitiesMock(...args), + }; +}); + vi.mock('@kilocode/db', async importOriginal => { const actual = await importOriginal(); return { ...actual, getWorkerDb }; @@ -192,10 +206,14 @@ function decideRequest(payload: unknown) { } describe('auto routing worker', () => { - beforeEach(() => { + beforeEach(async () => { clearClassifierConfigCache(); clearRoutingTableCache(); clearModelCapabilitiesCache(); + const actualCaps = + await vi.importActual('./model-capabilities'); + getModelCapabilitiesMock.mockReset(); + getModelCapabilitiesMock.mockImplementation(actualCaps.getModelCapabilities); classifyNormalizedInput.mockReset(); classifyNormalizedInput.mockResolvedValue(mockClassifierResult); getWorkerDb.mockReset(); @@ -228,9 +246,26 @@ describe('auto routing worker', () => { modeConfigIdFromName.mockReset(); modeConfigIdFromName.mockImplementation((name: string) => name); modeConfigGet.mockReset(); - modeConfigGet.mockReturnValue({ - getMode: vi.fn(async () => null), - setMode: vi.fn(async () => undefined), + modeConfigGet.mockImplementation(() => { + let mode: AutoRoutingMode | null = null; + let pool: EfficientModelPool | null = null; + return { + getMode: vi.fn(async () => mode), + setMode: vi.fn(async (next: AutoRoutingMode | null) => { + mode = next; + }), + getPool: vi.fn(async () => pool), + setPool: vi.fn(async (next: EfficientModelPool | null) => { + pool = next; + }), + getSettings: vi.fn(async () => ({ mode, pool })), + setSettings: vi.fn( + async (settings: { mode: AutoRoutingMode | null; pool: EfficientModelPool | null }) => { + mode = settings.mode; + pool = settings.pool; + } + ), + }; }); benchmarkFetch.mockReset(); benchmarkFetch.mockImplementation(async (url: string) => { @@ -254,10 +289,12 @@ describe('auto routing worker', () => { cachePutEntry.mockResolvedValue(undefined); mockedFetch.mockReset(); globalThis.fetch = mockedFetch; + clearCustomRoutingTableCache(); }); afterEach(() => { globalThis.fetch = originalFetch; + clearCustomRoutingTableCache(); vi.restoreAllMocks(); }); @@ -640,12 +677,16 @@ describe('auto routing worker', () => { }); it('uses organization best-accuracy mode for auto routing decisions', async () => { - const modes = new Map([['org:org-1', 'best_accuracy']]); + const modes = new Map([['org:org-1', 'best_accuracy']]); modeConfigGet.mockImplementation((id: string) => ({ getMode: vi.fn(async () => modes.get(id) ?? null), - setMode: vi.fn(async (mode: string | null) => { + setMode: vi.fn(async (mode: AutoRoutingMode | null) => { modes.set(id, mode); }), + getPool: vi.fn(async () => null), + setPool: vi.fn(async () => undefined), + getSettings: vi.fn(async () => ({ mode: modes.get(id) ?? null, pool: null })), + setSettings: vi.fn(async () => undefined), })); benchmarkFetch.mockImplementation(async (url: string) => { if (String(url).includes('/admin/classifier-winner')) { @@ -694,12 +735,16 @@ describe('auto routing worker', () => { }); it('falls back to user auto routing mode when org mode is unset', async () => { - const modes = new Map([['user:user-1', 'best_accuracy']]); + const modes = new Map([['user:user-1', 'best_accuracy']]); modeConfigGet.mockImplementation((id: string) => ({ getMode: vi.fn(async () => modes.get(id) ?? null), - setMode: vi.fn(async (mode: string | null) => { + setMode: vi.fn(async (mode: AutoRoutingMode | null) => { modes.set(id, mode); }), + getPool: vi.fn(async () => null), + setPool: vi.fn(async () => undefined), + getSettings: vi.fn(async () => ({ mode: modes.get(id) ?? null, pool: null })), + setSettings: vi.fn(async () => undefined), })); benchmarkFetch.mockImplementation(async (url: string) => { if (String(url).includes('/admin/classifier-winner')) { @@ -748,13 +793,21 @@ describe('auto routing worker', () => { }); it('reads and updates owner routing mode through admin endpoints', async () => { - let storedMode: string | null = null; + let storedMode: AutoRoutingMode | null = null; modeConfigGet.mockImplementation(() => { return { getMode: vi.fn(async () => storedMode), - setMode: vi.fn(async (mode: string | null) => { + setMode: vi.fn(async (mode: AutoRoutingMode | null) => { storedMode = mode; }), + getPool: vi.fn(async () => null), + setPool: vi.fn(async () => undefined), + getSettings: vi.fn(async () => ({ mode: storedMode, pool: null })), + setSettings: vi.fn( + async (settings: { mode: AutoRoutingMode | null; pool: EfficientModelPool | null }) => { + storedMode = settings.mode; + } + ), }; }); @@ -789,6 +842,513 @@ describe('auto routing worker', () => { }); }); + describe('admin routing settings', () => { + const pool: EfficientModelPool = [{ model: 'custom/chat', variant: 'thinking' }]; + const statuses = [ + { entry: { model: 'custom/chat', variant: 'thinking' }, status: 'ready' as const }, + ]; + + function ownerSettingsStore(initial?: { + mode?: AutoRoutingMode | null; + pool?: EfficientModelPool | null; + }) { + let mode: AutoRoutingMode | null = initial?.mode ?? null; + let poolValue: EfficientModelPool | null = initial?.pool ?? null; + const setSettings = vi.fn( + async (settings: { mode: AutoRoutingMode | null; pool: EfficientModelPool | null }) => { + mode = settings.mode; + poolValue = settings.pool; + } + ); + modeConfigGet.mockImplementation(() => ({ + getMode: vi.fn(async () => mode), + setMode: vi.fn(async (next: AutoRoutingMode | null) => { + mode = next; + }), + getPool: vi.fn(async () => poolValue), + setPool: vi.fn(async (next: EfficientModelPool | null) => { + poolValue = next; + }), + getSettings: vi.fn(async () => ({ mode, pool: poolValue })), + setSettings, + })); + return { + get mode() { + return mode; + }, + get pool() { + return poolValue; + }, + setSettings, + }; + } + + it('GET without pool skips the status fetch', async () => { + ownerSettingsStore({ mode: 'best_accuracy', pool: null }); + const response = await request('/admin/routing-settings?ownerType=user&ownerId=user-1', { + headers: { authorization: 'Bearer classifier-token' }, + }); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + configuredMode: 'best_accuracy', + configuredPool: null, + poolStatuses: [], + }); + expect(benchmarkFetch).not.toHaveBeenCalledWith( + expect.stringContaining('/admin/profiles/status'), + expect.anything() + ); + }); + + it('GET with pool returns statuses from the benchmark worker', async () => { + ownerSettingsStore({ mode: null, pool }); + benchmarkFetch.mockImplementation(async (url: string) => { + if (String(url).includes('/admin/profiles/status')) { + return { + ok: true, + status: 200, + json: async () => ({ statuses }), + }; + } + return { + ok: true, + status: 200, + json: async () => ({ + table: benchmarkRoutingTable, + publishedAt: benchmarkRoutingTable.generatedAt, + }), + }; + }); + + const response = await request('/admin/routing-settings?ownerType=user&ownerId=user-1', { + headers: { authorization: 'Bearer classifier-token' }, + }); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + configuredPool: pool, + poolStatuses: statuses, + }); + expect(benchmarkFetch).toHaveBeenCalledWith( + expect.stringContaining('/admin/profiles/status'), + expect.objectContaining({ method: 'POST' }) + ); + }); + + it('GET returns 502 when status fetch fails for a configured pool', async () => { + ownerSettingsStore({ mode: null, pool }); + benchmarkFetch.mockImplementation(async (url: string) => { + if (String(url).includes('/admin/profiles/status')) { + return { ok: false, status: 500, text: async () => 'boom', json: async () => ({}) }; + } + return { + ok: true, + status: 200, + json: async () => ({ + table: benchmarkRoutingTable, + publishedAt: benchmarkRoutingTable.generatedAt, + }), + }; + }); + + const response = await request('/admin/routing-settings?ownerType=user&ownerId=user-1', { + headers: { authorization: 'Bearer classifier-token' }, + }); + expect(response.status).toBe(502); + await expect(response.json()).resolves.toEqual({ + error: 'Failed to load pool profile statuses', + }); + }); + + it('PUT persists only after successful register and reuses register statuses', async () => { + const store = ownerSettingsStore(); + benchmarkFetch.mockImplementation(async (url: string) => { + if (String(url).includes('/admin/profiles/register')) { + return { + ok: true, + status: 200, + json: async () => ({ statuses }), + }; + } + throw new Error(`unexpected benchmark call ${url}`); + }); + + const response = await request('/admin/routing-settings', { + method: 'PUT', + headers: { + authorization: 'Bearer classifier-token', + 'content-type': 'application/json', + }, + body: JSON.stringify({ + ownerType: 'user', + ownerId: 'user-1', + mode: 'best_accuracy', + pool, + }), + }); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + configuredMode: 'best_accuracy', + configuredPool: pool, + poolStatuses: statuses, + }); + expect(store.setSettings).toHaveBeenCalledWith({ mode: 'best_accuracy', pool }); + expect(store.pool).toEqual(pool); + expect(benchmarkFetch).toHaveBeenCalledWith( + expect.stringContaining('/admin/profiles/register'), + expect.objectContaining({ method: 'POST' }) + ); + expect(benchmarkFetch).not.toHaveBeenCalledWith( + expect.stringContaining('/admin/profiles/status'), + expect.anything() + ); + }); + + it('PUT maps register 429 to 429 with retryAt and persists nothing', async () => { + const store = ownerSettingsStore({ mode: 'cost_per_accuracy', pool: null }); + const retryAt = '2026-07-29T00:00:00.000Z'; + benchmarkFetch.mockImplementation(async (url: string) => { + if (String(url).includes('/admin/profiles/register')) { + return { + ok: false, + status: 429, + json: async () => ({ error: 'quota exceeded', retryAt }), + text: async () => 'quota exceeded', + }; + } + throw new Error(`unexpected benchmark call ${url}`); + }); + + const response = await request('/admin/routing-settings', { + method: 'PUT', + headers: { + authorization: 'Bearer classifier-token', + 'content-type': 'application/json', + }, + body: JSON.stringify({ + ownerType: 'user', + ownerId: 'user-1', + mode: 'best_accuracy', + pool, + }), + }); + expect(response.status).toBe(429); + await expect(response.json()).resolves.toEqual({ + error: 'quota exceeded', + retryAt, + }); + expect(store.setSettings).not.toHaveBeenCalled(); + expect(store.pool).toBeNull(); + expect(store.mode).toBe('cost_per_accuracy'); + }); + + it('PUT null pool clears without calling register', async () => { + const store = ownerSettingsStore({ mode: 'best_accuracy', pool }); + const response = await request('/admin/routing-settings', { + method: 'PUT', + headers: { + authorization: 'Bearer classifier-token', + 'content-type': 'application/json', + }, + body: JSON.stringify({ + ownerType: 'user', + ownerId: 'user-1', + mode: null, + pool: null, + }), + }); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + configuredMode: null, + configuredPool: null, + poolStatuses: [], + }); + expect(store.setSettings).toHaveBeenCalledWith({ mode: null, pool: null }); + expect(benchmarkFetch).not.toHaveBeenCalledWith( + expect.stringContaining('/admin/profiles/register'), + expect.anything() + ); + }); + + it('PUT forwards retryEntries on register and persists only after success', async () => { + const store = ownerSettingsStore(); + const retryEntries = [{ model: 'custom/chat', variant: 'thinking' as const }]; + let registerBody: unknown; + benchmarkFetch.mockImplementation(async (url: string, init?: RequestInit) => { + if (String(url).includes('/admin/profiles/register')) { + const body = init?.body; + registerBody = + typeof body === 'string' + ? JSON.parse(body) + : body instanceof Uint8Array + ? JSON.parse(new TextDecoder().decode(body)) + : undefined; + return { + ok: true, + status: 200, + json: async () => ({ statuses }), + }; + } + throw new Error(`unexpected benchmark call ${url}`); + }); + + const response = await request('/admin/routing-settings', { + method: 'PUT', + headers: { + authorization: 'Bearer classifier-token', + 'content-type': 'application/json', + }, + body: JSON.stringify({ + ownerType: 'user', + ownerId: 'user-1', + mode: 'best_accuracy', + pool, + retryEntries, + }), + }); + expect(response.status).toBe(200); + expect(registerBody).toEqual({ + ownerType: 'user', + ownerId: 'user-1', + entries: pool, + retryEntries, + }); + expect(store.setSettings).toHaveBeenCalledTimes(1); + expect(store.setSettings).toHaveBeenCalledWith({ mode: 'best_accuracy', pool }); + expect(store.pool).toEqual(pool); + }); + }); + + describe('custom pool decide path', () => { + const pool: EfficientModelPool = [{ model: 'pool/chat', variant: 'thinking' }]; + const customTable = { + version: 'custom-1', + generatedAt: '2026-06-12T00:00:00.000Z', + minAccuracy: 0.7, + switchCostFactor: 3, + bestAccuracySwitchThreshold: 0.05, + source: 'benchmark' as const, + routes: { + 'implementation/feature_development': [ + { + model: 'pool/chat', + accuracy: 0.92, + avgCostUsd: 0.001, + meetsThreshold: true, + variant: 'thinking', + }, + ], + }, + }; + + function setOwnerPool(nextPool: EfficientModelPool | null) { + modeConfigGet.mockImplementation(() => ({ + getMode: vi.fn(async () => null), + setMode: vi.fn(async () => undefined), + getPool: vi.fn(async () => nextPool), + setPool: vi.fn(async () => undefined), + getSettings: vi.fn(async () => ({ mode: null, pool: nextPool })), + setSettings: vi.fn(async () => undefined), + })); + } + + it('uses the custom routing table when a pool is configured', async () => { + setOwnerPool(pool); + dbWhereCaps.mockResolvedValue([ + { + openrouterId: 'pool/chat', + inputModalities: [], + contextLength: 1_000_000, + isActive: true, + }, + ]); + benchmarkFetch.mockImplementation(async (url: string) => { + if (String(url).includes('/admin/custom-routing-table')) { + return { + ok: true, + status: 200, + json: async () => ({ table: customTable }), + }; + } + if (String(url).includes('/admin/classifier-winner')) { + return { ok: true, status: 200, json: async () => ({ winner: null }) }; + } + // Platform table must not be used for the decision. + return { + ok: true, + status: 200, + json: async () => ({ + table: benchmarkRoutingTable, + publishedAt: benchmarkRoutingTable.generatedAt, + }), + }; + }); + + const response = await decideRequest(mirrorPayload()); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + decision: { + model: 'pool/chat', + variant: 'thinking', + source: 'benchmark', + tableVersion: 'custom-1', + }, + }); + expect(benchmarkFetch).toHaveBeenCalledWith( + expect.stringContaining('/admin/custom-routing-table'), + expect.objectContaining({ method: 'POST' }) + ); + expect(cachePutEntry).toHaveBeenCalledWith( + 'sticky', + expect.objectContaining({ + model: 'pool/chat', + variant: 'thinking', + }) + ); + }); + + it('returns a null decision when custom table lookup fails', async () => { + setOwnerPool(pool); + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}); + benchmarkFetch.mockImplementation(async (url: string) => { + if (String(url).includes('/admin/custom-routing-table')) { + return { ok: false, status: 500, text: async () => 'fail', json: async () => ({}) }; + } + if (String(url).includes('/admin/classifier-winner')) { + return { ok: true, status: 200, json: async () => ({ winner: null }) }; + } + return { + ok: true, + status: 200, + json: async () => ({ + table: benchmarkRoutingTable, + publishedAt: benchmarkRoutingTable.generatedAt, + }), + }; + }); + + const response = await decideRequest(mirrorPayload()); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ decision: null }); + expect(warn).toHaveBeenCalledWith( + expect.stringContaining('auto_routing_custom_table_read_failed') + ); + warn.mockRestore(); + }); + + it('keeps the null-pool path on the platform cached table', async () => { + setOwnerPool(null); + const response = await decideRequest(mirrorPayload()); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + decision: { + model: 'google/gemini-2.5-flash-lite', + source: 'benchmark', + tableVersion: 'bench-run-1', + }, + }); + expect(benchmarkFetch).not.toHaveBeenCalledWith( + expect.stringContaining('/admin/custom-routing-table'), + expect.anything() + ); + }); + + it('loads capabilities once with pool model ids for constraints+pool', async () => { + setOwnerPool(pool); + dbWhereCaps.mockResolvedValue([ + { + openrouterId: 'pool/chat', + inputModalities: ['image'], + contextLength: 1_000_000, + isActive: true, + }, + ]); + benchmarkFetch.mockImplementation(async (url: string) => { + if (String(url).includes('/admin/custom-routing-table')) { + return { + ok: true, + status: 200, + json: async () => ({ table: customTable }), + }; + } + if (String(url).includes('/admin/classifier-winner')) { + return { ok: true, status: 200, json: async () => ({ winner: null }) }; + } + return { + ok: true, + status: 200, + json: async () => ({ + table: benchmarkRoutingTable, + publishedAt: benchmarkRoutingTable.generatedAt, + }), + }; + }); + + const response = await decideRequest( + mirrorPayload({ constraints: { requiredInputModalities: ['image'] } }) + ); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + decision: { + model: 'pool/chat', + variant: 'thinking', + source: 'benchmark', + }, + }); + expect(getModelCapabilitiesMock).toHaveBeenCalledTimes(1); + expect(getModelCapabilitiesMock).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ + additionalModelIds: ['pool/chat'], + }) + ); + }); + + it('takes coding-plan short-circuit from the single pool-aware capability load', async () => { + setOwnerPool(pool); + 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 + ); + dbWhereCaps.mockResolvedValue([ + { + openrouterId: 'minimax/minimax-m3', + inputModalities: ['image'], + contextLength: 1_000_000, + isActive: true, + }, + { + openrouterId: 'pool/chat', + inputModalities: ['image'], + contextLength: 1_000_000, + isActive: true, + }, + ]); + + const response = await decideRequest( + mirrorPayload({ constraints: { requiredInputModalities: ['image'] } }) + ); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + decision: { model: 'minimax/minimax-m3', source: 'coding_plan_default' }, + }); + expect(classifyNormalizedInput).not.toHaveBeenCalled(); + expect(getModelCapabilitiesMock).toHaveBeenCalledTimes(1); + expect(getModelCapabilitiesMock).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ + codingPlanModelId: 'minimax/minimax-m3', + additionalModelIds: ['pool/chat'], + }) + ); + }); + }); + it('filters denied routing-policy models from the full decide path', async () => { vi.spyOn(Math, 'random').mockReturnValue(0); @@ -956,6 +1516,7 @@ describe('auto routing worker', () => { expect(cachePutEntry).toHaveBeenCalledTimes(1); expect(cachePutEntry).toHaveBeenCalledWith('sticky', { model: expect.any(String), + variant: null, routeKey: 'implementation/feature_development', }); expect(writeDataPoint).toHaveBeenCalledWith( diff --git a/services/auto-routing/src/index.ts b/services/auto-routing/src/index.ts index 6c16e2adbc..0aa3d44b92 100644 --- a/services/auto-routing/src/index.ts +++ b/services/auto-routing/src/index.ts @@ -4,6 +4,7 @@ import { authMiddleware } from './auth'; import { classifierAnalyticsHandler } from './admin-classifier-analytics'; import { getClassifierModelHandler, putClassifierModelHandler } from './admin-classifier-model'; import { getRoutingModeHandler, putRoutingModeHandler } from './admin-routing-mode'; +import { getRoutingSettingsHandler, putRoutingSettingsHandler } from './admin-routing-settings'; import { decideHandler } from './decide'; import type { HonoEnv } from './hono-env'; @@ -20,6 +21,8 @@ app.get('/admin/classifier-model', getClassifierModelHandler); app.put('/admin/classifier-model', putClassifierModelHandler); app.get('/admin/routing-mode', getRoutingModeHandler); app.put('/admin/routing-mode', putRoutingModeHandler); +app.get('/admin/routing-settings', getRoutingSettingsHandler); +app.put('/admin/routing-settings', putRoutingSettingsHandler); app.get('/admin/classifier-analytics', classifierAnalyticsHandler); app.post('/decide', decideHandler); diff --git a/services/auto-routing/src/model-capabilities.test.ts b/services/auto-routing/src/model-capabilities.test.ts index 9933e1ecfd..d690c824a3 100644 --- a/services/auto-routing/src/model-capabilities.test.ts +++ b/services/auto-routing/src/model-capabilities.test.ts @@ -168,7 +168,7 @@ describe('getModelCapabilities', () => { await getModelCapabilities(env); - expect(put).toHaveBeenCalledWith('model_capabilities_v1', expect.stringContaining('"a/chat"'), { + expect(put).toHaveBeenCalledWith('model_capabilities_v2', expect.stringContaining('"a/chat"'), { expirationTtl: 3600, }); }); @@ -198,11 +198,11 @@ describe('getModelCapabilities', () => { const result = await getModelCapabilities(env); expect(result.size).toBe(0); - // The model_capabilities_v1 key is never written; the routing-table + // The model_capabilities_v2 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' + (call: unknown[]) => call[0] === 'model_capabilities_v2' ); expect(capabilityPuts).toEqual([]); expect(warn).toHaveBeenCalled(); @@ -340,13 +340,13 @@ describe('getModelCapabilities', () => { 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 + // value for kvReadThrough is null and the model_capabilities_v2 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' + (call: unknown[]) => call[0] === 'model_capabilities_v2' ); expect(capabilityPutsBefore).toEqual([]); @@ -355,7 +355,7 @@ describe('getModelCapabilities', () => { expect(second.size).toBe(0); expect(get).toHaveBeenCalledTimes(2); const capabilityPutsAfter = put.mock.calls.filter( - (call: unknown[]) => call[0] === 'model_capabilities_v1' + (call: unknown[]) => call[0] === 'model_capabilities_v2' ); expect(capabilityPutsAfter).toEqual([]); @@ -372,7 +372,7 @@ describe('getModelCapabilities', () => { 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' + call => call[0] === 'model_capabilities_v2' ); expect(capabilityPutsEmpty).toHaveLength(1); expect(JSON.parse(capabilityPutsEmpty[0][1] as unknown as string)).toEqual({}); @@ -399,4 +399,58 @@ describe('getModelCapabilities', () => { const result = await getModelCapabilities(env); expect(result.size).toBe(0); }); + + it('selects isActive from model_stats and normalizes absent KV isActive to null', async () => { + dbWhere.mockImplementation(() => + Promise.resolve([ + { + openrouterId: 'a/chat', + inputModalities: ['text'], + contextLength: 8192, + isActive: true, + }, + { + openrouterId: 'b/chat', + inputModalities: ['text'], + contextLength: 4096, + isActive: false, + }, + ]) + ); + const env = makeEnv(null); + const result = await getModelCapabilities(env); + expect(result.get('a/chat')?.isActive).toBe(true); + expect(result.get('b/chat')?.isActive).toBe(false); + + clearModelCapabilitiesCache(); + const envFromKv = makeEnv( + JSON.stringify({ + 'a/chat': { inputModalities: ['text'], contextLength: 8192 }, + 'b/chat': { inputModalities: ['text'], contextLength: 4096, isActive: null }, + }) + ); + const fromKv = await getModelCapabilities(envFromKv); + expect(fromKv.get('a/chat')?.isActive).toBeNull(); + expect(fromKv.get('b/chat')?.isActive).toBeNull(); + }); + + it('unions additionalModelIds into the queried id set', async () => { + dbWhere.mockImplementation(() => + Promise.resolve([ + { openrouterId: 'a/chat', inputModalities: ['text'], contextLength: 8192, isActive: true }, + { + openrouterId: 'pool/extra', + inputModalities: ['text'], + contextLength: 4096, + isActive: true, + }, + ]) + ); + const env = makeEnv(null); + const result = await getModelCapabilities(env, { + additionalModelIds: ['pool/extra', 'a/chat'], + }); + expect(result.get('pool/extra')?.contextLength).toBe(4096); + expect(result.get('a/chat')?.isActive).toBe(true); + }); }); diff --git a/services/auto-routing/src/model-capabilities.ts b/services/auto-routing/src/model-capabilities.ts index d04dcef7a3..892826d8f6 100644 --- a/services/auto-routing/src/model-capabilities.ts +++ b/services/auto-routing/src/model-capabilities.ts @@ -8,9 +8,11 @@ import { getRoutingTable } from './routing-table'; // 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. +// `isActive` is the model_stats soft-active flag; null when unknown/absent. export type ModelCapabilities = { inputModalities: ReadonlySet; contextLength: number | null; + isActive: boolean | null; }; // An empty Map signals "no capability data" to callers: a request carrying @@ -46,14 +48,16 @@ function foldModalities(raw: ReadonlyArray | null | undefined): Set; function isCacheValue(value: unknown): value is ModelCapabilitiesCacheValue { @@ -79,9 +83,17 @@ function isCacheValue(value: unknown): value is ModelCapabilitiesCacheValue { 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 }; + const row = entry as { + inputModalities?: unknown; + contextLength?: unknown; + isActive?: unknown; + }; if (!Array.isArray(row.inputModalities)) return false; if (row.contextLength !== null && typeof row.contextLength !== 'number') return false; + // Accept boolean, null, or absent (absent → null on merge). + if (row.isActive !== undefined && row.isActive !== null && typeof row.isActive !== 'boolean') { + return false; + } } return true; } @@ -97,6 +109,7 @@ async function queryModelCapabilities( openrouterId: modelStats.openrouterId, inputModalities: modelStats.inputModalities, contextLength: modelStats.contextLength, + isActive: modelStats.isActive, }) .from(modelStats) .where(inArray(modelStats.openrouterId, modelIds as string[])); @@ -106,6 +119,7 @@ async function queryModelCapabilities( out[row.openrouterId] = { inputModalities: Array.isArray(row.inputModalities) ? row.inputModalities : [], contextLength: typeof row.contextLength === 'number' ? row.contextLength : null, + isActive: typeof row.isActive === 'boolean' ? row.isActive : null, }; } return out; @@ -124,6 +138,7 @@ function mergeInto( target.set(modelId, { inputModalities: foldModalities(row.inputModalities), contextLength: row.contextLength, + isActive: row.isActive ?? null, }); } } @@ -153,7 +168,16 @@ async function loadAll(env: ModelCapabilitiesEnv): Promise; + } = {} ): Promise { const load = async (): Promise> => { // We derive the id set inside the module so the caller (decide.ts) does @@ -208,6 +234,11 @@ export async function getModelCapabilities( if (options.codingPlanModelId) { ids.add(options.codingPlanModelId); } + if (options.additionalModelIds) { + for (const id of options.additionalModelIds) { + ids.add(id); + } + } const idList = Array.from(ids); if (idList.length === 0) { return new Map(); diff --git a/services/auto-routing/src/routing-mode.test.ts b/services/auto-routing/src/routing-mode.test.ts index f7014fbfe7..4fdc522630 100644 --- a/services/auto-routing/src/routing-mode.test.ts +++ b/services/auto-routing/src/routing-mode.test.ts @@ -2,28 +2,60 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { AutoRoutingModeConfigDO, getConfiguredAutoRoutingMode, + getConfiguredAutoRoutingSettings, getAutoRoutingMode, + getEffectiveAutoRoutingSettings, setAutoRoutingMode, + setAutoRoutingSettings, } from './routing-mode'; -import type { AutoRoutingMode } from '@kilocode/auto-routing-contracts'; +import type { AutoRoutingMode, EfficientModelPool } from '@kilocode/auto-routing-contracts'; +import type { AutoRoutingOwnerSettings } from './routing-mode'; type ModeStub = { getMode: ReturnType Promise>>; setMode: ReturnType Promise>>; + getPool: ReturnType Promise>>; + setPool: ReturnType Promise>>; + getSettings: ReturnType Promise>>; + setSettings: ReturnType Promise>>; }; -function makeEnv(initialModes: Record = {}) { - const modes = new Map(Object.entries(initialModes)); +const SAMPLE_POOL: EfficientModelPool = [ + { model: 'a/chat', variant: 'thinking' }, + { model: 'b/chat', variant: null }, +]; + +function makeEnv( + initial: Record = {} +) { + const modes = new Map(); + const pools = new Map(); + for (const [key, value] of Object.entries(initial)) { + modes.set(key, value.mode ?? null); + pools.set(key, value.pool ?? null); + } const stubs = new Map(); const idFromName = vi.fn((name: string) => name); const get = vi.fn((id: string) => { const existing = stubs.get(id); if (existing) return existing; - const stub = { + const stub: ModeStub = { getMode: vi.fn(async () => modes.get(id) ?? null), setMode: vi.fn(async (mode: AutoRoutingMode | null) => { modes.set(id, mode); }), + getPool: vi.fn(async () => pools.get(id) ?? null), + setPool: vi.fn(async (pool: EfficientModelPool | null) => { + pools.set(id, pool); + }), + getSettings: vi.fn(async () => ({ + mode: modes.get(id) ?? null, + pool: pools.get(id) ?? null, + })), + setSettings: vi.fn(async (settings: AutoRoutingOwnerSettings) => { + modes.set(id, settings.mode); + pools.set(id, settings.pool); + }), }; stubs.set(id, stub); return stub; @@ -35,22 +67,72 @@ function makeEnv(initialModes: Record = {}) { }, } as unknown as Pick; - return { env, modes, stubs, idFromName, get }; + return { env, modes, pools, stubs, idFromName, get }; +} + +function applyPut( + target: Map, + keyOrEntries: string | Record, + value?: unknown +) { + if (typeof keyOrEntries === 'string') { + target.set(keyOrEntries, value); + return; + } + for (const [key, entryValue] of Object.entries(keyOrEntries)) { + target.set(key, entryValue); + } +} + +function applyDelete(target: Map, keyOrKeys: string | string[]) { + if (typeof keyOrKeys === 'string') { + target.delete(keyOrKeys); + return; + } + for (const key of keyOrKeys) { + target.delete(key); + } } function createFakeStorage() { const entries = new Map(); - return { + const storage = { entries, get: async (key: string) => entries.get(key), - put: async (key: string, value: unknown) => { - entries.set(key, value); + put: async (keyOrEntries: string | Record, value?: unknown) => { + applyPut(entries, keyOrEntries, value); + }, + delete: async (keyOrKeys: string | string[]) => { + applyDelete(entries, keyOrKeys); }, - delete: async (key: string) => { - entries.delete(key); + transaction: async ( + callback: (txn: { + put: (keyOrEntries: string | Record, value?: unknown) => Promise; + delete: (keyOrKeys: string | string[]) => Promise; + get: (key: string) => Promise; + }) => Promise + ): Promise => { + const shadow = new Map(entries); + const txn = { + get: async (key: string) => shadow.get(key), + put: async (keyOrEntries: string | Record, value?: unknown) => { + applyPut(shadow, keyOrEntries, value); + }, + delete: async (keyOrKeys: string | string[]) => { + applyDelete(shadow, keyOrKeys); + }, + }; + const result = await callback(txn); + entries.clear(); + for (const [key, value] of shadow) { + entries.set(key, value); + } + return result; }, }; + + return storage; } function createModeDO() { @@ -76,6 +158,115 @@ describe('AutoRoutingModeConfigDO', () => { await modeDO.setMode(null); expect(storage.entries.has('mode')).toBe(false); }); + + it('round-trips pool get/set/clear and rejects malformed stored pools', async () => { + const { modeDO, storage } = createModeDO(); + + await expect(modeDO.getPool()).resolves.toBeNull(); + await modeDO.setPool(SAMPLE_POOL); + await expect(modeDO.getPool()).resolves.toEqual(SAMPLE_POOL); + + storage.entries.set('pool', [{ model: '', variant: null }]); + await expect(modeDO.getPool()).resolves.toBeNull(); + + await modeDO.setPool(null); + expect(storage.entries.has('pool')).toBe(false); + }); + + it('reads and writes combined settings in one RPC shape', async () => { + const { modeDO } = createModeDO(); + + await expect(modeDO.getSettings()).resolves.toEqual({ mode: null, pool: null }); + await modeDO.setSettings({ mode: 'best_accuracy', pool: SAMPLE_POOL }); + await expect(modeDO.getSettings()).resolves.toEqual({ + mode: 'best_accuracy', + pool: SAMPLE_POOL, + }); + await modeDO.setSettings({ mode: null, pool: null }); + await expect(modeDO.getSettings()).resolves.toEqual({ mode: null, pool: null }); + }); + + it('commits mode and pool via one multi-key put and paired multi-key delete', async () => { + const { modeDO, storage } = createModeDO(); + const putCalls: unknown[] = []; + const deleteCalls: unknown[] = []; + const originalTransaction = storage.transaction.bind(storage); + const transactionSpy = vi.spyOn(storage, 'transaction').mockImplementation(async callback => + originalTransaction(async txn => { + const put = txn.put.bind(txn); + const del = txn.delete.bind(txn); + txn.put = async (keyOrEntries, value?) => { + putCalls.push(keyOrEntries); + return put(keyOrEntries, value); + }; + txn.delete = async keyOrKeys => { + deleteCalls.push(keyOrKeys); + return del(keyOrKeys); + }; + return callback(txn); + }) + ); + + await modeDO.setSettings({ mode: 'best_accuracy', pool: SAMPLE_POOL }); + expect(transactionSpy).toHaveBeenCalledTimes(1); + expect(putCalls).toEqual([{ mode: 'best_accuracy', pool: SAMPLE_POOL }]); + expect(deleteCalls).toEqual([]); + + putCalls.length = 0; + deleteCalls.length = 0; + transactionSpy.mockClear(); + await modeDO.setSettings({ mode: null, pool: null }); + expect(transactionSpy).toHaveBeenCalledTimes(1); + expect(deleteCalls).toEqual([['mode', 'pool']]); + expect(putCalls).toEqual([]); + + putCalls.length = 0; + deleteCalls.length = 0; + transactionSpy.mockClear(); + await modeDO.setSettings({ mode: 'cost_per_accuracy', pool: null }); + expect(transactionSpy).toHaveBeenCalledTimes(1); + expect(putCalls).toEqual([{ mode: 'cost_per_accuracy' }]); + expect(deleteCalls).toEqual([['pool']]); + }); + + it('rolls back mode and pool when a transactional delete fails', async () => { + const { modeDO, storage } = createModeDO(); + await modeDO.setSettings({ mode: 'best_accuracy', pool: SAMPLE_POOL }); + + const originalTransaction = storage.transaction.bind(storage); + vi.spyOn(storage, 'transaction').mockImplementation(async callback => + originalTransaction(async txn => { + txn.delete = async () => { + throw new Error('delete failed'); + }; + return callback(txn); + }) + ); + + await expect(modeDO.setSettings({ mode: 'cost_per_accuracy', pool: null })).rejects.toThrow( + 'delete failed' + ); + expect(storage.entries.get('mode')).toBe('best_accuracy'); + expect(storage.entries.get('pool')).toEqual(SAMPLE_POOL); + await expect(modeDO.getSettings()).resolves.toEqual({ + mode: 'best_accuracy', + pool: SAMPLE_POOL, + }); + }); + + it('atomically commits a mixed mode write and pool delete', async () => { + const { modeDO, storage } = createModeDO(); + await modeDO.setSettings({ mode: 'best_accuracy', pool: SAMPLE_POOL }); + + await modeDO.setSettings({ mode: 'cost_per_accuracy', pool: null }); + + expect(storage.entries.get('mode')).toBe('cost_per_accuracy'); + expect(storage.entries.has('pool')).toBe(false); + await expect(modeDO.getSettings()).resolves.toEqual({ + mode: 'cost_per_accuracy', + pool: null, + }); + }); }); describe('auto routing mode config', () => { @@ -97,19 +288,19 @@ describe('auto routing mode config', () => { it('uses organization mode before user mode', async () => { const { env, idFromName } = makeEnv({ - 'user:user-1': 'best_accuracy', - 'org:org-1': 'cost_per_accuracy', + 'user:user-1': { mode: 'best_accuracy' }, + 'org:org-1': { mode: 'cost_per_accuracy' }, }); await expect( getAutoRoutingMode(env, { userId: 'user-1', organizationId: 'org-1' }) ).resolves.toBe('cost_per_accuracy'); - expect(idFromName).toHaveBeenNthCalledWith(1, 'org:org-1'); + expect(idFromName).toHaveBeenCalledWith('org:org-1'); }); it('falls back to user mode when organization mode is absent', async () => { const { env } = makeEnv({ - 'user:user-1': 'best_accuracy', + 'user:user-1': { mode: 'best_accuracy' }, }); await expect( @@ -119,7 +310,7 @@ describe('auto routing mode config', () => { it('reads the owner object on every lookup instead of serving a stale module value', async () => { const { env, modes, stubs } = makeEnv({ - 'user:user-1': 'best_accuracy', + 'user:user-1': { mode: 'best_accuracy' }, }); await expect( @@ -159,4 +350,89 @@ describe('auto routing mode config', () => { expect.stringContaining('auto_routing_config_read_failed') ); }); + + it('resolves mode and pool independently (org pool + personal mode)', async () => { + const { env } = makeEnv({ + 'user:user-1': { mode: 'best_accuracy', pool: null }, + 'org:org-1': { mode: null, pool: SAMPLE_POOL }, + }); + + await expect( + getEffectiveAutoRoutingSettings(env, { userId: 'user-1', organizationId: 'org-1' }) + ).resolves.toEqual({ + mode: 'best_accuracy', + pool: SAMPLE_POOL, + }); + }); + + it('resolves mode and pool independently (org mode + personal pool)', async () => { + const personalPool: EfficientModelPool = [{ model: 'personal/chat', variant: 'max' }]; + const { env } = makeEnv({ + 'user:user-1': { mode: null, pool: personalPool }, + 'org:org-1': { mode: 'best_accuracy', pool: null }, + }); + + await expect( + getEffectiveAutoRoutingSettings(env, { userId: 'user-1', organizationId: 'org-1' }) + ).resolves.toEqual({ + mode: 'best_accuracy', + pool: personalPool, + }); + }); + + it('lets org pool override personal pool and clearing restore inheritance', async () => { + const personalPool: EfficientModelPool = [{ model: 'personal/chat', variant: null }]; + const orgPool: EfficientModelPool = [{ model: 'org/chat', variant: 'thinking' }]; + const { env, modes, pools } = makeEnv({ + 'user:user-1': { mode: 'cost_per_accuracy', pool: personalPool }, + 'org:org-1': { mode: null, pool: orgPool }, + }); + + await expect( + getEffectiveAutoRoutingSettings(env, { userId: 'user-1', organizationId: 'org-1' }) + ).resolves.toEqual({ mode: 'cost_per_accuracy', pool: orgPool }); + + pools.set('org:org-1', null); + modes.set('org:org-1', null); + + await expect( + getEffectiveAutoRoutingSettings(env, { userId: 'user-1', organizationId: 'org-1' }) + ).resolves.toEqual({ mode: 'cost_per_accuracy', pool: personalPool }); + }); + + it('writes combined settings through the owner stub', async () => { + const { env, modes, pools, stubs } = makeEnv(); + + await setAutoRoutingSettings( + env, + { ownerType: 'user', ownerId: 'user-1' }, + { mode: 'best_accuracy', pool: SAMPLE_POOL } + ); + expect(modes.get('user:user-1')).toBe('best_accuracy'); + expect(pools.get('user:user-1')).toEqual(SAMPLE_POOL); + expect(stubs.get('user:user-1')?.setSettings).toHaveBeenCalledWith({ + mode: 'best_accuracy', + pool: SAMPLE_POOL, + }); + + await expect( + getConfiguredAutoRoutingSettings(env, { ownerType: 'user', ownerId: 'user-1' }) + ).resolves.toEqual({ mode: 'best_accuracy', pool: SAMPLE_POOL }); + }); + + it('degrades pool to null when configured settings read fails', async () => { + const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {}); + const { env, stubs } = makeEnv({ + 'user:user-1': { mode: 'best_accuracy', pool: SAMPLE_POOL }, + }); + await getConfiguredAutoRoutingSettings(env, { ownerType: 'user', ownerId: 'user-1' }); + stubs.get('user:user-1')?.getSettings.mockRejectedValueOnce(new Error('do unavailable')); + + await expect( + getConfiguredAutoRoutingSettings(env, { ownerType: 'user', ownerId: 'user-1' }) + ).resolves.toEqual({ mode: null, pool: null }); + expect(warnSpy).toHaveBeenCalledWith( + expect.stringContaining('auto_routing_config_read_failed') + ); + }); }); diff --git a/services/auto-routing/src/routing-mode.ts b/services/auto-routing/src/routing-mode.ts index 407f79be78..4961ae0570 100644 --- a/services/auto-routing/src/routing-mode.ts +++ b/services/auto-routing/src/routing-mode.ts @@ -1,8 +1,10 @@ import { AutoRoutingModeSchema, DEFAULT_AUTO_ROUTING_MODE, + EfficientModelPoolSchema, type AutoRoutingMode, type AutoRoutingModeOwnerType, + type EfficientModelPool, } from '@kilocode/auto-routing-contracts'; import { formatError } from '@kilocode/worker-utils'; import { DurableObject } from 'cloudflare:workers'; @@ -10,6 +12,12 @@ import { DurableObject } from 'cloudflare:workers'; type AutoRoutingModeEnv = Pick; const MODE_STORAGE_KEY = 'mode'; +const POOL_STORAGE_KEY = 'pool'; + +export type AutoRoutingOwnerSettings = { + mode: AutoRoutingMode | null; + pool: EfficientModelPool | null; +}; function modeKey(ownerType: AutoRoutingModeOwnerType, ownerId: string): string { return `${ownerType}:${ownerId}`; @@ -20,6 +28,11 @@ function parseStoredMode(raw: unknown): AutoRoutingMode | null { return parsed.success ? parsed.data : null; } +function parseStoredPool(raw: unknown): EfficientModelPool | null { + const parsed = EfficientModelPoolSchema.safeParse(raw); + return parsed.success ? parsed.data : null; +} + export class AutoRoutingModeConfigDO extends DurableObject { async getMode(): Promise { return parseStoredMode(await this.ctx.storage.get(MODE_STORAGE_KEY)); @@ -32,6 +45,51 @@ export class AutoRoutingModeConfigDO extends DurableObject { } await this.ctx.storage.put(MODE_STORAGE_KEY, mode); } + + async getPool(): Promise { + return parseStoredPool(await this.ctx.storage.get(POOL_STORAGE_KEY)); + } + + async setPool(pool: EfficientModelPool | null): Promise { + if (pool === null) { + await this.ctx.storage.delete(POOL_STORAGE_KEY); + return; + } + await this.ctx.storage.put(POOL_STORAGE_KEY, pool); + } + + async getSettings(): Promise { + const [mode, pool] = await Promise.all([this.getMode(), this.getPool()]); + return { mode, pool }; + } + + async setSettings(settings: AutoRoutingOwnerSettings): Promise { + // Mode and pool commit or fail together in one storage transaction. + // setMode/setPool stay single-key for individual callers. + const toPut: Record = {}; + const toDelete: string[] = []; + if (settings.mode === null) { + toDelete.push(MODE_STORAGE_KEY); + } else { + toPut[MODE_STORAGE_KEY] = settings.mode; + } + if (settings.pool === null) { + toDelete.push(POOL_STORAGE_KEY); + } else { + toPut[POOL_STORAGE_KEY] = settings.pool; + } + if (Object.keys(toPut).length === 0 && toDelete.length === 0) { + return; + } + await this.ctx.storage.transaction(async txn => { + if (Object.keys(toPut).length > 0) { + await txn.put(toPut); + } + if (toDelete.length > 0) { + await txn.delete(toDelete); + } + }); + } } function modeStub(env: AutoRoutingModeEnv, ownerType: AutoRoutingModeOwnerType, ownerId: string) { @@ -57,23 +115,57 @@ export async function getConfiguredAutoRoutingMode( }); } +export async function getConfiguredAutoRoutingSettings( + env: AutoRoutingModeEnv, + owner: { ownerType: AutoRoutingModeOwnerType; ownerId: string } +): Promise { + return modeStub(env, owner.ownerType, owner.ownerId) + .getSettings() + .catch((error: unknown) => { + console.warn( + JSON.stringify({ + event: 'auto_routing_config_read_failed', + key: modeKey(owner.ownerType, owner.ownerId), + ...formatError(error), + }) + ); + return { mode: null, pool: null }; + }); +} + export async function getAutoRoutingMode( env: AutoRoutingModeEnv, owner: { userId: string; organizationId: string | null } ): Promise { - if (owner.organizationId) { - const orgMode = await getConfiguredAutoRoutingMode(env, { - ownerType: 'org', - ownerId: owner.organizationId, - }); - if (orgMode) return orgMode; - } + const settings = await getEffectiveAutoRoutingSettings(env, owner); + return settings.mode; +} - const userMode = await getConfiguredAutoRoutingMode(env, { +/** + * Resolve mode and pool independently with organization → personal → platform + * precedence per field. A configured org mode does not imply an org pool. + */ +export async function getEffectiveAutoRoutingSettings( + env: AutoRoutingModeEnv, + owner: { userId: string; organizationId: string | null } +): Promise<{ mode: AutoRoutingMode; pool: EfficientModelPool | null }> { + const userSettingsPromise = getConfiguredAutoRoutingSettings(env, { ownerType: 'user', ownerId: owner.userId, }); - return userMode ?? DEFAULT_AUTO_ROUTING_MODE; + const orgSettingsPromise = owner.organizationId + ? getConfiguredAutoRoutingSettings(env, { + ownerType: 'org', + ownerId: owner.organizationId, + }) + : Promise.resolve({ mode: null, pool: null } satisfies AutoRoutingOwnerSettings); + + const [orgSettings, userSettings] = await Promise.all([orgSettingsPromise, userSettingsPromise]); + + return { + mode: orgSettings.mode ?? userSettings.mode ?? DEFAULT_AUTO_ROUTING_MODE, + pool: orgSettings.pool ?? userSettings.pool ?? null, + }; } export async function setAutoRoutingMode( @@ -83,3 +175,11 @@ export async function setAutoRoutingMode( ): Promise { await modeStub(env, owner.ownerType, owner.ownerId).setMode(mode); } + +export async function setAutoRoutingSettings( + env: AutoRoutingModeEnv, + owner: { ownerType: AutoRoutingModeOwnerType; ownerId: string }, + settings: AutoRoutingOwnerSettings +): Promise { + await modeStub(env, owner.ownerType, owner.ownerId).setSettings(settings); +}