diff --git a/.plans/experimental-models-1.md b/.plans/experimental-models-1.md index 0f98a70b19e..63a1ff22855 100644 --- a/.plans/experimental-models-1.md +++ b/.plans/experimental-models-1.md @@ -10,14 +10,14 @@ [todo] not started --> -| Phase | Status | Notes | -| ----------------------------------- | ------------- | --------------------------------------------------------------------------------------------------------------------------------------- | -| Phase 1 — Schema + Migration | [done] | Migration `0134_black_union_jack.sql` in main. PR #3299. | -| Phase 2 — Gateway Header Capture | [todo] | | -| Phase 3 — Variant Picker + Routing | [todo] | | -| Phase 4 — Usage, Metrics, Reporting | [todo] | | -| Phase 5 — Admin tRPC + UI | [in-progress] | Branch `mark/experimental-models-admin`. tRPC router + UI tab landed; `getLiveStats` and `getPromptByHash` deferred to Phase 4/R2 work. | -| Phase 6 — Specs + Tests | [todo] | Spec file `.specs/model-experiments.md` not yet created. | +| Phase | Status | Notes | +| ----------------------------------- | ---------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| Phase 1 — Schema + Migration | [done] | Migration `0134_black_union_jack.sql` in main. PR #3299. | +| Phase 2 — Gateway Header Capture | [merged-pending] | Branch `mark/experimental-models-gateway`, PR #3325 (draft). | +| Phase 3 — Variant Picker + Routing | [merged-pending] | Branch `mark/experimental-models-gateway`, PR #3325 (draft). | +| Phase 4 — Usage, Metrics, Reporting | [in-progress] | 4a (R2 helper) + 4b (capture/persist) + 4c (after-hook wiring) on PR #3325. 4d (`model_experiment_request_stats` view) intentionally deferred until a real report consumer needs it. | +| Phase 5 — Admin tRPC + UI | [in-progress] | tRPC router + UI tab landed in `#3302`; `getLiveStats` and `getPromptByHash` deferred to follow-up alongside admin UI to inflate prompts. | +| Phase 6 — Specs + Tests | [todo] | Spec file `.specs/model-experiments.md` not yet created. | **Phase 1 — concrete output (landed in `packages/db/src/schema.ts` + `0134_black_union_jack.sql`):** @@ -30,9 +30,48 @@ **Not yet done in Phase 1 (still owed by future PRs):** - ~~`ExperimentUpstreamSchema` zod schema (lives in app code, not in `packages/db`).~~ Landed with Phase 5 admin work in `apps/web/src/lib/ai-gateway/experiments/upstream-schema.ts`. -- The `model_experiment_request_stats` reporting view (Phase 4). +- The `model_experiment_request_stats` reporting view — explicitly deferred. See Phase 4d note below. - Convert `model_experiment_request` to a declarative-partitioned table (range on `created_at`, monthly). Phase 1 landed it as a plain table; partitioning is a follow-up — see "Partitioning (follow-up)" under the schema definition for rationale and conversion mechanics. +**Phase 2 — concrete output (in PR #3325):** + +- `apps/web/src/app/api/openrouter/[...path]/route.ts` — extracts `x-kilo-request` into `clientRequestId` and `x-kilo-session` as fallback for `session_id` when `x-kilocode-taskid` is absent. Captures `x-kilocode-machineid` once into `machineIdHeader` and threads it (plus the resolved client IP) into `getProvider`. +- `apps/web/src/lib/ai-gateway/processUsage.types.ts` — extends `MicrodollarUsageContext` with optional `clientRequestId`, `modelExperimentVariantVersionId`, `modelExperimentAllocationSubject`, and `experimentPromptCapture`. Optional so the dozens of construction sites in routes/tests/helpers don't need touching. Adds `ExperimentPromptCapture` type. + +**Phase 3 — concrete output (in PR #3325):** + +- `apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts` — `buildDirectProvider(input)` + `inferSupportedChatApis(...)`. Used by both the new experiment branch and the existing `kilo-internal/...` (custom_llm2) path so direct-to-upstream traffic shares one implementation. Custom_llm passes `extra_headers`; experiments deliberately don't (excluded from `ExperimentUpstreamSchema`). +- `apps/web/src/lib/ai-gateway/experiments/pick-variant.ts` — `isPublicIdExperimented(publicId)` (Redis membership pre-check with Postgres fallback; returns `unavailable` when both fail), `getRoutingExperimentForPublicId(publicId)` (Redis-cached resolved experiment with a single `SELECT DISTINCT ON (variant_id)` to pick each variant's current version, decrypts `encrypted_api_key` once at cache-build time, returns `none` / `experiment` / `unavailable`), and `pickModelExperimentVariant(input)` (deterministic `getRandomNumber` cumulative-weight walk in id-asc order, allocation subject precedence user → machine → ip; missing all subjects returns `unavailable`). +- `apps/web/src/lib/ai-gateway/experiments/index.ts` — public exports. +- `apps/web/src/lib/ai-gateway/providers/get-provider.ts` — refactored to return discriminated `GetProviderResult` (`provider` / `not-found` / `unavailable`). Adds the experiment branch after BYOK and before `kilo-internal/...` and the `kiloExclusiveModels` lookup. Active selections set new flags `skipBalanceCheck`, `skipProviderPin`, `skipKiloExclusiveModelSettings` on the result and attach `experiment` metadata. Custom_llm path refactored to use `buildDirectProvider`. +- `apps/web/src/app/api/openrouter/[...path]/route.ts` — calls the new `getProvider({...})` signature with `clientIp` + `machineId`, handles `not-found` (local model-unavailable) and `unavailable` (503 temporarily-unavailable) before reading `provider.supportedChatApis`. Honors `skipBalanceCheck` and `skipProviderPin` while still running organization model/data-collection policy. Sets `usageContext.modelExperimentVariantVersionId` + `modelExperimentAllocationSubject` from the result and calls `buildExperimentPromptCapture` after provider transforms for experimented requests only. +- `apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.ts` — accepts an optional options bag with `skipKiloExclusiveModelSettings` so the registry's `internal_id`/provider rewrite doesn't override the variant's upstream. Generic provider-specific request fixes and `provider.transformRequest` still run. +- `apps/web/src/lib/ai-gateway/auto-model/resolution.ts` — no auto-router changes. `autoFreeModels` and the frontier preset list are hand-curated and don't overlap with experiment preview ids; the explicit-opt-in property is preserved by construction. Avoids paying per-candidate Redis membership checks on every `kilo-auto/free` request. + +**Phase 4a–c — concrete output (in PR #3325):** + +- `apps/web/src/lib/r2/experiment-prompts.ts` — `putPromptIfAbsent(content)` (HEAD-then-PUT under sha256 hex key for automatic dedup), `getPromptByHash(sha)` (admin out-of-band reads with strict 64-char hex validation), and `sha256Hex(content)`. Throws on R2 errors; callers translate to `__failed__` sentinel. +- `apps/web/src/lib/r2/client.ts` — adds `R2_EXPERIMENT_PROMPTS_BUCKET_NAME` env var and `r2ExperimentPromptsBucketName` export. Per-environment buckets `kilo-experiment-prompts-dev` / `kilo-experiment-prompts-prod`. +- `apps/web/src/lib/ai-gateway/experiments/persist.ts` — `buildExperimentPromptCapture(request)` (chat_completions: extracts leading `role: 'system'` message; messages/responses: folds whole body into the body side; both sides tail-truncated at 4 MB with Sentry breadcrumb on truncation). `persistExperimentAttribution(input)` (parallel `Promise.allSettled` R2 puts via `putPromptSafely`, single-row insert into `model_experiment_request` with `__absent__`/`__failed__` sentinels per side; Sentry-reports failures and swallows — best-effort analytics, must not roll back the billing write). +- `apps/web/src/lib/ai-gateway/processUsage.ts` — `logMicrodollarUsage` and `processTokenData` now return `{ usageId, createdAt }` so the experiment attribution row keys onto the same usage row. Existing callers ignoring the return value are unaffected. +- `apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts` — `accountForMicrodollarUsage` chains `persistExperimentAttribution` after the microdollar write inside the same `after()` hook, only for experimented requests. + +**Phase 4d — explicitly deferred:** + +The `model_experiment_request_stats` reporting view was implemented and reverted out of PR #3325 because there is no consumer yet. The two real arguments for the view are (a) centralizing the 4-table join + JSON extraction and (b) physically excluding `upstream->>'api_key'` at the column level so any consumer of the view cannot leak the key. (a) is also achievable inline with Drizzle queries, and (b) can be enforced with a code-review rule + grep test until a real consumer materializes. Re-add when `getLiveStats` (Phase 5 deferred) or another report needs it; by then the right column set will be obvious. + +**Stale membership cache — follow-up (next step after PR #3325 lands):** + +Surfaced in PR #3325 review. Two compounding issues in the current routing pre-check: + +- `EXPERIMENTED_PUBLIC_IDS_REDIS_KEY` is written without a TTL by `recomputeExperimentedPublicIds()` (`apps/web/src/routers/admin/model-experiments-router.ts`), so a stale value persists indefinitely. +- `isPublicIdExperimented(publicId)` (`apps/web/src/lib/ai-gateway/experiments/pick-variant.ts`) trusts a populated cache absolutely on a "not in list" read; it only falls through to Postgres when Redis returns null or throws. +- Cache invalidation in the admin router catches and swallows errors. + +Failure mode: admin activates an experiment, Postgres write succeeds, Redis recompute fails (network blip), the membership key keeps the previous list. Until another admin mutation succeeds against Redis, every request to the experimented public id silently routes through default routing instead of the experiment. + +Fix: add a TTL (e.g. 300s) on `redisSet(EXPERIMENTED_PUBLIC_IDS_REDIS_KEY, …)` so the next reader after expiry gets `null` from Redis, falls through to the existing DB query path, and writes a fresh list. Self-healing with bounded staleness. A "DB fallback when populated cache says no" was considered and rejected — most requests are non-experiment, so per-miss DB queries would defeat the cache's purpose. Targeted test: simulate Redis being unavailable during activate, then a subsequent request after TTL expiry should pick up the experiment via the DB fallback. + **Phase 5 — concrete output (in PR — see status table for branch/PR):** - `apps/web/src/lib/ai-gateway/experiments/upstream-schema.ts` — `ExperimentUpstreamSchema` (strict subset of `CustomLlmDefinitionSchema`, no `api_key`, no `extra_headers`). @@ -352,7 +391,7 @@ In `apps/web/src/app/api/openrouter/[...path]/route.ts`: Add `apps/web/src/lib/ai-gateway/experiments/`: - `pick-variant.ts` - - `isPublicIdExperimented(publicId)`: fast membership check through helpers added to `apps/web/src/lib/redis-keys.ts` (`EXPERIMENTED_PUBLIC_IDS_REDIS_KEY` / `modelExperimentRedisKey(publicId)`). The membership value contains every `public_model_id` with `status IN ('active', 'paused')`. Used by `getProvider` (see below) as a fast pre-check before the per-public-id fetch, and by the `kilo-auto` candidate-set construction. On Redis error, the function queries Postgres for that `public_model_id` instead of falling through. If both Redis and Postgres are unavailable, return an explicit `unavailable` result so explicit requests to preview experiment ids receive the gateway's "temporarily unavailable" response instead of silently routing as non-experimented traffic. + - `isPublicIdExperimented(publicId)`: fast membership check through helpers added to `apps/web/src/lib/redis-keys.ts` (`EXPERIMENTED_PUBLIC_IDS_REDIS_KEY` / `modelExperimentRedisKey(publicId)`). The membership value contains every `public_model_id` with `status IN ('active', 'paused')`. Used by `getProvider` (see below) as a fast pre-check before the per-public-id fetch. On Redis error, the function queries Postgres for that `public_model_id` instead of falling through. If both Redis and Postgres are unavailable, return an explicit `unavailable` result so explicit requests to preview experiment ids receive the gateway's "temporarily unavailable" response instead of silently routing as non-experimented traffic. - `getRoutingExperimentForPublicId(publicId)`: returns the routing-relevant experiment with its current status (`active` or `paused`) and resolved variant + version data, `null` when Postgres proves there is no routing-relevant experiment, or `unavailable` when cache/database/config failures prevent a safe routing decision. For each variant, the cached payload contains the current `variant_version_id`, the `upstream` JSONB blob (no key), and the **decrypted** `api_key` merged in alongside as a separate field (in-memory shape: `{ ...upstream, api_key }`). Per-public-id cache at `modelExperimentRedisKey(publicId)`, Redis-cached for 10 minutes. Pre-checks `isPublicIdExperimented` to avoid fetching when no experiment exists. The cache build resolves "current version" per variant via `SELECT DISTINCT ON (variant_id) id, variant_id, upstream, encrypted_api_key, effective_at FROM model_experiment_variant_version WHERE variant_id IN (...) AND effective_at <= now() ORDER BY variant_id, effective_at DESC, id DESC` (Postgres-specific; one query for the experiment, no per-variant round trips), then calls `decryptApiKey(encrypted_api_key, BYOK_ENCRYPTION_KEY)` per row before serialising to Redis. If `BYOK_ENCRYPTION_KEY` is unset for an active/paused experiment, return `unavailable` and log a single warn-level error per process boot. - `pickModelExperimentVariant({ publicModelId, userId, machineId, clientIp })`: calls `getRoutingExperimentForPublicId`. Behavior depends on returned experiment status: - `active`: pick a variant and return `{ status: 'active', experimentId, variantId, variantVersionId, upstream, allocationSubject }`. If no allocation subject is available (no userId/machineId/clientIp), capture the invariant violation and return `{ status: 'unavailable' }`. @@ -433,7 +472,7 @@ Integration in `getProvider` (`apps/web/src/lib/ai-gateway/providers/get-provide Routing scope: - Applies only when the request's resolved public id is in the experimented SET. Under Dedicated mode v1 these are dedicated testing public ids (e.g. `kilo/preview-experiment-foo`) that clients select explicitly. -- `kilo-auto` resolution does not feed experimented public ids: the auto-router's candidate-set construction excludes any public id where `isPublicIdExperimented(publicId)` is true (one-line guard near `applyResolvedAutoModel`). Dedicated testing ids never get silently selected by auto-routing. +- `kilo-auto` resolution does not feed experimented public ids by construction: `autoFreeModels` and the frontier preset list are hand-curated, and dedicated preview ids are never added to either. No runtime guard is required (and adding one would force per-candidate Redis membership checks on every `kilo-auto/free` request). The invariant lives in code review of those static lists. - Does not apply to BYOK requests or `kilo-internal/...` traffic (those branches are matched first / by id prefix and never reach the experiment branch). - Balance checks are skipped for experimented preview ids because v1 traffic is free/provider-funded. Server-side organization allow/deny and data-collection policy checks still run against the public model id; direct experiment routing ignores only the provider-pinning side effect because the upstream is selected by the experiment variant. - Experimented traffic goes **direct to `upstream.base_url`** — OpenRouter and Vercel are never contacted. No gateway pin needed. @@ -575,9 +614,57 @@ These constraints exist because of how the gateway is built today. The spec must - Identifier-less traffic; under v1 such requests fail closed as temporarily unavailable because missing IP is treated as a gateway invariant violation. - A/B variants spanning entirely different public model ids. - Client-visible variant ids or variant-aware UI behavior. +- Response-side rewriting of the served `internal_id` back to the requested `public_id` for experiment traffic — see "Followup: rewrite checkpoint identity in experiment responses" below. - Partner trace export, redaction, HMAC webhooks, partner auth, and warehouse coordination (see [Part 2](./experimental-models-2.md)). - Replay bundles, SWE-bench/OpenHands adapters, and held-out replay-eval service (see [Part 2](./experimental-models-2.md)). +## Followup: rewrite checkpoint identity in experiment responses + +Surfaced in PR #3325 review (P1). Experiment routing rewrites the outbound request to the variant's `upstream.internal_id` (in `buildDirectProvider`, applied via `body.model` before the partner fetch), but the response is returned to the client unchanged. + +The existing response-rewriting branch in `apps/web/src/app/api/openrouter/[...path]/route.ts:715,733` only runs for `kilo-exclusive` free traffic flowing through OpenRouter or Vercel — experiment providers carry `provider.id === 'custom'` and bypass it. As a result, OpenAI- and Anthropic-shape partner responses echo `internal_id` in the JSON body and in streaming `model:` events, disclosing the served checkpoint and variant to the client. + +This violates the client-blinding requirement from the Accepted Design ("Clients keep sending the same public model id" — line 112) and from the spec ("client blinding" — `Phase 6 — Specs + Tests`). A user could diff response payloads across requests to deduce their bucket assignment and observe checkpoint hot-swaps. + +Fix: rewrite `model` back to the requested `public_id` in experiment responses on the way out, mirroring the existing kilo-exclusive rewrite. Both response shapes need coverage: + +- Non-streaming JSON: replace `model` in the parsed body before returning. +- Streaming SSE/event-stream: rewrite the per-chunk `model` field in chat-completions deltas, Anthropic `message_start` / `message_delta` events, and Responses-API `response.created` / `response.completed` events. The existing `rewriteFreeModelResponse_*` helpers in `apps/web/src/lib/ai-gateway/providers/openrouter/responses.ts` (and siblings) already implement this for the gateway-routed path; experiment traffic should reuse the same rewriters keyed on `experiment` rather than provider id, or the predicate at `route.ts:715` should be widened to "rewrite when the served model id differs from the requested public id" so the kilo-exclusive and experiment paths share one rule. + +Targeted test: end-to-end an experimented chat-completions and messages request, assert the streamed and final-JSON `model` values match the requested public id and never the variant's `internal_id`. + +## Followup: unify direct-upstream routing abstraction + +Three routing paths now bypass parts of the OpenRouter policy machinery: +`custom_llm2`, `direct-byok`, and experiments. Each does so through different +flags on `GetProviderProviderResult` (`bypassAccessCheck`, `skipProviderPin`, +`skipKiloExclusiveModelSettings`) and ad-hoc `if (experiment && ...)` policy +refusals in `route.ts` (data collection, provider allow-list). + +The proliferation is a symptom of treating each direct-upstream caller as a +special case rather than as instances of one abstraction. A followup PR +should: + +- Collapse the per-caller flags into a single notion (e.g. `routingMode: +'gateway' | 'direct'`) on the provider result. `gateway` flows through + OpenRouter/Vercel and accepts the full `body.provider` policy machinery; + `direct` does not. +- Move the policy-refusal points (currently `if (experiment && settings?.data_collection === 'deny')`, + `if (experiment && providerConfig?.only !== undefined)`) into a single + `checkPolicyEnforcableOnDirect` step that runs for every `direct`-mode + request and returns the appropriate refusal when the org has explicit + policy that the gateway can't enforce on a direct partner endpoint. +- Reconsider `custom_llm2`'s `bypassAccessCheck: true`. Today it skips the + whole org-restrictions block (per the AI Gateway `AGENTS.md`: "enabling + requires explicit admin action, so the org allow-list doesn't apply"). + That justification holds for per-org admin-enabled custom LLMs but not + for globally-routed experiment public ids; the unified abstraction + should make that distinction explicit rather than burying it in flag + combinations. + +This refactor is out of scope for the experiment-routing PR and is tracked +here so the next PR touching the gateway routing surface can address it. + ## Files Touched Core experiment implementation: diff --git a/apps/web/src/app/api/openrouter/[...path]/route.ts b/apps/web/src/app/api/openrouter/[...path]/route.ts index f15667c47a2..ee49a96162f 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.ts @@ -18,6 +18,8 @@ import type { } from '@/lib/ai-gateway/providers/openrouter/types'; import { applyProviderSpecificLogic } from '@/lib/ai-gateway/providers/apply-provider-specific-logic'; import { getProvider } from '@/lib/ai-gateway/providers/get-provider'; +import { buildExperimentPromptCapture } from '@/lib/ai-gateway/experiments/persist'; +import { isPublicIdExperimented } from '@/lib/ai-gateway/experiments/membership'; import { upstreamRequest } from '@/lib/ai-gateway/providers/upstream-request'; import { debugSaveProxyRequest } from '@/lib/debugUtils'; import { setTag, startInactiveSpan } from '@sentry/nextjs'; @@ -43,6 +45,7 @@ import { malformedJsonResponse, makeErrorReadable, modelDoesNotExistResponse, + modelNotAllowedResponse, extractHeaderAndLimitLength, noFreeModelsAvailableResponse, temporarilyUnavailableResponse, @@ -276,7 +279,9 @@ export async function POST(request: NextRequest): Promise( + async () => { + try { + const cached = await redisGet(EXPERIMENTED_PUBLIC_IDS_REDIS_KEY); + if (cached === null) return []; + return parseStringArray(cached) ?? []; + } catch { + // captureException already invoked by the redis helper + return []; + } + }, + EXPERIMENTED_PUBLIC_IDS_LOCAL_CACHE_TTL_MS, + [] +); + +/** + * Fast pre-check: does any active|paused experiment target this public id? + * + * Reads from a short-lived in-process cache backed by Redis. + */ +export async function isPublicIdExperimented(publicId: string): Promise { + const ids = await getExperimentedPublicIds(); + return ids.includes(publicId); +} + +function parseStringArray(raw: string): string[] | null { + try { + const parsed = JSON.parse(raw); + if (!Array.isArray(parsed)) return null; + if (parsed.some(v => typeof v !== 'string')) return null; + return parsed as string[]; + } catch { + return null; + } +} diff --git a/apps/web/src/lib/ai-gateway/experiments/persist.test.ts b/apps/web/src/lib/ai-gateway/experiments/persist.test.ts new file mode 100644 index 00000000000..5f2335b2233 --- /dev/null +++ b/apps/web/src/lib/ai-gateway/experiments/persist.test.ts @@ -0,0 +1,132 @@ +import { describe, expect, it } from '@jest/globals'; +import { + buildExperimentPromptCapture, + REQUEST_BODY_CAP_BYTES, + truncateToUtf8Bytes, +} from './persist'; +import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types'; + +describe('truncateToUtf8Bytes', () => { + it('returns the input unchanged when already within the cap', () => { + expect(truncateToUtf8Bytes('hello', 100)).toBe('hello'); + }); + + it('returns valid UTF-8 with length <= maxBytes for ASCII', () => { + const out = truncateToUtf8Bytes('a'.repeat(1000), 16); + expect(Buffer.byteLength(out, 'utf8')).toBe(16); + expect(out).toBe('a'.repeat(16)); + }); + + it('does not split a multi-byte UTF-8 codepoint', () => { + // '日' is 3 bytes in UTF-8 (0xE6 0x97 0xA5). With maxBytes=4 the cut + // would otherwise land mid-codepoint and produce invalid UTF-8. + const input = '日日日'; // 9 UTF-8 bytes + const out = truncateToUtf8Bytes(input, 4); + // Result must be valid UTF-8 and <= 4 bytes. The only valid prefixes + // here are '' (0 bytes) and '日' (3 bytes). + expect(Buffer.byteLength(out, 'utf8')).toBeLessThanOrEqual(4); + expect(out).toBe('日'); + }); + + it('handles 4-byte emoji codepoints correctly', () => { + // '😀' is 4 bytes in UTF-8 (a surrogate pair in UTF-16). + const input = '😀😀😀'; + const out = truncateToUtf8Bytes(input, 6); + // Only '😀' (4 bytes) fits cleanly; cutting at byte 6 would split + // the second emoji. + expect(out).toBe('😀'); + expect(Buffer.byteLength(out, 'utf8')).toBe(4); + }); + + it('produces deterministic output across calls', () => { + const input = 'hello 日本語 world 🎉 こんにちは'; + const a = truncateToUtf8Bytes(input, 12); + const b = truncateToUtf8Bytes(input, 12); + expect(a).toBe(b); + expect(Buffer.byteLength(a, 'utf8')).toBeLessThanOrEqual(12); + }); + + it('cuts cleanly when the cap exactly aligns with a codepoint boundary', () => { + const input = '日日'; // 6 bytes + const out = truncateToUtf8Bytes(input, 3); + expect(out).toBe('日'); + }); + + it('handles maxBytes of 0', () => { + expect(truncateToUtf8Bytes('hello', 0)).toBe(''); + }); +}); + +describe('buildExperimentPromptCapture', () => { + function chatRequest(body: Record): GatewayRequest { + // Tests may pass synthetic message shapes that don't fully match + // the production OpenAI types; double-cast through unknown is the + // pragmatic test-only escape hatch. + return { kind: 'chat_completions', body } as unknown as GatewayRequest; + } + + it('captures the full serialized body and records the request kind', () => { + const cap = buildExperimentPromptCapture( + chatRequest({ + model: 'kilo/preview-foo', + messages: [ + { role: 'system', content: 'you are a helpful assistant' }, + { role: 'user', content: 'hi' }, + ], + }) + ); + expect(cap.requestKind).toBe('chat_completions'); + expect(cap.requestBodyContent).toContain('"role":"system"'); + expect(cap.requestBodyContent).toContain('"role":"user"'); + expect(cap.wasTruncated).toBe(false); + }); + + it('records `messages` for an Anthropic-shape request', () => { + const cap = buildExperimentPromptCapture({ + kind: 'messages', + body: { + model: 'kilo/preview-foo', + system: 'you are a helpful assistant', + messages: [{ role: 'user', content: 'hi' }], + }, + } as unknown as GatewayRequest); + expect(cap.requestKind).toBe('messages'); + expect(cap.requestBodyContent).toContain('"system":"you are a helpful assistant"'); + }); + + it('records `responses` for a Responses-API request', () => { + const cap = buildExperimentPromptCapture({ + kind: 'responses', + body: { model: 'kilo/preview-foo', input: 'hi' }, + } as unknown as GatewayRequest); + expect(cap.requestKind).toBe('responses'); + expect(cap.requestBodyContent).toContain('"input":"hi"'); + }); + + it('truncates the body by UTF-8 bytes and marks wasTruncated', () => { + const huge = 'a'.repeat(REQUEST_BODY_CAP_BYTES + 1024); + const cap = buildExperimentPromptCapture( + chatRequest({ + model: 'kilo/preview-foo', + messages: [{ role: 'user', content: huge }], + }) + ); + expect(cap.wasTruncated).toBe(true); + expect(Buffer.byteLength(cap.requestBodyContent, 'utf8')).toBeLessThanOrEqual( + REQUEST_BODY_CAP_BYTES + ); + }); + + it('produces deterministic output for identical inputs', () => { + // Content addressing on top of this capture relies on byte-for-byte + // determinism; a JSON.stringify implementation that reordered keys + // would silently break dedup. + const body = { + model: 'kilo/preview-foo', + messages: [{ role: 'user', content: 'hi' }], + }; + const a = buildExperimentPromptCapture(chatRequest({ ...body })); + const b = buildExperimentPromptCapture(chatRequest({ ...body })); + expect(a.requestBodyContent).toBe(b.requestBodyContent); + }); +}); diff --git a/apps/web/src/lib/ai-gateway/experiments/persist.ts b/apps/web/src/lib/ai-gateway/experiments/persist.ts new file mode 100644 index 00000000000..c057594a7e7 --- /dev/null +++ b/apps/web/src/lib/ai-gateway/experiments/persist.ts @@ -0,0 +1,112 @@ +import { captureException } from '@sentry/nextjs'; +import { db } from '@/lib/drizzle'; +import { model_experiment_request } from '@kilocode/db/schema'; +import { putPromptOrNull } from '@/lib/r2/experiment-prompts'; +import type { ExperimentPromptCapture } from '@/lib/ai-gateway/processUsage.types'; +import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types'; + +/** + * Sentinel values written to `model_experiment_request.request_body_sha256` + * instead of a real 64-char hex digest. The schema CHECK constraint accepts + * either a 64-char lowercase hex digest or one of these sentinels. + * + * - `__failed__`: R2 storage failed; the attribution row still lands. + * - `__deleted__`: prompt content intentionally wiped while retaining + * experiment attribution. + */ +export const PROMPT_HASH_FAILED = '__failed__'; +export const PROMPT_HASH_DELETED = '__deleted__'; + +/** 4 MB measured as UTF-8 bytes (matching what `putPromptIfAbsent` actually + * uploads). Comfortably above any current 1M-token-context request. */ +export const REQUEST_BODY_CAP_BYTES = 4 * 1024 * 1024; + +/** + * Tail-truncates a string to at most `maxBytes` UTF-8 bytes, stepping back + * to a UTF-8 codepoint boundary so the result is always valid UTF-8. + * + * `string.length` would have measured UTF-16 code units, which under-counts + * non-ASCII content by up to 4× when encoded as UTF-8 — a 4 MB code-unit + * cap on CJK content can balloon to ~12 MB on the wire. Slicing UTF-16 + * code units could also split surrogate pairs and produce U+FFFD on encode, + * making the resulting hash non-deterministic across runtimes. + */ +export function truncateToUtf8Bytes(s: string, maxBytes: number): string { + const buf = Buffer.from(s, 'utf8'); + if (buf.length <= maxBytes) return s; + // UTF-8 continuation bytes have the form 0b10xxxxxx. Walk backward from + // `maxBytes` until we land on a non-continuation byte (a codepoint + // start), then slice at that boundary. + let end = maxBytes; + while (end > 0 && (buf[end] & 0xc0) === 0x80) end--; + return buf.subarray(0, end).toString('utf8'); +} + +/** + * Build a bounded capture of the canonical post-`transformRequest` upstream + * request body. The full serialized body is content-addressed as a single + * blob; the upstream API shape (`chat_completions` / `messages` / + * `responses`) is recorded on `requestKind` so downstream reporting can + * tell what was sent without re-parsing. + * + * The 4 MB cap ensures we never retain unbounded data through the async + * `after()` path. If the body exceeds the cap it is tail-truncated + * deterministically and `wasTruncated` is set to true. + */ +export function buildExperimentPromptCapture(request: GatewayRequest): ExperimentPromptCapture { + let bodyContent = typeof request.body === 'string' ? request.body : JSON.stringify(request.body); + let wasTruncated = false; + if (Buffer.byteLength(bodyContent, 'utf8') > REQUEST_BODY_CAP_BYTES) { + bodyContent = truncateToUtf8Bytes(bodyContent, REQUEST_BODY_CAP_BYTES); + wasTruncated = true; + } + return { + requestKind: request.kind, + requestBodyContent: bodyContent, + wasTruncated, + }; +} + +export type PersistExperimentAttributionInput = { + usageId: string; + createdAt: string; + variantVersionId: string; + allocationSubject: 'user' | 'machine' | 'ip'; + clientRequestId: string | null; + capture: ExperimentPromptCapture | null; +}; + +/** + * Insert one row into `model_experiment_request` for an experimented + * request. The R2 put is best-effort; on failure we still write the + * attribution row with `__failed__` as the body hash. + * + * Failures here MUST NOT roll back the microdollar usage write. Errors + * are reported to Sentry and swallowed. + */ +export async function persistExperimentAttribution( + input: PersistExperimentAttributionInput +): Promise { + const requestKind = input.capture?.requestKind ?? 'chat_completions'; + const wasTruncated = input.capture?.wasTruncated ?? false; + const storedPromptHash = + input.capture && (await putPromptOrNull(input.capture.requestBodyContent)); + + try { + await db.insert(model_experiment_request).values({ + usage_id: input.usageId, + variant_version_id: input.variantVersionId, + allocation_subject: input.allocationSubject, + client_request_id: input.clientRequestId, + request_kind: requestKind, + request_body_sha256: storedPromptHash ?? PROMPT_HASH_FAILED, + was_truncated: wasTruncated, + created_at: input.createdAt, + }); + } catch (err) { + captureException(err, { + tags: { source: 'model-experiments', operation: 'persistExperimentAttribution' }, + extra: { usageId: input.usageId, variantVersionId: input.variantVersionId }, + }); + } +} diff --git a/apps/web/src/lib/ai-gateway/experiments/pick-variant.test.ts b/apps/web/src/lib/ai-gateway/experiments/pick-variant.test.ts new file mode 100644 index 00000000000..2ea839860c9 --- /dev/null +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.test.ts @@ -0,0 +1,330 @@ +import { afterEach, beforeEach, describe, expect, it } from '@jest/globals'; +import { cleanupDbForTest, db } from '@/lib/drizzle'; +import { insertTestUser } from '@/tests/helpers/user.helper'; +import { createCallerForUser } from '@/routers/test-utils'; +import { model_experiment, model_experiment_variant_version } from '@kilocode/db/schema'; +import { eq } from 'drizzle-orm'; +import { isPublicIdExperimented } from './membership'; +import { pickModelExperimentVariant } from './pick-variant'; +import { redisDel, redisSet } from '@/lib/redis'; +import { EXPERIMENTED_PUBLIC_IDS_REDIS_KEY } from '@/lib/redis-keys'; +import type { User } from '@kilocode/db/schema'; + +let admin: User; + +const upstreamA = { + internal_id: 'partner-checkpoint-a', + base_url: 'https://partner.example.com/v1', +}; +const upstreamB = { + internal_id: 'partner-checkpoint-b', + base_url: 'https://partner.example.com/v1', +}; +const redisIt = process.env.REDIS_URL ? it : it.skip; + +beforeEach(async () => { + await cleanupDbForTest(); + admin = await insertTestUser({ + google_user_email: `admin-${Math.random()}@admin.example.com`, + is_admin: true, + }); +}); + +async function clearRoutingCaches() { + // Tests share the dev Redis instance across runs; flush the membership key + // so each test sees a fresh load path. + await redisDel(EXPERIMENTED_PUBLIC_IDS_REDIS_KEY); +} + +async function seedExperimentedPublicIds(ids: string[]): Promise { + return await redisSet(EXPERIMENTED_PUBLIC_IDS_REDIS_KEY, JSON.stringify(ids)); +} + +afterEach(async () => { + await redisDel(EXPERIMENTED_PUBLIC_IDS_REDIS_KEY); +}); + +async function makeActiveExperiment(opts: { + publicId: string; + weights?: [number, number]; + apiKeys?: [string, string]; +}) { + const { publicId, weights = [1, 1], apiKeys = ['sk-a', 'sk-b'] } = opts; + const caller = await createCallerForUser(admin.id); + const exp = await caller.admin.modelExperiments.create({ + public_model_id: publicId, + name: `exp-${publicId}`, + }); + const a = await caller.admin.modelExperiments.addVariant({ + id: exp.id, + label: 'control', + weight: weights[0], + }); + const b = await caller.admin.modelExperiments.addVariant({ + id: exp.id, + label: 'treatment', + weight: weights[1], + }); + await caller.admin.modelExperiments.swapVariantVersion({ + variantId: a.id, + upstream: upstreamA, + apiKey: apiKeys[0], + }); + await caller.admin.modelExperiments.swapVariantVersion({ + variantId: b.id, + upstream: upstreamB, + apiKey: apiKeys[1], + }); + await caller.admin.modelExperiments.activate({ id: exp.id }); + return { experimentId: exp.id, variantA: a.id, variantB: b.id }; +} + +describe('isPublicIdExperimented', () => { + it('returns false for an unknown public id', async () => { + await clearRoutingCaches(); + expect(await isPublicIdExperimented('kilo/preview-not-experimented')).toBe(false); + }); + + redisIt('returns true when the public id has an active experiment', async () => { + await makeActiveExperiment({ publicId: 'kilo/preview-iset-active' }); + expect(await seedExperimentedPublicIds(['kilo/preview-iset-active'])).toBe(true); + expect(await isPublicIdExperimented('kilo/preview-iset-active')).toBe(true); + }); + + redisIt('returns true when the public id has only a paused experiment', async () => { + const { experimentId } = await makeActiveExperiment({ publicId: 'kilo/preview-iset-paused' }); + const caller = await createCallerForUser(admin.id); + await caller.admin.modelExperiments.pause({ id: experimentId }); + expect(await seedExperimentedPublicIds(['kilo/preview-iset-paused'])).toBe(true); + expect(await isPublicIdExperimented('kilo/preview-iset-paused')).toBe(true); + }); +}); + +describe('pickModelExperimentVariant', () => { + it('returns null for a public id with no routing-relevant experiment', async () => { + await clearRoutingCaches(); + const result = await pickModelExperimentVariant({ + publicModelId: 'kilo/preview-pick-none', + userId: 'user-1', + machineId: null, + clientIp: null, + }); + expect(result).toBeNull(); + }); + + it('produces stable assignments for the same userId', async () => { + await makeActiveExperiment({ publicId: 'kilo/preview-stable' }); + const first = await pickModelExperimentVariant({ + publicModelId: 'kilo/preview-stable', + userId: 'user-1', + machineId: null, + clientIp: null, + }); + const second = await pickModelExperimentVariant({ + publicModelId: 'kilo/preview-stable', + userId: 'user-1', + machineId: null, + clientIp: null, + }); + expect(first?.status).toBe('active'); + expect(second?.status).toBe('active'); + if (first?.status !== 'active' || second?.status !== 'active') return; + expect(first.variantId).toBe(second.variantId); + expect(first.variantVersionId).toBe(second.variantVersionId); + }); + + it('decrypts and returns the partner-issued api key for the chosen variant', async () => { + const { variantA, variantB } = await makeActiveExperiment({ + publicId: 'kilo/preview-key', + apiKeys: ['sk-control-secret', 'sk-treatment-secret'], + }); + const result = await pickModelExperimentVariant({ + publicModelId: 'kilo/preview-key', + userId: 'user-key', + machineId: null, + clientIp: null, + }); + expect(result?.status).toBe('active'); + if (result?.status !== 'active') return; + expect(result.upstream.api_key).toMatch(/^sk-(control|treatment)-secret$/); + expect([variantA, variantB]).toContain(result.variantId); + }); + + it('respects allocation-subject precedence: user > machine > ip', async () => { + await makeActiveExperiment({ publicId: 'kilo/preview-alloc' }); + const userPick = await pickModelExperimentVariant({ + publicModelId: 'kilo/preview-alloc', + userId: 'user-z', + machineId: 'machine-z', + clientIp: '1.2.3.4', + }); + const machinePick = await pickModelExperimentVariant({ + publicModelId: 'kilo/preview-alloc', + userId: null, + machineId: 'machine-z', + clientIp: '1.2.3.4', + }); + const ipPick = await pickModelExperimentVariant({ + publicModelId: 'kilo/preview-alloc', + userId: null, + machineId: null, + clientIp: '1.2.3.4', + }); + expect(userPick?.status).toBe('active'); + expect(machinePick?.status).toBe('active'); + expect(ipPick?.status).toBe('active'); + if (userPick?.status !== 'active') return; + if (machinePick?.status !== 'active') return; + if (ipPick?.status !== 'active') return; + expect(userPick.allocationSubject).toBe('user'); + expect(machinePick.allocationSubject).toBe('machine'); + expect(ipPick.allocationSubject).toBe('ip'); + }); + + it('returns unavailable when no allocation subject is available', async () => { + await makeActiveExperiment({ publicId: 'kilo/preview-noalloc' }); + const result = await pickModelExperimentVariant({ + publicModelId: 'kilo/preview-noalloc', + userId: null, + machineId: null, + clientIp: null, + }); + expect(result?.status).toBe('unavailable'); + }); + + it('returns not-found for a paused experiment so traffic does not silently fall through', async () => { + const { experimentId } = await makeActiveExperiment({ publicId: 'kilo/preview-paused' }); + const caller = await createCallerForUser(admin.id); + await caller.admin.modelExperiments.pause({ id: experimentId }); + await clearRoutingCaches(); + const result = await pickModelExperimentVariant({ + publicModelId: 'kilo/preview-paused', + userId: 'user-q', + machineId: null, + clientIp: null, + }); + expect(result?.status).toBe('not-found'); + }); + + it('hot-swap: serves the new variant_version_id but keeps the same bucket', async () => { + const { experimentId, variantA, variantB } = await makeActiveExperiment({ + publicId: 'kilo/preview-hotswap', + }); + const caller = await createCallerForUser(admin.id); + const before = await pickModelExperimentVariant({ + publicModelId: 'kilo/preview-hotswap', + userId: 'user-pinned', + machineId: null, + clientIp: null, + }); + expect(before?.status).toBe('active'); + if (before?.status !== 'active') return; + + // Hot-swap whichever variant the user landed on with a new RC. + await caller.admin.modelExperiments.swapVariantVersion({ + variantId: before.variantId, + upstream: { ...upstreamA, internal_id: 'partner-checkpoint-rc-next' }, + apiKey: 'sk-rc-next', + }); + await clearRoutingCaches(); + + const after = await pickModelExperimentVariant({ + publicModelId: 'kilo/preview-hotswap', + userId: 'user-pinned', + machineId: null, + clientIp: null, + }); + expect(after?.status).toBe('active'); + if (after?.status !== 'active') return; + // Same slot (variantId), new RC (variantVersionId), new internal_id. + expect(after.variantId).toBe(before.variantId); + expect(after.variantVersionId).not.toBe(before.variantVersionId); + expect(after.upstream.internal_id).toBe('partner-checkpoint-rc-next'); + expect(after.upstream.api_key).toBe('sk-rc-next'); + // sanity: experiment + slots unchanged + expect([variantA, variantB]).toContain(after.variantId); + expect(experimentId).toBeTruthy(); + }); + + it('weighted distribution lands roughly on configured weights', async () => { + // 1:3 split. With 200 distinct seeds, control should be near 25%. + await makeActiveExperiment({ + publicId: 'kilo/preview-weighted', + weights: [1, 3], + }); + const counts = { control: 0, treatment: 0 }; + for (let i = 0; i < 200; i++) { + const r = await pickModelExperimentVariant({ + publicModelId: 'kilo/preview-weighted', + userId: `user-${i}`, + machineId: null, + clientIp: null, + }); + if (r?.status !== 'active') throw new Error('expected active'); + if (r.upstream.internal_id === upstreamA.internal_id) counts.control++; + else counts.treatment++; + } + // Loose bounds: 1:3 ≈ 25/75. Allow ±10pp for n=200. + const controlPct = counts.control / 200; + expect(controlPct).toBeGreaterThan(0.15); + expect(controlPct).toBeLessThanOrEqual(0.35); + }); + + it('historical attribution survives hot-swap: old variant_version_id still resolves to old upstream via DB', async () => { + const { experimentId } = await makeActiveExperiment({ + publicId: 'kilo/preview-attr', + }); + const before = await pickModelExperimentVariant({ + publicModelId: 'kilo/preview-attr', + userId: 'user-attr', + machineId: null, + clientIp: null, + }); + expect(before?.status).toBe('active'); + if (before?.status !== 'active') return; + + const oldVersionId = before.variantVersionId; + const oldInternalId = before.upstream.internal_id; + + const caller = await createCallerForUser(admin.id); + await caller.admin.modelExperiments.swapVariantVersion({ + variantId: before.variantId, + upstream: { ...upstreamA, internal_id: 'rc-newer' }, + apiKey: 'sk-newer', + }); + + // Old version row is still in the DB and still points at the old + // upstream — `model_experiment_request` rows that referenced + // `oldVersionId` remain attributable to the original RC. + const [row] = await db + .select() + .from(model_experiment_variant_version) + .where(eq(model_experiment_variant_version.id, oldVersionId)); + expect(row).toBeDefined(); + const oldUpstream = row.upstream as { internal_id: string }; + expect(oldUpstream.internal_id).toBe(oldInternalId); + expect(experimentId).toBeTruthy(); + }); + + it('completed experiments are not returned by the picker (status none after completion)', async () => { + const { experimentId } = await makeActiveExperiment({ + publicId: 'kilo/preview-completed', + }); + const caller = await createCallerForUser(admin.id); + await caller.admin.modelExperiments.complete({ id: experimentId }); + await clearRoutingCaches(); + const result = await pickModelExperimentVariant({ + publicModelId: 'kilo/preview-completed', + userId: 'user-c', + machineId: null, + clientIp: null, + }); + expect(result).toBeNull(); + // sanity: row still exists, just not routing-relevant + const [row] = await db + .select() + .from(model_experiment) + .where(eq(model_experiment.id, experimentId)); + expect(row.status).toBe('completed'); + }); +}); diff --git a/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts new file mode 100644 index 00000000000..bc73ede4fb0 --- /dev/null +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts @@ -0,0 +1,290 @@ +import { and, desc, eq, inArray, lte, sql } from 'drizzle-orm'; +import { captureException, captureMessage } from '@sentry/nextjs'; +import { + model_experiment, + model_experiment_variant, + model_experiment_variant_version, +} from '@kilocode/db/schema'; +import { db } from '@/lib/drizzle'; +import { decryptApiKey, type EncryptedData } from '@/lib/ai-gateway/byok/encryption'; +import { BYOK_ENCRYPTION_KEY } from '@/lib/config.server'; +import { getRandomNumber } from '@/lib/ai-gateway/getRandomNumber'; +import { ExperimentUpstreamSchema } from '@/lib/ai-gateway/experiments/upstream-schema'; +import type { + AllocationSubject, + ExperimentStatus, + PickVariantInput, + PickVariantResult, + ResolveResult, + RoutingVariant, +} from '@/lib/ai-gateway/experiments/pick-variant.types'; + +/** + * Returns the routing-relevant experiment for `publicId` (status active or + * paused) with all variants resolved to their current + * `model_experiment_variant_version`. This only runs after the Redis-backed + * membership pre-check (`isPublicIdExperimented` in `./membership.ts`) says + * the public id is experiment-routed, so preview traffic pays the Postgres + * lookup while normal models avoid it. + */ +export async function getRoutingExperimentForPublicId(publicId: string): Promise { + try { + return await loadExperimentFromDb(publicId); + } catch (err) { + captureException(err, { + tags: { source: 'model-experiments', operation: 'getRoutingExperimentForPublicId' }, + extra: { publicId }, + }); + return { kind: 'unavailable' }; + } +} + +async function loadExperimentFromDb(publicId: string): Promise { + const [experiment] = await db + .select({ + id: model_experiment.id, + public_model_id: model_experiment.public_model_id, + status: model_experiment.status, + }) + .from(model_experiment) + .where( + and( + eq(model_experiment.public_model_id, publicId), + inArray(model_experiment.status, ['active', 'paused']) + ) + ) + .limit(1); + + if (!experiment) { + return { kind: 'none' }; + } + + const status = experiment.status as ExperimentStatus; + + // Variants for this experiment. + const variantRows = await db + .select({ + id: model_experiment_variant.id, + weight: model_experiment_variant.weight, + }) + .from(model_experiment_variant) + .where(eq(model_experiment_variant.experiment_id, experiment.id)) + .orderBy(model_experiment_variant.id); + + if (variantRows.length === 0) { + // Active/paused with no variants is an invariant violation; admin + // activation should never permit this. Fail closed. + captureMessage('Routing-relevant experiment has no variants', { + level: 'error', + tags: { source: 'model-experiments' }, + extra: { experimentId: experiment.id, publicId }, + }); + return { kind: 'unavailable' }; + } + + const variantIds = variantRows.map(v => v.id); + + // Resolve "current version" per variant: the latest version row whose + // effective_at <= now() per (variant_id, effective_at desc, id desc). + // Postgres SELECT DISTINCT ON for one query, no per-variant round trips. + const { rows: versionRows } = await db.execute<{ + id: string; + variant_id: string; + upstream: unknown; + encrypted_api_key: EncryptedData; + }>(sql` + SELECT DISTINCT ON (${model_experiment_variant_version.variant_id}) + ${model_experiment_variant_version.id} AS id, + ${model_experiment_variant_version.variant_id} AS variant_id, + ${model_experiment_variant_version.upstream} AS upstream, + ${model_experiment_variant_version.encrypted_api_key} AS encrypted_api_key + FROM ${model_experiment_variant_version} + WHERE ${inArray(model_experiment_variant_version.variant_id, variantIds)} + AND ${lte(model_experiment_variant_version.effective_at, sql`now()`)} + ORDER BY ${model_experiment_variant_version.variant_id}, + ${desc(model_experiment_variant_version.effective_at)}, + ${desc(model_experiment_variant_version.id)} + `); + + const versionByVariantId = new Map< + string, + { id: string; upstream: unknown; encryptedApiKey: EncryptedData } + >(); + for (const r of versionRows) { + versionByVariantId.set(r.variant_id, { + id: r.id, + upstream: r.upstream, + encryptedApiKey: r.encrypted_api_key, + }); + } + + // Every variant must have a current version. + for (const v of variantRows) { + if (!versionByVariantId.has(v.id)) { + captureMessage('Routing-relevant experiment variant missing current version', { + level: 'error', + tags: { source: 'model-experiments' }, + extra: { experimentId: experiment.id, variantId: v.id, publicId }, + }); + return { kind: 'unavailable' }; + } + } + + // We deliberately DO NOT call decryptApiKey here. Decryption happens + // per-pick so key rotation takes effect on the next request. + const variants: RoutingVariant[] = []; + for (const v of variantRows) { + const ver = versionByVariantId.get(v.id); + if (!ver) { + // Already verified above; this is a defensive belt-and-suspenders. + return { kind: 'unavailable' }; + } + const parsedUpstream = ExperimentUpstreamSchema.safeParse(ver.upstream); + if (!parsedUpstream.success) { + captureMessage('Failed to parse experiment variant upstream blob', { + level: 'error', + tags: { source: 'model-experiments' }, + extra: { + experimentId: experiment.id, + variantId: v.id, + variantVersionId: ver.id, + issues: parsedUpstream.error.issues, + }, + }); + return { kind: 'unavailable' }; + } + variants.push({ + variantId: v.id, + weight: v.weight, + variantVersionId: ver.id, + upstream: parsedUpstream.data, + encryptedApiKey: ver.encryptedApiKey, + }); + } + + return { + kind: 'experiment', + experiment: { + experimentId: experiment.id, + publicModelId: experiment.public_model_id, + status, + variants, + }, + }; +} + +/** + * Picks a variant for the request, or returns a status that the caller can + * map to a local error response. + * + * - `active`: variant chosen, return the upstream blob and metadata. + * - `not-found`: experiment exists but is paused; route as 404 instead of + * silently falling through to default routing. + * - `unavailable`: cache/DB/config failure or no allocation subject. Map to + * 503 temporarily-unavailable. + * + * Returns `null` only when Postgres/cache state proves the public id is not + * currently routed by an experiment (caller should continue with non- + * experiment routing). + */ +export async function pickModelExperimentVariant( + input: PickVariantInput +): Promise { + const resolved = await getRoutingExperimentForPublicId(input.publicModelId); + if (resolved.kind === 'none') return null; + if (resolved.kind === 'unavailable') return { status: 'unavailable' }; + + const exp = resolved.experiment; + if (exp.status === 'paused') { + return { status: 'not-found' }; + } + + const allocationSubject = pickAllocationSubject(input); + if (!allocationSubject) { + captureMessage('Experiment request missing all allocation subjects', { + level: 'error', + tags: { source: 'model-experiments' }, + extra: { experimentId: exp.experimentId, publicModelId: exp.publicModelId }, + }); + return { status: 'unavailable' }; + } + + const totalWeight = exp.variants.reduce((sum, v) => sum + v.weight, 0); + if (totalWeight <= 0) { + captureMessage('Experiment total weight is non-positive', { + level: 'error', + tags: { source: 'model-experiments' }, + extra: { experimentId: exp.experimentId }, + }); + return { status: 'unavailable' }; + } + + const seed = `model_exp_${exp.experimentId}_${allocationSubject.subject}_${allocationSubject.value}`; + const bucket = getRandomNumber(seed, totalWeight); + + // Walk variants in id-asc order (cache layer already sorted them) and + // pick the variant whose cumulative weight first exceeds the bucket. + let cumulative = 0; + for (const v of exp.variants) { + cumulative += v.weight; + if (bucket < cumulative) { + // Decrypt only the chosen variant's key, here and now. This is the + // ONLY decryption point for experiment routing — the cache holds + // ciphertext exclusively. Doing it per-pick (a) keeps the + // key rotation effective immediately (next request after the key + // flips fails closed), and (b) means we never decrypt N-1 keys we + // won't use. + let apiKey: string; + try { + apiKey = decryptApiKey(v.encryptedApiKey, BYOK_ENCRYPTION_KEY); + } catch (err) { + captureException(err, { + tags: { source: 'model-experiments', operation: 'decryptApiKey' }, + extra: { + experimentId: exp.experimentId, + variantId: v.variantId, + variantVersionId: v.variantVersionId, + }, + }); + return { status: 'unavailable' }; + } + return { + status: 'active', + experimentId: exp.experimentId, + variantId: v.variantId, + variantVersionId: v.variantVersionId, + upstream: { ...v.upstream, api_key: apiKey }, + allocationSubject: allocationSubject.subject, + }; + } + } + + // Defensive fail-closed: `getRandomNumber` is contractually expected to + // return a finite value in `[0, totalWeight)`, so the loop above should + // always select a variant. We still handle the no-selection case rather + // than silently routing as unexperimented if that contract ever breaks + // (NaN, negative, weight rounding) — billing-touching code shouldn't + // depend on a math invariant we can't statically prove. + captureMessage('Experiment bucket walk did not select a variant', { + level: 'error', + tags: { source: 'model-experiments' }, + extra: { experimentId: exp.experimentId, bucket, totalWeight }, + }); + return { status: 'unavailable' }; +} + +function pickAllocationSubject( + input: PickVariantInput +): { subject: AllocationSubject; value: string } | null { + if (input.userId) return { subject: 'user', value: input.userId }; + if (input.machineId) return { subject: 'machine', value: input.machineId }; + if (input.clientIp) return { subject: 'ip', value: input.clientIp }; + return null; +} + +export type { + AllocationSubject, + ModelExperiment, + PickVariantInput, + PickVariantResult, +} from '@/lib/ai-gateway/experiments/pick-variant.types'; diff --git a/apps/web/src/lib/ai-gateway/experiments/pick-variant.types.ts b/apps/web/src/lib/ai-gateway/experiments/pick-variant.types.ts new file mode 100644 index 00000000000..1e4c669d37d --- /dev/null +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.types.ts @@ -0,0 +1,49 @@ +import type { ModelExperiment } from '@kilocode/db/schema'; +import type { EncryptedData } from '@/lib/ai-gateway/byok/encryption'; +import type { ResolvedExperimentUpstream } from '@/lib/ai-gateway/experiments/build-direct-provider'; +import type { ExperimentUpstream } from '@/lib/ai-gateway/experiments/upstream-schema'; + +export type ExperimentStatus = 'active' | 'paused'; + +export type RoutingVariant = { + variantId: string; + weight: number; + variantVersionId: string; + upstream: ExperimentUpstream; + encryptedApiKey: EncryptedData; +}; + +export type RoutingExperiment = { + experimentId: string; + publicModelId: string; + status: ExperimentStatus; + variants: RoutingVariant[]; +}; + +export type ResolveResult = + | { kind: 'experiment'; experiment: RoutingExperiment } + | { kind: 'none' } + | { kind: 'unavailable' }; + +export type AllocationSubject = 'user' | 'machine' | 'ip'; + +export type PickVariantInput = { + publicModelId: string; + userId: string | null; + machineId: string | null; + clientIp: string | null; +}; + +export type PickVariantResult = + | { + status: 'active'; + experimentId: string; + variantId: string; + variantVersionId: string; + upstream: ResolvedExperimentUpstream; + allocationSubject: AllocationSubject; + } + | { status: 'not-found' } + | { status: 'unavailable' }; + +export type { ModelExperiment }; diff --git a/apps/web/src/lib/ai-gateway/is-free-model.ts b/apps/web/src/lib/ai-gateway/is-free-model.ts index c19ede168aa..cbd8eab678b 100644 --- a/apps/web/src/lib/ai-gateway/is-free-model.ts +++ b/apps/web/src/lib/ai-gateway/is-free-model.ts @@ -1,16 +1,16 @@ import { KILO_AUTO_FREE_MODEL } from '@/lib/ai-gateway/auto-model'; import { isKiloExclusiveFreeModel, isOpenRouterStealthModel } from '@/lib/ai-gateway/models'; +import { isPublicIdExperimented } from '@/lib/ai-gateway/experiments/membership'; /** * Returns true if `model` should be treated as free for the requesting user - * this request. + * this request — including dedicated experimented public ids, which are + * partner/Kilo-funded for v1. * - * Server-only: future implementations will consult external state (e.g. a - * Redis-backed membership set for partner-funded preview model ids) to - * answer for models that are conditionally free. Lives outside `models.ts` - * so client bundles importing the model-id constants - * (`PRIMARY_DEFAULT_MODEL`, `preferredModels`, …) from `models.ts` don't - * transitively pull in those server-only dependencies. + * Server-only: consults a Redis-backed membership set for experiment routing. + * Lives outside `models.ts` so client bundles importing the model-id + * constants (`PRIMARY_DEFAULT_MODEL`, `preferredModels`, …) from `models.ts` + * don't transitively pull in the Redis client. */ export async function isFreeModel(model: string): Promise { return ( @@ -18,6 +18,7 @@ export async function isFreeModel(model: string): Promise { model === KILO_AUTO_FREE_MODEL.id || (model ?? '').endsWith(':free') || model === 'openrouter/free' || - isOpenRouterStealthModel(model ?? '') + isOpenRouterStealthModel(model ?? '') || + (await isPublicIdExperimented(model ?? '')) ); } diff --git a/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts b/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts index 2b5f027732d..96461b82d89 100644 --- a/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts +++ b/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts @@ -43,6 +43,7 @@ import { detectContextOverflow } from '@/lib/ai-gateway/context-overflow'; import { KILO_AUTO_BALANCED_MODEL, KILO_AUTO_FREE_MODEL } from '@/lib/ai-gateway/auto-model'; import type { GatewayChatApiKind, ProviderId } from '@/lib/ai-gateway/providers/types'; import { computeOpenRouterCostFields } from '@/lib/ai-gateway/processUsage.shared'; +import { persistExperimentAttribution } from '@/lib/ai-gateway/experiments/persist'; export { proxyErrorTypeSchema, ProxyErrorType } from '@/lib/proxy-error-types'; import { ProxyErrorType } from '@/lib/proxy-error-types'; @@ -308,7 +309,31 @@ export function accountForMicrodollarUsage( ) { const logFileExtension = usageContext.isStreaming ? '.log.resp.sse' : '.log.resp.json'; debugSaveProxyResponseStream(clonedReponse, logFileExtension); - after(countAndStoreUsage(clonedReponse, usageContext, openrouterRequestSpan)); + after( + countAndStoreUsage(clonedReponse, usageContext, openrouterRequestSpan).then( + async usageIdentity => { + // Chain the experiment-attribution write after the microdollar + // write. This is best-effort analytics: failures here MUST NOT + // roll back the billing write, which has already succeeded by + // the time we reach here. `persistExperimentAttribution` + // swallows errors internally. + if ( + usageIdentity && + usageContext.modelExperimentVariantVersionId && + usageContext.modelExperimentAllocationSubject + ) { + await persistExperimentAttribution({ + usageId: usageIdentity.usageId, + createdAt: usageIdentity.createdAt, + variantVersionId: usageContext.modelExperimentVariantVersionId, + allocationSubject: usageContext.modelExperimentAllocationSubject, + clientRequestId: usageContext.clientRequestId ?? null, + capture: usageContext.experimentPromptCapture ?? null, + }); + } + } + ) + ); } export async function captureProxyError(params: { diff --git a/apps/web/src/lib/ai-gateway/processUsage.types.ts b/apps/web/src/lib/ai-gateway/processUsage.types.ts index e7bbba77db3..7ac89d06e2c 100644 --- a/apps/web/src/lib/ai-gateway/processUsage.types.ts +++ b/apps/web/src/lib/ai-gateway/processUsage.types.ts @@ -151,6 +151,30 @@ export type MicrodollarUsageContext = { * do not need to know about it. */ clientRequestId?: string | null; + /** + * The exact `model_experiment_variant_version` row that served this request. + * Set only when an experiment was applied. + */ + modelExperimentVariantVersionId?: string; + /** 'user' | 'machine' | 'ip' — recorded for reporting filters when an experiment was applied. */ + modelExperimentAllocationSubject?: 'user' | 'machine' | 'ip'; + /** + * Bounded (size-capped) capture of the canonical post-`transformRequest` + * upstream request body. Set only for experimented requests. Consumed + * by `persistExperimentAttribution` after the microdollar write. + */ + experimentPromptCapture?: ExperimentPromptCapture; +}; + +/** + * Bounded prompt capture used by the experiment attribution path. The + * full serialized body is stored as a single content-addressed blob. + * `requestKind` records which upstream API shape it was serialized for. + */ +export type ExperimentPromptCapture = { + requestKind: 'chat_completions' | 'messages' | 'responses'; + requestBodyContent: string; + wasTruncated: boolean; }; export type CoreUsageWithMetaData = { diff --git a/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.ts b/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.ts index 883c38e9175..ef4692ceae0 100644 --- a/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.ts +++ b/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.ts @@ -102,16 +102,30 @@ function applyPreferredProvider( } } +export type ApplyProviderSpecificLogicOptions = { + /** + * When true, skip the kilo-exclusive `internal_id` rewrite + provider pin + * normally applied to public ids registered in `kiloExclusiveModels`. + * Generic provider-specific request fixes and `provider.transformRequest` + * still run. + * + * Set by experiment routing because the partner upstream is selected by + * the variant version, not by the registry. + */ + skipKiloExclusiveModelSettings?: boolean; +}; + export function applyProviderSpecificLogic( provider: Provider, requestedModel: string, requestToMutate: GatewayRequest, extraHeaders: Record, userByok: BYOKResult[] | null, - originalHeaders: FraudDetectionHeaders + originalHeaders: FraudDetectionHeaders, + options: ApplyProviderSpecificLogicOptions = {} ) { const kiloExclusiveModel = kiloExclusiveModels.find(m => m.public_id === requestedModel); - if (kiloExclusiveModel) { + if (kiloExclusiveModel && !options.skipKiloExclusiveModelSettings) { applyKiloExclusiveModelSettings(requestToMutate, kiloExclusiveModel); } diff --git a/apps/web/src/lib/ai-gateway/providers/get-provider.ts b/apps/web/src/lib/ai-gateway/providers/get-provider.ts index 285f9cf84c9..1d920529bad 100644 --- a/apps/web/src/lib/ai-gateway/providers/get-provider.ts +++ b/apps/web/src/lib/ai-gateway/providers/get-provider.ts @@ -19,6 +19,23 @@ import { buildDirectProvider, inferSupportedChatApis, } from '@/lib/ai-gateway/experiments/build-direct-provider'; +import { isPublicIdExperimented } from '@/lib/ai-gateway/experiments/membership'; +import { + pickModelExperimentVariant, + type AllocationSubject, +} from '@/lib/ai-gateway/experiments/pick-variant'; + +/** + * Metadata about the experiment that resolved this provider, attached when + * routing chose a model_experiment_variant_version. Persisted in the + * `model_experiment_request` row by Phase 4 attribution. + */ +export type ExperimentRouting = { + experimentId: string; + variantId: string; + variantVersionId: string; + allocationSubject: AllocationSubject; +}; export type GetProviderProviderResult = { kind: 'provider'; @@ -28,15 +45,22 @@ export type GetProviderProviderResult = { * by direct-byok and custom_llm2 because both already require explicit * admin opt-in. */ bypassAccessCheck: boolean; + /** Skip pinning `body.provider` from the organization-determined config. + * Set for direct experiment upstreams where the partner endpoint is + * selected by the variant, not by gateway routing. */ + skipProviderPin?: boolean; + /** Skip the kilo-exclusive `internal_id` rewrite + provider pin in + * `applyProviderSpecificLogic`. Generic provider-specific request fixes + * and `provider.transformRequest` still run. */ + skipKiloExclusiveModelSettings?: boolean; + /** Present when this provider was resolved through a model experiment. */ + experiment?: ExperimentRouting; }; /** * Discriminated routing result. `not-found` maps to the local - * model-unavailable response; `unavailable` maps to a 503 - * temporarily-unavailable response. Today only the `provider` shape is - * produced; the other kinds are scaffolding for upcoming model-experiment - * routing (PR #3325) where paused/unreachable experiments need to - * short-circuit without falling through to default routing. + * model-unavailable response (used by paused experiments); `unavailable` + * maps to a 503 temporarily-unavailable response (cache/DB/config failure). */ export type GetProviderResult = | GetProviderProviderResult @@ -140,10 +164,17 @@ export type GetProviderInput = { user: User | AnonymousUserContext; organizationId: string | undefined; taskId: string | undefined; + /** Resolved client IP from the route handler. Used as the IP-cohort + * allocation subject for experiment routing when no userId/machineId + * is available. */ + clientIp: string | null; + /** Machine identifier from `x-kilocode-machineid`. Used as the machine- + * cohort allocation subject for experiment routing. */ + machineId: string | null; }; export async function getProvider(input: GetProviderInput): Promise { - const { requestedModel, request, user, organizationId, taskId } = input; + const { requestedModel, request, user, organizationId, taskId, clientIp, machineId } = input; const directByokByok = await checkDirectBYOK(user, requestedModel, organizationId); if (directByokByok) { @@ -160,6 +191,44 @@ export async function getProvider(input: GetProviderInput): Promise 0; +} + +/** + * Write `content` under its sha256 hex digest if the object does not + * already exist. Returns the digest used as the object key. + * + * Concurrent races between two requests writing the same content are + * harmless — the worst case is two simultaneous PUTs with identical bytes + * and the same final state in R2. + * + * Throws on R2 errors; callers must catch and store `__failed__` so the + * experiment attribution row still lands. + */ +export async function putPromptIfAbsent(content: string): Promise { + if (!isBucketConfigured()) { + throw new Error('R2_EXPERIMENT_PROMPTS_BUCKET_NAME is not configured'); + } + const sha = sha256Hex(content); + // HEAD-then-PUT is the same shape used by cli-sessions copyBlobs and is + // sufficient for content-addressed dedup. Strong consistency on R2 is not + // guaranteed across regions, but the put-on-miss collisions are + // idempotent (same content, same key) so any race is harmless. + try { + await r2Client.send( + new HeadObjectCommand({ + Bucket: r2ExperimentPromptsBucketName, + Key: sha, + }) + ); + return sha; + } catch (err) { + if (!isNotFoundError(err)) { + throw err; + } + } + await r2Client.send( + new PutObjectCommand({ + Bucket: r2ExperimentPromptsBucketName, + Key: sha, + Body: content, + ContentType: 'application/json; charset=utf-8', + }) + ); + return sha; +} + +export async function putPromptOrNull(content: string): Promise { + try { + return await putPromptIfAbsent(content); + } catch (err) { + captureException(err, { + tags: { source: 'experiment-prompts', operation: 'putPromptIfAbsent' }, + }); + return null; + } +} + +/** + * Reads the prompt content for a given sha256 hex digest, or returns null + * when the object does not exist. Sentinels (`__absent__`, `__failed__`, + * `__deleted__`) MUST be filtered before calling this. + */ +export async function getPromptByHash(sha: string): Promise { + if (!isBucketConfigured()) { + return null; + } + if (!/^[0-9a-f]{64}$/.test(sha)) { + throw new Error('getPromptByHash requires a 64-char lowercase hex sha256'); + } + try { + const response = await r2Client.send( + new GetObjectCommand({ + Bucket: r2ExperimentPromptsBucketName, + Key: sha, + }) + ); + if (!response.Body) return null; + return await response.Body.transformToString(); + } catch (err) { + if (isNotFoundError(err)) return null; + captureException(err, { + tags: { source: 'experiment-prompts', operation: 'getPromptByHash' }, + extra: { sha }, + }); + throw err; + } +} + +function isNotFoundError(err: unknown): boolean { + if (typeof err !== 'object' || err === null) return false; + const e = err as { name?: string; $metadata?: { httpStatusCode?: number } }; + return e.name === 'NotFound' || e.$metadata?.httpStatusCode === 404; +} diff --git a/apps/web/src/lib/redis-keys.ts b/apps/web/src/lib/redis-keys.ts index 65e72a22b55..d60790697f1 100644 --- a/apps/web/src/lib/redis-keys.ts +++ b/apps/web/src/lib/redis-keys.ts @@ -48,16 +48,5 @@ export const EXPERIMENTED_PUBLIC_IDS_REDIS_KEY = redisKey( 'ai-gateway.model-experiments:experimented-public-ids' ); -/** - * Per-public-id resolved experiment payload (variants + current versions, - * with decrypted api keys merged in for hot-path use). - * - * Cached for ~10 minutes; invalidated by every admin mutation that affects - * routing for the experiment. The cached value contains decrypted upstream - * api keys, so the TTL doubles as a key-rotation lag bound. - */ -export const modelExperimentRedisKey = (publicId: string) => - redisKey(`ai-gateway.model-experiments:by-public-id:${publicId}`); - export const gitLabOAuthCredentialsRedisKey = (credentialRef: string) => redisKey(`auth-credentials:gitlab:${credentialRef}`); diff --git a/apps/web/src/routers/admin/model-experiments-router.ts b/apps/web/src/routers/admin/model-experiments-router.ts index 6fba3c40a8d..8d027674b42 100644 --- a/apps/web/src/routers/admin/model-experiments-router.ts +++ b/apps/web/src/routers/admin/model-experiments-router.ts @@ -10,8 +10,8 @@ import { import { encryptApiKey } from '@/lib/ai-gateway/byok/encryption'; import { BYOK_ENCRYPTION_KEY } from '@/lib/config.server'; import { ExperimentUpstreamSchema } from '@/lib/ai-gateway/experiments/upstream-schema'; -import { EXPERIMENTED_PUBLIC_IDS_REDIS_KEY, modelExperimentRedisKey } from '@/lib/redis-keys'; -import { redisDel, redisSet } from '@/lib/redis'; +import { EXPERIMENTED_PUBLIC_IDS_REDIS_KEY } from '@/lib/redis-keys'; +import { redisSet } from '@/lib/redis'; import { TRPCError } from '@trpc/server'; import { and, asc, desc, eq, inArray, sql } from 'drizzle-orm'; import * as z from 'zod'; @@ -111,7 +111,7 @@ async function applyExperimentTransition(opts: { .where(eq(model_experiment.id, opts.id)) .returning(); if (!updated) notFound('Experiment'); - await invalidateExperimentCaches(updated.public_model_id); + await refreshExperimentedPublicIdsCache(); return updated; } catch (err) { if (isUniqueViolation(err)) { @@ -124,15 +124,8 @@ async function applyExperimentTransition(opts: { } } -async function invalidateExperimentCaches(publicModelId: string) { - // Per-public-id cache and the membership set are both touched by every - // routing-affecting mutation. Best-effort — Redis being down does not block - // admin writes. - try { - await redisDel(modelExperimentRedisKey(publicModelId)); - } catch { - // already captured by redis helper - } +async function refreshExperimentedPublicIdsCache() { + // Best-effort — Redis being down does not block admin writes. try { await recomputeExperimentedPublicIds(); } catch { @@ -143,7 +136,7 @@ async function invalidateExperimentCaches(publicModelId: string) { // ---- Selectors ---------------------------------------------------------- // NEVER select encrypted_api_key here. Plaintext keys are decrypted only by -// the gateway-side cache loader (Phase 3); admin reads must not see them. +// gateway request routing; admin reads must not see them. const variantVersionPublicColumns = { id: model_experiment_variant_version.id, variant_id: model_experiment_variant_version.variant_id, @@ -371,11 +364,10 @@ export const adminModelExperimentsRouter = createTRPCRouter({ .set(next) .where(eq(model_experiment.id, input.id)) .returning(); - // Only routing-relevant edits touch the per-public-id cache; cosmetic - // name/description-only changes don't invalidate. + // Only routing-relevant edits touch the experimented-public-id cache; + // cosmetic name/description-only changes don't refresh it. if (existing.public_model_id !== updated.public_model_id) { - await invalidateExperimentCaches(updated.public_model_id); - await invalidateExperimentCaches(existing.public_model_id); + await refreshExperimentedPublicIdsCache(); } return updated; }), @@ -519,7 +511,7 @@ export const adminModelExperimentsRouter = createTRPCRouter({ created_by: ctx.user.id, }) .returning(variantVersionPublicColumns); - await invalidateExperimentCaches(experiment.public_model_id); + await refreshExperimentedPublicIdsCache(); return inserted; }), @@ -560,7 +552,7 @@ export const adminModelExperimentsRouter = createTRPCRouter({ created_by: ctx.user.id, }) .returning(variantVersionPublicColumns); - await invalidateExperimentCaches(experiment.public_model_id); + await refreshExperimentedPublicIdsCache(); return inserted; }), });