From eff9e059372fe5ffe4624334bffb7e3dbe81ef85 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 19 May 2026 12:06:34 +0200 Subject: [PATCH 01/49] feat(gateway): capture x-kilo-request and x-kilo-session headers Phase 2 of the model-experiments rollout. Captures the per-message client request id (used as the join key for PostHog feedback and for `model_experiment_request.client_request_id` on experimented traffic) and accepts `x-kilo-session` as a fallback session id when `x-kilocode-taskid` is absent. Also extends MicrodollarUsageContext with the experiment routing fields (`modelExperimentVariantVersionId`, `modelExperimentAllocationSubject`, `experimentPromptCapture`, `clientRequestId`) so later phases can plumb them through without touching every usage-context construction site. --- .../src/app/api/openrouter/[...path]/route.ts | 10 ++++++- .../src/lib/ai-gateway/processUsage.types.ts | 30 +++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/apps/web/src/app/api/openrouter/[...path]/route.ts b/apps/web/src/app/api/openrouter/[...path]/route.ts index b063c4bb7f2..a9be16e7ac2 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.ts @@ -223,6 +223,13 @@ export async function POST(request: NextRequest): Promise Date: Tue, 19 May 2026 12:09:23 +0200 Subject: [PATCH 02/49] refactor(gateway): extract buildDirectProvider used by custom_llm path Phase 3a of the model-experiments rollout. Pulls the kilo-internal/... (custom_llm2) Provider construction into a shared `buildDirectProvider` helper under `apps/web/src/lib/ai-gateway/experiments/`. The experiment routing path added in Phase 3b will consume the same helper, so direct-to- upstream traffic for custom LLMs and experiments share one implementation. Behavior is unchanged for custom_llm callers. --- .../experiments/build-direct-provider.ts | 101 ++++++++++++++++++ .../lib/ai-gateway/providers/get-provider.ts | 79 ++++---------- 2 files changed, 122 insertions(+), 58 deletions(-) create mode 100644 apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts diff --git a/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts b/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts new file mode 100644 index 00000000000..b0836b52fe7 --- /dev/null +++ b/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts @@ -0,0 +1,101 @@ +import { + addCacheBreakpoints, + injectReasoningIntoContent, +} from '@/lib/ai-gateway/providers/openrouter/request-helpers'; +import type { CustomLlmProvider, OpenClawApiAdapter } from '@kilocode/db'; +import type { GatewayChatApiKind, Provider } from '@/lib/ai-gateway/providers/types'; +import type { ExperimentUpstream } from '@/lib/ai-gateway/experiments/upstream-schema'; + +/** + * Maps adapter settings to the supported chat APIs for the resulting Provider. + * Mirrors the same inference performed for `custom_llm2` rows in + * `apps/web/src/lib/ai-gateway/providers/get-provider.ts`. + */ +export function inferSupportedChatApis( + aiSdkProvider: CustomLlmProvider | undefined, + openClawApiAdapter: OpenClawApiAdapter | undefined +): ReadonlyArray { + const result = new Array(); + if (aiSdkProvider === 'openai' || openClawApiAdapter === 'openai-responses') { + result.push('responses'); + } + if (aiSdkProvider === 'anthropic' || openClawApiAdapter === 'anthropic-messages') { + result.push('messages'); + } + if ( + aiSdkProvider === 'openai-compatible' || + aiSdkProvider === 'alibaba' || + aiSdkProvider === 'openrouter' || + openClawApiAdapter === 'openai-completions' || + result.length === 0 + ) { + result.push('chat_completions'); + } + return result; +} + +/** + * Plain in-memory shape: an `ExperimentUpstream` (no key) merged with the + * decrypted partner-issued api key. + * + * The cache loader (`getRoutingExperimentForPublicId`) decrypts + * `model_experiment_variant_version.encrypted_api_key` once and stores the + * resulting plaintext alongside the rest of the upstream blob in this shape + * for hot-path use. The plaintext NEVER touches Postgres or any tRPC + * response — only this in-memory record and the per-public-id Redis cache. + */ +export type ResolvedExperimentUpstream = ExperimentUpstream & { api_key: string }; + +/** + * Input to `buildDirectProvider`. A superset of `ResolvedExperimentUpstream` + * that also accepts `extra_headers`, which is used by the custom_llm2 + * (`kilo-internal/...`) code path but not by experiments. Experiment + * upstreams must NOT pass `extra_headers` — see + * `ExperimentUpstreamSchema`. + */ +export type DirectProviderInput = ResolvedExperimentUpstream & { + extra_headers?: Record; +}; + +/** + * Builds a `Provider` that points directly at a partner-issued upstream. + * + * Used by both the experiment routing path and the existing + * `kilo-internal/...` (custom_llm2) path: `id: 'custom'`, `apiUrl` = + * upstream `base_url`, `apiKey` = upstream-issued key, supported chat APIs + * inferred from adapter settings. + * + * Direct traffic goes to `apiUrl` — OpenRouter and Vercel are never + * contacted. The route layer is responsible for not applying provider + * pinning or kilo-exclusive model rewrites on top of this provider. + */ +export function buildDirectProvider(upstream: DirectProviderInput): Provider { + return { + id: 'custom', + apiUrl: upstream.base_url, + apiKey: upstream.api_key, + supportedChatApis: inferSupportedChatApis( + upstream.opencode_settings?.ai_sdk_provider, + upstream.openclaw_settings?.api_adapter + ), + transformRequest(context) { + if (upstream.remove_from_body) { + const body = context.request.body as Record; + for (const key of upstream.remove_from_body) { + delete body[key]; + } + } + Object.assign(context.request.body, upstream.extra_body ?? {}); + if (upstream.extra_headers) { + Object.assign(context.extraHeaders, upstream.extra_headers); + } + context.request.body.model = upstream.internal_id; + if (upstream.add_cache_breakpoints) { + addCacheBreakpoints(context.request); + } + if (upstream.inject_reasoning_into_content) { + injectReasoningIntoContent(context.request); + } + }, + }; +} 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 5800fb78c63..9ac26891bdf 100644 --- a/apps/web/src/lib/ai-gateway/providers/get-provider.ts +++ b/apps/web/src/lib/ai-gateway/providers/get-provider.ts @@ -11,41 +11,14 @@ import { db } from '@/lib/drizzle'; import { eq } from 'drizzle-orm'; import type { AnonymousUserContext } from '@/lib/anonymous'; import { isAnonymousContext } from '@/lib/anonymous'; -import type { BYOKResult, GatewayChatApiKind, Provider } from '@/lib/ai-gateway/providers/types'; +import type { BYOKResult, Provider } from '@/lib/ai-gateway/providers/types'; import PROVIDERS from '@/lib/ai-gateway/providers/provider-definitions'; import { getDirectByokModel } from '@/lib/ai-gateway/providers/direct-byok'; +import { CustomLlmDefinitionSchema } from '@kilocode/db'; import { - CustomLlmDefinitionSchema, - type OpenClawApiAdapter, - type CustomLlmProvider, -} from '@kilocode/db'; -import { - addCacheBreakpoints, - injectReasoningIntoContent, -} from '@/lib/ai-gateway/providers/openrouter/request-helpers'; - -function inferSupportedChatApis( - aiSdkProvider: CustomLlmProvider | undefined, - openClawApiAdapter: OpenClawApiAdapter | undefined -): ReadonlyArray { - const result = new Array(); - if (aiSdkProvider === 'openai' || openClawApiAdapter === 'openai-responses') { - result.push('responses'); - } - if (aiSdkProvider === 'anthropic' || openClawApiAdapter === 'anthropic-messages') { - result.push('messages'); - } - if ( - aiSdkProvider === 'openai-compatible' || - aiSdkProvider === 'alibaba' || - aiSdkProvider === 'openrouter' || - openClawApiAdapter === 'openai-completions' || - result.length === 0 - ) { - result.push('chat_completions'); - } - return result; -} + buildDirectProvider, + inferSupportedChatApis, +} from '@/lib/ai-gateway/experiments/build-direct-provider'; async function checkDirectBYOK( user: User | AnonymousUserContext, @@ -95,32 +68,22 @@ async function checkCustomLlm( return null; } return { - provider: { - id: 'custom', - apiUrl: customLlm.base_url, - apiKey: customLlm.api_key, - supportedChatApis: inferSupportedChatApis( - customLlm.opencode_settings?.ai_sdk_provider, - customLlm.openclaw_settings?.api_adapter - ), - transformRequest(context) { - if (customLlm.remove_from_body) { - const body = context.request.body as Record; - for (const key of customLlm.remove_from_body ?? []) { - delete body[key]; - } - } - Object.assign(context.request.body, customLlm.extra_body ?? {}); - Object.assign(context.extraHeaders, customLlm.extra_headers ?? {}); - context.request.body.model = customLlm.internal_id; - if (customLlm.add_cache_breakpoints) { - addCacheBreakpoints(context.request); - } - if (customLlm.inject_reasoning_into_content) { - injectReasoningIntoContent(context.request); - } - }, - }, + provider: buildDirectProvider({ + internal_id: customLlm.internal_id, + base_url: customLlm.base_url, + api_key: customLlm.api_key, + opencode_settings: customLlm.opencode_settings + ? { ai_sdk_provider: customLlm.opencode_settings.ai_sdk_provider } + : undefined, + openclaw_settings: customLlm.openclaw_settings + ? { api_adapter: customLlm.openclaw_settings.api_adapter } + : undefined, + extra_body: customLlm.extra_body, + extra_headers: customLlm.extra_headers, + remove_from_body: customLlm.remove_from_body, + add_cache_breakpoints: customLlm.add_cache_breakpoints, + inject_reasoning_into_content: customLlm.inject_reasoning_into_content, + }), userByok: null, bypassAccessCheck: true, }; From 59a7924b34725d40e2f5ced3c37b4d4bcc1b6ffa Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 19 May 2026 12:13:52 +0200 Subject: [PATCH 03/49] feat(gateway): add experiment variant picker and routing-cache loader MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 3b of the model-experiments rollout. Adds `apps/web/src/lib/ai-gateway/experiments/pick-variant.ts` with: - `isPublicIdExperimented(publicId)` — fast Redis-backed membership check with Postgres fallback, returning `unavailable` when both cache and DB are inaccessible so the gateway can fail closed. - `getRoutingExperimentForPublicId(publicId)` — Redis-cached resolved experiment for a public id, including a single Postgres `DISTINCT ON (variant_id)` query that picks each variant's current `model_experiment_variant_version` and decrypts `encrypted_api_key` once at cache-build time. Returns `none` (no routing-relevant experiment), `experiment` (active or paused), or `unavailable` (config/decrypt/validation failure or BYOK_ENCRYPTION_KEY missing). - `pickModelExperimentVariant(input)` — deterministic picker using `getRandomNumber`. Allocation subject precedence is user → machine → ip; missing all subjects returns `unavailable`. Variants are walked by id-asc cumulative weight so live label edits never rebucket users. The `upstream` JSONB blob is validated by ExperimentUpstreamSchema before use; partner-issued api keys live only in the encrypted column and the cached payload, never in tRPC responses. --- .../src/lib/ai-gateway/experiments/index.ts | 15 + .../ai-gateway/experiments/pick-variant.ts | 461 ++++++++++++++++++ 2 files changed, 476 insertions(+) create mode 100644 apps/web/src/lib/ai-gateway/experiments/index.ts create mode 100644 apps/web/src/lib/ai-gateway/experiments/pick-variant.ts diff --git a/apps/web/src/lib/ai-gateway/experiments/index.ts b/apps/web/src/lib/ai-gateway/experiments/index.ts new file mode 100644 index 00000000000..ec106d4a054 --- /dev/null +++ b/apps/web/src/lib/ai-gateway/experiments/index.ts @@ -0,0 +1,15 @@ +export { + buildDirectProvider, + inferSupportedChatApis, + type ResolvedExperimentUpstream, + type DirectProviderInput, +} from './build-direct-provider'; +export { + isPublicIdExperimented, + getRoutingExperimentForPublicId, + pickModelExperimentVariant, + type AllocationSubject, + type PickVariantInput, + type PickVariantResult, +} from './pick-variant'; +export { ExperimentUpstreamSchema, type ExperimentUpstream } from './upstream-schema'; 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..c1b0f9f8cbd --- /dev/null +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts @@ -0,0 +1,461 @@ +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, + type ModelExperiment, +} from '@kilocode/db/schema'; +import { db } from '@/lib/drizzle'; +import { redisGet, redisSet } from '@/lib/redis'; +import { EXPERIMENTED_PUBLIC_IDS_REDIS_KEY, modelExperimentRedisKey } from '@/lib/redis-keys'; +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, + type ExperimentUpstream, +} from '@/lib/ai-gateway/experiments/upstream-schema'; +import type { ResolvedExperimentUpstream } from '@/lib/ai-gateway/experiments/build-direct-provider'; + +/** TTL for the per-public-id experiment cache. Doubles as the upper bound on + * api-key rotation lag (the cache holds decrypted plaintext). */ +const PER_PUBLIC_ID_CACHE_TTL_SECONDS = 600; + +type ExperimentStatus = 'active' | 'paused'; + +/** Per-variant payload cached in Redis. `upstream` already includes the + * decrypted `api_key` merged in for hot-path use. */ +type CachedVariant = { + variantId: string; + weight: number; + variantVersionId: string; + upstream: ResolvedExperimentUpstream; +}; + +type CachedExperiment = { + experimentId: string; + publicModelId: string; + status: ExperimentStatus; + variants: CachedVariant[]; +}; + +type ResolveResult = + | { kind: 'experiment'; experiment: CachedExperiment } + | { 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' }; + +let warnedMissingEncryptionKey = false; + +function warnMissingEncryptionKeyOnce() { + if (warnedMissingEncryptionKey) return; + warnedMissingEncryptionKey = true; + captureMessage('BYOK_ENCRYPTION_KEY is unset; model experiment routing returning unavailable', { + level: 'error', + tags: { source: 'model-experiments' }, + }); + console.error( + '[model-experiments] BYOK_ENCRYPTION_KEY is unset; experiment routing will return unavailable' + ); +} + +/** + * Fast pre-check: does any active|paused experiment target this public id? + * + * Reads from a JSON-encoded array in Redis. On Redis error, falls back to a + * single Postgres query for that public id. If both fail, returns + * `unavailable` so the caller can fail closed instead of routing as + * unexperimented traffic. + */ +export async function isPublicIdExperimented( + publicId: string +): Promise { + try { + const cached = await redisGet(EXPERIMENTED_PUBLIC_IDS_REDIS_KEY); + if (cached !== null) { + const ids = parseStringArray(cached); + if (ids === null) { + // Corrupt cache; fall through to DB. (Already captured by Sentry below.) + } else { + return ids.includes(publicId); + } + } + } catch { + // captureException already invoked by the redis helper + } + + try { + const rows = await db + .select({ id: model_experiment.id }) + .from(model_experiment) + .where( + and( + eq(model_experiment.public_model_id, publicId), + inArray(model_experiment.status, ['active', 'paused']) + ) + ) + .limit(1); + return rows.length > 0; + } catch (err) { + captureException(err, { + tags: { source: 'model-experiments', operation: 'isPublicIdExperimented' }, + extra: { publicId }, + }); + return 'unavailable'; + } +} + +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; + } +} + +/** + * Returns the routing-relevant experiment for `publicId` (status active or + * paused) with all variants resolved to their current + * `model_experiment_variant_version` and decrypted upstream blob. Returns + * `none` only after Postgres has proven there is no routing-relevant + * experiment for this id. + * + * Cached per-publicId for ~10 minutes. The cached payload contains + * decrypted partner-issued api keys, so the TTL doubles as a key-rotation + * lag bound. The cache is invalidated by every admin mutation that affects + * routing. + */ +export async function getRoutingExperimentForPublicId(publicId: string): Promise { + const cacheKey = modelExperimentRedisKey(publicId); + try { + const cached = await redisGet(cacheKey); + if (cached !== null) { + const parsed = tryParseCachedExperiment(cached); + if (parsed === 'none' || parsed === null) { + // null = corrupt; fall through to DB rebuild. + if (parsed === 'none') return { kind: 'none' }; + } else { + return { kind: 'experiment', experiment: parsed }; + } + } + } catch { + // captured by redis helper; fall through to DB + } + + let resolved: ResolveResult; + try { + resolved = await loadExperimentFromDb(publicId); + } catch (err) { + captureException(err, { + tags: { source: 'model-experiments', operation: 'getRoutingExperimentForPublicId' }, + extra: { publicId }, + }); + return { kind: 'unavailable' }; + } + + // Best-effort cache write. A failure here does not affect routing. + try { + if (resolved.kind === 'experiment') { + await redisSet( + cacheKey, + JSON.stringify(serializeCachedExperiment(resolved.experiment)), + PER_PUBLIC_ID_CACHE_TTL_SECONDS + ); + } else if (resolved.kind === 'none') { + await redisSet(cacheKey, JSON.stringify({ kind: 'none' }), PER_PUBLIC_ID_CACHE_TTL_SECONDS); + } + } catch { + // already captured + } + + return resolved; +} + +type SerializedCachedExperiment = { + kind: 'experiment'; + experimentId: string; + publicModelId: string; + status: ExperimentStatus; + variants: CachedVariant[]; +}; + +type SerializedCacheEntry = SerializedCachedExperiment | { kind: 'none' }; + +function serializeCachedExperiment(exp: CachedExperiment): SerializedCachedExperiment { + return { kind: 'experiment', ...exp }; +} + +function tryParseCachedExperiment(raw: string): CachedExperiment | 'none' | null { + try { + const parsed = JSON.parse(raw) as SerializedCacheEntry; + if (parsed.kind === 'none') return 'none'; + if (parsed.kind === 'experiment') { + return { + experimentId: parsed.experimentId, + publicModelId: parsed.publicModelId, + status: parsed.status, + variants: parsed.variants, + }; + } + return null; + } catch { + return null; + } +} + +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' }; + } + } + + if (!BYOK_ENCRYPTION_KEY) { + warnMissingEncryptionKeyOnce(); + return { kind: 'unavailable' }; + } + + const variants: CachedVariant[] = []; + 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' }; + } + let apiKey: string; + try { + apiKey = decryptApiKey(ver.encryptedApiKey, BYOK_ENCRYPTION_KEY); + } catch (err) { + captureException(err, { + tags: { source: 'model-experiments', operation: 'decryptApiKey' }, + extra: { experimentId: experiment.id, variantId: v.id, variantVersionId: ver.id }, + }); + return { kind: 'unavailable' }; + } + variants.push({ + variantId: v.id, + weight: v.weight, + variantVersionId: ver.id, + upstream: { ...(parsedUpstream.data satisfies ExperimentUpstream), api_key: apiKey }, + }); + } + + 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 subjectPick = pickAllocationSubject(input); + if (!subjectPick) { + 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}_${subjectPick.subject}_${subjectPick.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) { + return { + status: 'active', + experimentId: exp.experimentId, + variantId: v.variantId, + variantVersionId: v.variantVersionId, + upstream: v.upstream, + allocationSubject: subjectPick.subject, + }; + } + } + + // Unreachable: getRandomNumber returns [0, totalWeight). Fail closed. + 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; +} + +// Re-export for callers that need the raw experiment row type. +export type { ModelExperiment }; From 72dd511718c3505c5191cd0160aa65e87ffb5a0d Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 19 May 2026 12:31:03 +0200 Subject: [PATCH 04/49] feat(gateway): route experimented public ids to direct upstream MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 3c + 3d of the model-experiments rollout. `getProvider` now returns a discriminated `GetProviderResult` union (`provider` | `not-found` | `unavailable`) and adds a new branch for experimented public ids that runs after the BYOK checks and before kilo-internal/custom_llm and the kiloExclusiveModels lookup. Active experiments produce a direct provider built via `buildDirectProvider` plus experiment routing metadata; paused experiments map to a local model-unavailable response so traffic doesn't silently fall through to default routing; cache/DB/config failures map to a 503 temporarily-unavailable response. Experiment traffic uses three new flags on the provider result: `skipBalanceCheck` (preview traffic is provider-funded), `skipProviderPin` (the partner endpoint is selected by the variant version, not by the org config), and `skipKiloExclusiveModelSettings` (which `applyProviderSpecificLogic` honors so the registry does not override the variant's `internal_id`). Server-side organization model/data-collection policy still runs against the public id. Also excludes experimented preview public ids from `kilo-auto/free` candidate sets so they are never silently selected on a user's behalf — users only ever reach the experiment path by opting in to the dedicated preview model id. The route handler captures `x-kilocode-machineid` once into `machineIdHeader` and threads it (plus the resolved client IP) into `getProvider` so the variant picker can fall back from userId to machineId to clientIp. --- .../src/app/api/openrouter/[...path]/route.ts | 47 +++++-- .../lib/ai-gateway/auto-model/resolution.ts | 19 ++- .../apply-provider-specific-logic.ts | 18 ++- .../lib/ai-gateway/providers/get-provider.ts | 132 ++++++++++++++++-- 4 files changed, 194 insertions(+), 22 deletions(-) diff --git a/apps/web/src/app/api/openrouter/[...path]/route.ts b/apps/web/src/app/api/openrouter/[...path]/route.ts index a9be16e7ac2..d54fa900fc1 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.ts @@ -379,13 +379,33 @@ export async function POST(request: NextRequest): Promise 0, botId, @@ -462,7 +482,7 @@ export async function POST(request: NextRequest): Promise +): Promise> { + const checks = await Promise.all( + candidates.map(async id => ({ id, experimented: await isPublicIdExperimented(id) })) + ); + return checks.filter(c => c.experimented !== true).map(c => c.id); +} + function gatewaySupportsApiKind( gateway: ProviderId, apiKind: GatewayRequest['kind'] | null @@ -97,7 +112,9 @@ export async function resolveAutoModel( const { model, modeHeader, featureHeader, sessionId, apiKind, clientIp } = params; if (model === KILO_AUTO_FREE_MODEL.id) { const openRouterModels = await getOpenRouterModels(); - const candidates = getAutoFreeCandidates(openRouterModels, apiKind); + const candidates = await filterOutExperimentedPublicIds( + getAutoFreeCandidates(openRouterModels, apiKind) + ); if (candidates.length === 0) { return { kind: 'no_free_models_available' }; } 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 9ac26891bdf..15fb31cb05b 100644 --- a/apps/web/src/lib/ai-gateway/providers/get-provider.ts +++ b/apps/web/src/lib/ai-gateway/providers/get-provider.ts @@ -19,12 +19,62 @@ import { buildDirectProvider, inferSupportedChatApis, } from '@/lib/ai-gateway/experiments/build-direct-provider'; +import { + isPublicIdExperimented, + 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'; + provider: Provider; + userByok: BYOKResult[] | null; + /** Skip balance, paid-auth, and organization policy checks entirely. Used + * by direct-byok and custom_llm2 because both already require explicit + * admin opt-in. */ + bypassAccessCheck: boolean; + /** Skip only the balance check while still enforcing organization model/ + * provider policy. Set for experiment traffic (preview/provider-funded). */ + skipBalanceCheck?: 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 (used by paused experiments); `unavailable` + * maps to a 503 temporarily-unavailable response (cache/DB/config failure). + */ +export type GetProviderResult = + | GetProviderProviderResult + | { kind: 'not-found' } + | { kind: 'unavailable' }; async function checkDirectBYOK( user: User | AnonymousUserContext, requestedModel: string, organizationId: string | undefined -) { +): Promise { const { provider: directByok, model: directByokModel } = await getDirectByokModel(requestedModel); if (!directByok || !directByokModel) { return null; @@ -36,6 +86,7 @@ async function checkDirectBYOK( return null; } return { + kind: 'provider', provider: { id: 'direct-byok', apiUrl: directByok.base_url, @@ -54,7 +105,7 @@ async function checkDirectBYOK( async function checkCustomLlm( requestedModel: string, organizationId: string -): Promise<{ provider: Provider; userByok: null; bypassAccessCheck: true } | null> { +): Promise { const [row] = await db .select() .from(custom_llm2) @@ -68,6 +119,7 @@ async function checkCustomLlm( return null; } return { + kind: 'provider', provider: buildDirectProvider({ internal_id: customLlm.internal_id, base_url: customLlm.base_url, @@ -109,13 +161,24 @@ async function checkVercelBYOK( : getBYOKforUser(db, user.id, modelProviders); } -export async function getProvider( - requestedModel: string, - request: GatewayRequest, - user: User | AnonymousUserContext, - organizationId: string | undefined, - taskId: string | undefined -): Promise<{ provider: Provider; userByok: BYOKResult[] | null; bypassAccessCheck: boolean }> { +export type GetProviderInput = { + requestedModel: string; + request: GatewayRequest; + 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, clientIp, machineId } = input; + const directByokByok = await checkDirectBYOK(user, requestedModel, organizationId); if (directByokByok) { return directByokByok; @@ -124,12 +187,55 @@ export async function getProvider( const vercelByok = await checkVercelBYOK(user, requestedModel, organizationId); if (vercelByok) { return { + kind: 'provider', provider: PROVIDERS.VERCEL_AI_GATEWAY, userByok: vercelByok, bypassAccessCheck: false, }; } + // Model experiment routing for dedicated preview public ids. Runs before + // `kilo-internal/...` and the `kiloExclusiveModels` lookup so an + // experimented public id never falls through to OpenRouter/Vercel. + const experimented = await isPublicIdExperimented(requestedModel); + if (experimented === 'unavailable') { + return { kind: 'unavailable' }; + } + if (experimented === true) { + const userId = isAnonymousContext(user) ? null : user.id; + const selection = await pickModelExperimentVariant({ + publicModelId: requestedModel, + userId, + machineId, + clientIp, + }); + if (selection?.status === 'not-found') { + return { kind: 'not-found' }; + } + if (selection?.status === 'unavailable') { + return { kind: 'unavailable' }; + } + if (selection?.status === 'active') { + return { + kind: 'provider', + provider: buildDirectProvider(selection.upstream), + userByok: null, + bypassAccessCheck: false, + skipBalanceCheck: true, + skipProviderPin: true, + skipKiloExclusiveModelSettings: true, + experiment: { + experimentId: selection.experimentId, + variantId: selection.variantId, + variantVersionId: selection.variantVersionId, + allocationSubject: selection.allocationSubject, + }, + }; + } + // selection === null: cache+DB say no routing-relevant experiment for + // this id. Fall through to non-experiment routing. + } + if (requestedModel.startsWith('kilo-internal/') && organizationId) { const customLlmResult = await checkCustomLlm(requestedModel, organizationId); if (customLlmResult) { @@ -145,10 +251,16 @@ export async function getProvider( eligibleForVercelRouting && (await shouldRouteToVercel(requestedModel, request, taskId || user.id)) ) { - return { provider: PROVIDERS.VERCEL_AI_GATEWAY, userByok: null, bypassAccessCheck: false }; + return { + kind: 'provider', + provider: PROVIDERS.VERCEL_AI_GATEWAY, + userByok: null, + bypassAccessCheck: false, + }; } return { + kind: 'provider', provider: Object.values(PROVIDERS).find(p => p.id === kiloExclusiveModel?.gateway) ?? PROVIDERS.OPENROUTER, From b81a5b18c7e5327e413c325d9b302909a5fadd40 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 19 May 2026 12:31:58 +0200 Subject: [PATCH 05/49] feat(r2): add experiment-prompts content-addressed storage helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 4a of the model-experiments rollout. New helper module `apps/web/src/lib/r2/experiment-prompts.ts` exposes `putPromptIfAbsent(content)` (HEAD-then-PUT under a sha256 hex key for automatic dedup) and `getPromptByHash(sha)` (admin out-of-band reads). The bucket comes from a new env var `R2_EXPERIMENT_PROMPTS_BUCKET_NAME` exported from the existing R2 client module; per-environment buckets are `kilo-experiment-prompts-dev` and `kilo-experiment-prompts-prod`. The helper is wired up but not yet called from the request hot path — Phase 4b will add the bounded prompt capture and attribution writer that consumes it. --- apps/web/src/lib/r2/client.ts | 8 ++ apps/web/src/lib/r2/experiment-prompts.ts | 114 ++++++++++++++++++++++ 2 files changed, 122 insertions(+) create mode 100644 apps/web/src/lib/r2/experiment-prompts.ts diff --git a/apps/web/src/lib/r2/client.ts b/apps/web/src/lib/r2/client.ts index 2dcdade71f6..1ee58a78692 100644 --- a/apps/web/src/lib/r2/client.ts +++ b/apps/web/src/lib/r2/client.ts @@ -9,6 +9,13 @@ const R2_CLI_SESSIONS_BUCKET_NAME = getEnvVariable('R2_CLI_SESSIONS_BUCKET_NAME' const CLOUD_AGENT_R2_ATTACHMENTS_BUCKET_NAME = getEnvVariable( 'CLOUD_AGENT_R2_ATTACHMENTS_BUCKET_NAME' ); +// Per-environment buckets: +// dev: kilo-experiment-prompts-dev +// prod: kilo-experiment-prompts-prod +// Optional: when unset, the experiment prompt-storage path is a no-op and +// `model_experiment_request` rows record the `__failed__` sentinel for the +// affected side. Experiment attribution still lands. +const R2_EXPERIMENT_PROMPTS_BUCKET_NAME = getEnvVariable('R2_EXPERIMENT_PROMPTS_BUCKET_NAME'); if (!R2_ACCOUNT_ID) { throw new Error('R2_ACCOUNT_ID environment variable is required'); @@ -43,3 +50,4 @@ export const r2Client = new S3Client({ export const r2CliSessionsBucketName = R2_CLI_SESSIONS_BUCKET_NAME; export const r2CloudAgentAttachmentsBucketName = CLOUD_AGENT_R2_ATTACHMENTS_BUCKET_NAME; +export const r2ExperimentPromptsBucketName = R2_EXPERIMENT_PROMPTS_BUCKET_NAME; diff --git a/apps/web/src/lib/r2/experiment-prompts.ts b/apps/web/src/lib/r2/experiment-prompts.ts new file mode 100644 index 00000000000..ced4592e5ce --- /dev/null +++ b/apps/web/src/lib/r2/experiment-prompts.ts @@ -0,0 +1,114 @@ +import { createHash } from 'node:crypto'; +import { GetObjectCommand, HeadObjectCommand, PutObjectCommand } from '@aws-sdk/client-s3'; +import { captureException } from '@sentry/nextjs'; +import { r2Client, r2ExperimentPromptsBucketName } from './client'; + +/** + * Content-addressed prompt storage for model experiments. + * + * Each unique blob is written once under its sha256 hex digest as the + * object key. `model_experiment_request` rows reference only the hash, so + * deduplication is automatic across requests sharing identical content + * (e.g. a system prompt repeated across thousands of agentic turns). + * + * - Bucket: `R2_EXPERIMENT_PROMPTS_BUCKET_NAME` (one per environment). + * - Trust boundary: same R2 credentials as cli-sessions and + * cloud-agent-attachments; prompts collected under explicit experiment + * opt-in are not subject to the default microdollar_usage_metadata + * soft-delete policy. See `.plans/experimental-models-1.md` "Prompt + * Storage > GDPR and consent". + * + * Failures here MUST NOT roll back the microdollar usage write or the + * experiment attribution row. Callers store `__failed__` for the affected + * side and continue. + */ + +/** Lowercase hex sha256. Used both as content hash and as the R2 object key. */ +export function sha256Hex(content: string): string { + return createHash('sha256').update(content, 'utf8').digest('hex'); +} + +function isBucketConfigured(): boolean { + return r2ExperimentPromptsBucketName.length > 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; +} + +/** + * 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; +} From 9bf9d4e44ba971d824d37d0b786905a3566df898 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 19 May 2026 12:34:26 +0200 Subject: [PATCH 06/49] feat(gateway): persist experiment attribution after microdollar write MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 4b + 4c of the model-experiments rollout. Adds `apps/web/src/lib/ai-gateway/experiments/persist.ts` with two exports: - `buildExperimentPromptCapture(request)` produces a bounded capture of the canonical post-`transformRequest` upstream body, splitting out the leading `role: 'system'` message for chat_completions and tail-truncating either side at 4 MB. The capture is stored on `MicrodollarUsageContext` so the async `after()` path never retains a reference to the full uncapped request body. - `persistExperimentAttribution(input)` writes one row into `model_experiment_request` after running two parallel content-addressed R2 puts via `putPromptIfAbsent`, using `__absent__` / `__failed__` sentinels for missing/failed sides. Errors are reported to Sentry and swallowed — the billing write must succeed independently of analytics. `processUsage.ts` now surfaces the inserted `{ usageId, createdAt }` out of `logMicrodollarUsage` and `processTokenData` so the `accountForMicrodollarUsage` after-hook can chain `persistExperimentAttribution` on the same row. The route handler calls `buildExperimentPromptCapture` after provider-specific transforms and stashes the result on the usage context only for experimented requests. --- .../src/app/api/openrouter/[...path]/route.ts | 9 + .../src/lib/ai-gateway/experiments/index.ts | 9 + .../src/lib/ai-gateway/experiments/persist.ts | 186 ++++++++++++++++++ .../src/lib/ai-gateway/llm-proxy-helpers.ts | 28 ++- apps/web/src/lib/ai-gateway/processUsage.ts | 10 +- 5 files changed, 237 insertions(+), 5 deletions(-) create mode 100644 apps/web/src/lib/ai-gateway/experiments/persist.ts diff --git a/apps/web/src/app/api/openrouter/[...path]/route.ts b/apps/web/src/app/api/openrouter/[...path]/route.ts index d54fa900fc1..40c824c0be9 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.ts @@ -18,6 +18,7 @@ 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 { upstreamRequest } from '@/lib/ai-gateway/providers/upstream-request'; import { debugSaveProxyRequest } from '@/lib/debugUtils'; import { setTag, startInactiveSpan } from '@sentry/nextjs'; @@ -559,6 +560,14 @@ export async function POST(request: NextRequest): Promise 0 && messages[0]?.role === 'system') { + systemContent = serialize(messages[0].content ?? ''); + nonSystemMessages = messages.slice(1); + } + const remainder = { ...request.body, messages: nonSystemMessages }; + let bodyContent = serialize(remainder); + if (systemContent !== null && systemContent.length > SYSTEM_PROMPT_CAP_BYTES) { + systemContent = systemContent.slice(0, SYSTEM_PROMPT_CAP_BYTES); + wasTruncated = true; + } + if (bodyContent.length > REQUEST_BODY_CAP_BYTES) { + bodyContent = bodyContent.slice(0, REQUEST_BODY_CAP_BYTES); + wasTruncated = true; + } + if (wasTruncated) { + noteTruncation(request.body.model); + } + return { + systemPromptContent: systemContent, + requestBodyContent: bodyContent, + wasTruncated, + }; + } + + // messages / responses: fold the entire serialized body into the body + // side; system extraction for those API kinds is left for later if it + // turns out to dedup meaningfully. + let bodyContent = serialize(request.body); + if (bodyContent.length > REQUEST_BODY_CAP_BYTES) { + bodyContent = bodyContent.slice(0, REQUEST_BODY_CAP_BYTES); + wasTruncated = true; + noteTruncation(request.body.model); + } + return { + systemPromptContent: null, + requestBodyContent: bodyContent, + wasTruncated, + }; +} + +function serialize(value: unknown): string { + if (typeof value === 'string') return value; + return JSON.stringify(value); +} + +function noteTruncation(model: string | undefined) { + // Sentry breadcrumb only (no payload content) — we want to know if this + // ever fires at non-trivial rates so we can re-evaluate caps. + captureMessage('Experiment prompt capture truncated', { + level: 'info', + tags: { source: 'model-experiments', operation: 'buildExperimentPromptCapture' }, + extra: { model: model ?? null }, + }); +} + +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. Best-effort R2 puts run in parallel; only the failed side + * gets `__failed__`. + * + * Failures here MUST NOT roll back the microdollar usage write. Errors + * are reported to Sentry and swallowed. + */ +export async function persistExperimentAttribution( + input: PersistExperimentAttributionInput +): Promise { + try { + const { systemHash, bodyHash } = await storePrompts(input.capture); + 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, + system_prompt_sha256: systemHash, + request_body_sha256: bodyHash, + was_truncated: input.capture?.wasTruncated ?? false, + created_at: input.createdAt, + }); + } catch (err) { + captureException(err, { + tags: { source: 'model-experiments', operation: 'persistExperimentAttribution' }, + extra: { usageId: input.usageId, variantVersionId: input.variantVersionId }, + }); + } +} + +async function storePrompts( + capture: ExperimentPromptCapture | null +): Promise<{ systemHash: string; bodyHash: string }> { + if (!capture) { + return { systemHash: PROMPT_HASH_ABSENT, bodyHash: PROMPT_HASH_FAILED }; + } + const systemPromise = + capture.systemPromptContent === null + ? Promise.resolve({ kind: 'absent' }) + : putPromptSafely(capture.systemPromptContent); + const bodyPromise = putPromptSafely(capture.requestBodyContent); + const [systemResult, bodyResult] = await Promise.allSettled([systemPromise, bodyPromise]); + return { + systemHash: resolveHash(systemResult, 'system'), + bodyHash: resolveHash(bodyResult, 'body'), + }; +} + +type PromptSlotResult = { kind: 'absent' } | { kind: 'sha'; sha: string } | { kind: 'failed' }; + +async function putPromptSafely(content: string): Promise { + try { + return { kind: 'sha', sha: await putPromptIfAbsent(content) }; + } catch (err) { + captureException(err, { + tags: { source: 'model-experiments', operation: 'putPromptIfAbsent' }, + }); + return { kind: 'failed' }; + } +} + +function resolveHash( + settled: PromiseSettledResult, + side: 'system' | 'body' +): string { + if (settled.status === 'rejected') { + captureException(settled.reason, { + tags: { source: 'model-experiments', operation: 'storePrompts', side }, + }); + return PROMPT_HASH_FAILED; + } + const result = settled.value; + if (result.kind === 'absent') return PROMPT_HASH_ABSENT; + if (result.kind === 'sha') return result.sha; + return PROMPT_HASH_FAILED; +} 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..7278a8d3c4d 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,32 @@ 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 + ) { + return; + } + 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.ts b/apps/web/src/lib/ai-gateway/processUsage.ts index 35005f0aed7..a4c82232103 100644 --- a/apps/web/src/lib/ai-gateway/processUsage.ts +++ b/apps/web/src/lib/ai-gateway/processUsage.ts @@ -245,7 +245,7 @@ export function toInsertableDbUsageRecord( export async function logMicrodollarUsage( usageStats: MicrodollarUsageStats, usageContext: MicrodollarUsageContext -) { +): Promise<{ usageId: string; createdAt: string }> { usageContext.status_code = usageStats.status_code; const contextInfo = extractUsageContextInfo(usageContext); const { core, metadata } = toInsertableDbUsageRecord(usageStats, contextInfo); @@ -256,6 +256,8 @@ export async function logMicrodollarUsage( usageContext.prior_microdollar_usage, usageContext.posthog_distinct_id ?? null ); + + return { usageId: core.id, createdAt: core.created_at }; } async function saveUsageRelatedData( @@ -909,14 +911,14 @@ export function calculateKiloExclusiveCost_mUsd( export async function processTokenData( usageStats: MicrodollarUsageStats | null, usageContext: MicrodollarUsageContext -) { +): Promise<{ usageId: string; createdAt: string } | null> { if (!usageStats) { captureMessage('SUSPICIOUS: No usage information', { level: 'error', tags: { source: 'usage_processing' }, extra: { usageContext }, }); - return; + return null; } const timer = createTimer(); @@ -987,7 +989,7 @@ export async function processTokenData( usageStats.cacheDiscount_mUsd = 0; } - await logMicrodollarUsage(usageStats, usageContext); + return logMicrodollarUsage(usageStats, usageContext); } function useAnthropicStyleTokenCounting(requestedModel: string, provider: ProviderId) { From ebceefb5f0959025fd2912ae39cc31c2dc929884 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 19 May 2026 12:36:01 +0200 Subject: [PATCH 07/49] chore(experiments): drop unused MicrodollarUsageContext import --- apps/web/src/lib/ai-gateway/experiments/persist.ts | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/apps/web/src/lib/ai-gateway/experiments/persist.ts b/apps/web/src/lib/ai-gateway/experiments/persist.ts index c49784d2837..fdd25a04318 100644 --- a/apps/web/src/lib/ai-gateway/experiments/persist.ts +++ b/apps/web/src/lib/ai-gateway/experiments/persist.ts @@ -2,10 +2,7 @@ import { captureException, captureMessage } from '@sentry/nextjs'; import { db } from '@/lib/drizzle'; import { model_experiment_request } from '@kilocode/db/schema'; import { putPromptIfAbsent } from '@/lib/r2/experiment-prompts'; -import type { - ExperimentPromptCapture, - MicrodollarUsageContext, -} from '@/lib/ai-gateway/processUsage.types'; +import type { ExperimentPromptCapture } from '@/lib/ai-gateway/processUsage.types'; import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types'; /** From 4cc8df64ce54ea54efb6ce03e510881393c43d8e Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 19 May 2026 12:53:11 +0200 Subject: [PATCH 08/49] docs(plan): update experimental-models-1 status for PR #3325 Marks Phases 2 + 3 as merged-pending against PR #3325 and Phase 4 as in-progress with 4a/4b/4c on the same PR. 4d (the `model_experiment_request_stats` view) is now explicitly deferred until a real report consumer needs it. Also records the concrete output of each phase as it landed. --- .plans/experimental-models-1.md | 45 ++++++++++++++++++++++++++------- 1 file changed, 36 insertions(+), 9 deletions(-) diff --git a/.plans/experimental-models-1.md b/.plans/experimental-models-1.md index d2beb3e0537..52fe9bbaf8a 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,7 +30,34 @@ **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. + +**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` — filters experimented public ids out of `kilo-auto/free` candidate sets via `filterOutExperimentedPublicIds`. Dedicated preview ids are never silently selected by auto-routing. + +**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. **Phase 5 — concrete output (in PR — see status table for branch/PR):** From 3f9519c22e05d494e611752e4fd97ca92975add5 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 19 May 2026 13:43:07 +0200 Subject: [PATCH 09/49] fix(experiments): zero billable cost for provider-funded experiment traffic Review fix #1. `skipBalanceCheck` only skipped the preflight zero-balance rejection; it did not mark the usage as free, so `processTokenData` would still record full upstream cost into `microdollar_usage.cost` and the insert CTE would increment `kilocode_users.microdollars_used`. A zero-balance user could be charged real upstream cost on a preview experiment and pushed negative. Adds an explicit `provider_funded` flag to `MicrodollarUsageContext` and treats it the same as `user_byok` in the cost-zeroing branch (zeroes `cost_mUsd` and `cacheDiscount_mUsd` while preserving `market_cost` for reporting). The route handler sets the flag for experimented requests, alongside the existing `skipBalanceCheck`. --- apps/web/src/app/api/openrouter/[...path]/route.ts | 5 +++++ apps/web/src/lib/ai-gateway/processUsage.ts | 6 +++++- apps/web/src/lib/ai-gateway/processUsage.types.ts | 7 +++++++ 3 files changed, 17 insertions(+), 1 deletion(-) diff --git a/apps/web/src/app/api/openrouter/[...path]/route.ts b/apps/web/src/app/api/openrouter/[...path]/route.ts index 40c824c0be9..9a6e114b0df 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.ts @@ -507,6 +507,11 @@ export async function POST(request: NextRequest): Promise Date: Tue, 19 May 2026 13:44:54 +0200 Subject: [PATCH 10/49] docs(plan): record stale membership cache as next-step follow-up Captures the review-#4 finding from PR #3325: no TTL on the EXPERIMENTED_PUBLIC_IDS_REDIS_KEY plus best-effort/swallowed invalidation can leave a stale list indefinitely. Documents the agreed fix (add a TTL so the cache self-heals via the existing DB fallback path) and defers it to a follow-up PR. --- .plans/experimental-models-1.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/.plans/experimental-models-1.md b/.plans/experimental-models-1.md index 52fe9bbaf8a..c4075b53dc5 100644 --- a/.plans/experimental-models-1.md +++ b/.plans/experimental-models-1.md @@ -59,6 +59,18 @@ 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`). From a354d37838b60b2cf305fa924de08c17fa986de6 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 19 May 2026 13:47:19 +0200 Subject: [PATCH 11/49] fix(experiments): refuse experiment routing when org disables data collection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review fix #2. Direct experiment upstreams bypass the gateway-pinned `data_collection` setting — that setting only applies when the request is forwarded through OpenRouter or Vercel with a pinned provider config. For experimented public ids the upstream is selected by the variant version, so the org's preference would be silently ignored while we still capture prompts to R2 for partner evaluation. When an org sets `data_collection: 'deny'` and the request resolves to an experiment, return the existing `dataCollectionRequiredResponse()` instead of routing. The existing model deny list still runs first via `checkOrganizationModelRestrictions`. Anonymous users and non-org authenticated users have no settings and continue to route normally. --- apps/web/src/app/api/openrouter/[...path]/route.ts | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/apps/web/src/app/api/openrouter/[...path]/route.ts b/apps/web/src/app/api/openrouter/[...path]/route.ts index 9a6e114b0df..c4f968a0d69 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.ts @@ -496,6 +496,15 @@ export async function POST(request: NextRequest): Promise Date: Tue, 19 May 2026 13:48:43 +0200 Subject: [PATCH 12/49] fix(experiments): cache encrypted upstream; decrypt only on pick Review fix #3. The Redis per-public-id experiment cache previously held decrypted partner-issued api keys for up to 10 minutes. That defeated the BYOK_ENCRYPTION_KEY fail-closed behavior: removing or rotating the env var did nothing for warm caches because `getRoutingExperimentForPublicId` returned the cached payload straight to the picker. The PR verification claim ('disabling BYOK_ENCRYPTION_KEY returns 503') only held on cold caches. Refactors the cache to hold CIPHERTEXT exclusively. Decryption + the missing-key check now run inside `pickModelExperimentVariant`, on the chosen variant only, on every request. Three knock-on benefits: 1. Key rotation is effective immediately on the next request; warm caches no longer serve stale plaintext keys. 2. We decrypt one variant per request instead of N at cache-build time. 3. A Redis dump no longer leaks usable api keys. Per-request cost is one AES-256-GCM decrypt (microseconds), which is negligible against an LLM round-trip. --- .../ai-gateway/experiments/pick-variant.ts | 84 ++++++++++++------- 1 file changed, 55 insertions(+), 29 deletions(-) diff --git a/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts index c1b0f9f8cbd..25968b590aa 100644 --- a/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts @@ -18,19 +18,30 @@ import { } from '@/lib/ai-gateway/experiments/upstream-schema'; import type { ResolvedExperimentUpstream } from '@/lib/ai-gateway/experiments/build-direct-provider'; -/** TTL for the per-public-id experiment cache. Doubles as the upper bound on - * api-key rotation lag (the cache holds decrypted plaintext). */ +/** TTL for the per-public-id experiment cache. The cache holds CIPHERTEXT + * (encrypted api keys + the upstream blob); decryption happens per-pick + * inside `pickModelExperimentVariant`, so this TTL only governs how stale + * the variant set / current-version-per-variant resolution can be, not + * how long plaintext keys can live in Redis. */ const PER_PUBLIC_ID_CACHE_TTL_SECONDS = 600; type ExperimentStatus = 'active' | 'paused'; -/** Per-variant payload cached in Redis. `upstream` already includes the - * decrypted `api_key` merged in for hot-path use. */ +/** + * Per-variant payload cached in Redis. + * + * The api key is stored ENCRYPTED — the cache never holds plaintext. + * Decryption happens once per request, only on the chosen variant, inside + * `pickModelExperimentVariant`. This keeps `BYOK_ENCRYPTION_KEY` rotation + * effective immediately (no 10-minute "plaintext lives in Redis" window) + * and limits the blast radius of a Redis dump. + */ type CachedVariant = { variantId: string; weight: number; variantVersionId: string; - upstream: ResolvedExperimentUpstream; + upstream: ExperimentUpstream; + encryptedApiKey: EncryptedData; }; type CachedExperiment = { @@ -140,14 +151,13 @@ function parseStringArray(raw: string): string[] | null { /** * Returns the routing-relevant experiment for `publicId` (status active or * paused) with all variants resolved to their current - * `model_experiment_variant_version` and decrypted upstream blob. Returns - * `none` only after Postgres has proven there is no routing-relevant - * experiment for this id. + * `model_experiment_variant_version`. The cached payload holds CIPHERTEXT + * for partner-issued api keys; decryption is deferred to + * `pickModelExperimentVariant`. Returns `none` only after Postgres has + * proven there is no routing-relevant experiment for this id. * - * Cached per-publicId for ~10 minutes. The cached payload contains - * decrypted partner-issued api keys, so the TTL doubles as a key-rotation - * lag bound. The cache is invalidated by every admin mutation that affects - * routing. + * Cached per-publicId for ~10 minutes. The cache is invalidated by every + * admin mutation that affects routing. */ export async function getRoutingExperimentForPublicId(publicId: string): Promise { const cacheKey = modelExperimentRedisKey(publicId); @@ -318,11 +328,11 @@ async function loadExperimentFromDb(publicId: string): Promise { } } - if (!BYOK_ENCRYPTION_KEY) { - warnMissingEncryptionKeyOnce(); - return { kind: 'unavailable' }; - } - + // We deliberately DO NOT call decryptApiKey here, even when + // BYOK_ENCRYPTION_KEY is missing. The encrypted blob is harmless on its + // own; decryption + the missing-key check happen per-pick in + // `pickModelExperimentVariant` so a warm cache cannot serve plaintext + // keys after key rotation/removal. const variants: CachedVariant[] = []; for (const v of variantRows) { const ver = versionByVariantId.get(v.id); @@ -344,21 +354,12 @@ async function loadExperimentFromDb(publicId: string): Promise { }); return { kind: 'unavailable' }; } - let apiKey: string; - try { - apiKey = decryptApiKey(ver.encryptedApiKey, BYOK_ENCRYPTION_KEY); - } catch (err) { - captureException(err, { - tags: { source: 'model-experiments', operation: 'decryptApiKey' }, - extra: { experimentId: experiment.id, variantId: v.id, variantVersionId: ver.id }, - }); - return { kind: 'unavailable' }; - } variants.push({ variantId: v.id, weight: v.weight, variantVersionId: ver.id, - upstream: { ...(parsedUpstream.data satisfies ExperimentUpstream), api_key: apiKey }, + upstream: parsedUpstream.data, + encryptedApiKey: ver.encryptedApiKey, }); } @@ -428,12 +429,37 @@ export async function pickModelExperimentVariant( 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 + // BYOK_ENCRYPTION_KEY missing-key fail-closed effective on warm + // caches, (b) lets key rotation invalidate access immediately + // (next request after the key flips fails closed), and (c) means + // we never decrypt N-1 keys we won't use. + if (!BYOK_ENCRYPTION_KEY) { + warnMissingEncryptionKeyOnce(); + return { status: 'unavailable' }; + } + 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, + upstream: { ...v.upstream, api_key: apiKey }, allocationSubject: subjectPick.subject, }; } From 67573606c0dcb9bf78ac2103c28f6acfcaf564c3 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 19 May 2026 13:49:33 +0200 Subject: [PATCH 13/49] fix(experiments): cap prompt capture by UTF-8 bytes, not UTF-16 units MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review fix #6. SYSTEM_PROMPT_CAP_BYTES / REQUEST_BODY_CAP_BYTES were checked with `string.length` and trimmed with `string.slice`, which both operate on UTF-16 code units. For non-ASCII content the actual UTF-8 byte payload uploaded to R2 could be up to ~3× the intended cap (a single CJK char is 1 UTF-16 unit but 3 UTF-8 bytes; an emoji is 2 UTF-16 units but 4 UTF-8 bytes), so the cap was not bounding what the plan claimed it bounded. Slicing UTF-16 code units could also split a surrogate pair, leading the next UTF-8 encode to emit U+FFFD and breaking the deterministic-prefix promise (truncated content would hash differently across runtimes). Adds `truncateToUtf8Bytes(s, maxBytes)` which encodes once, walks back to a UTF-8 codepoint boundary using the 0b10xxxxxx continuation-byte rule, and returns valid UTF-8 with bytes <= maxBytes. Replaces the three call sites in `buildExperimentPromptCapture` to measure with `Buffer.byteLength` and truncate via the new helper. --- .../src/lib/ai-gateway/experiments/persist.ts | 39 +++++++++++++++---- 1 file changed, 32 insertions(+), 7 deletions(-) diff --git a/apps/web/src/lib/ai-gateway/experiments/persist.ts b/apps/web/src/lib/ai-gateway/experiments/persist.ts index fdd25a04318..51901ac8cd3 100644 --- a/apps/web/src/lib/ai-gateway/experiments/persist.ts +++ b/apps/web/src/lib/ai-gateway/experiments/persist.ts @@ -14,10 +14,32 @@ export const PROMPT_HASH_ABSENT = '__absent__'; export const PROMPT_HASH_FAILED = '__failed__'; export const PROMPT_HASH_DELETED = '__deleted__'; -/** 4 MB. Comfortably above any current 1M-token-context request. */ +/** 4 MB measured as UTF-8 bytes (matching what `putPromptIfAbsent` actually + * uploads). Comfortably above any current 1M-token-context request. */ export const SYSTEM_PROMPT_CAP_BYTES = 4 * 1024 * 1024; 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, split into the leading system message and the remainder. @@ -46,12 +68,15 @@ export function buildExperimentPromptCapture(request: GatewayRequest): Experimen } const remainder = { ...request.body, messages: nonSystemMessages }; let bodyContent = serialize(remainder); - if (systemContent !== null && systemContent.length > SYSTEM_PROMPT_CAP_BYTES) { - systemContent = systemContent.slice(0, SYSTEM_PROMPT_CAP_BYTES); + if ( + systemContent !== null && + Buffer.byteLength(systemContent, 'utf8') > SYSTEM_PROMPT_CAP_BYTES + ) { + systemContent = truncateToUtf8Bytes(systemContent, SYSTEM_PROMPT_CAP_BYTES); wasTruncated = true; } - if (bodyContent.length > REQUEST_BODY_CAP_BYTES) { - bodyContent = bodyContent.slice(0, REQUEST_BODY_CAP_BYTES); + if (Buffer.byteLength(bodyContent, 'utf8') > REQUEST_BODY_CAP_BYTES) { + bodyContent = truncateToUtf8Bytes(bodyContent, REQUEST_BODY_CAP_BYTES); wasTruncated = true; } if (wasTruncated) { @@ -68,8 +93,8 @@ export function buildExperimentPromptCapture(request: GatewayRequest): Experimen // side; system extraction for those API kinds is left for later if it // turns out to dedup meaningfully. let bodyContent = serialize(request.body); - if (bodyContent.length > REQUEST_BODY_CAP_BYTES) { - bodyContent = bodyContent.slice(0, REQUEST_BODY_CAP_BYTES); + if (Buffer.byteLength(bodyContent, 'utf8') > REQUEST_BODY_CAP_BYTES) { + bodyContent = truncateToUtf8Bytes(bodyContent, REQUEST_BODY_CAP_BYTES); wasTruncated = true; noteTruncation(request.body.model); } From 986fde979df1967ce27a882d75790b8afdd4abe6 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 19 May 2026 13:52:25 +0200 Subject: [PATCH 14/49] test(experiments): cover picker, fail-closed paths, and UTF-8 truncation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds two test files covering the new experiment routing surface: - `persist.test.ts` (14 cases) — `truncateToUtf8Bytes` boundary cases (multi-byte CJK, surrogate-pair emoji, exact alignment, zero cap, determinism); `buildExperimentPromptCapture` extraction (leading vs mid-conversation system message, serialization of structured content, UTF-8-byte cap enforcement for both system and body sides). - `pick-variant.test.ts` (13 cases) — `isPublicIdExperimented` positive/negative + paused-also-counted; picker determinism for the same userId; api-key decryption returns the right partner secret per variant; allocation-subject precedence user > machine > ip; `unavailable` when no allocation subject; `not-found` for paused so traffic doesn't silently fall through; hot-swap preserves the variant_id slot while serving the new variant_version_id; weighted-distribution sanity (1:3 split lands ~25%); historical attribution via the old version row survives a hot-swap; completed experiments stop routing. Full suite: 373 passed (up from 371). --- .../ai-gateway/experiments/persist.test.ts | 177 ++++++++++ .../experiments/pick-variant.test.ts | 325 ++++++++++++++++++ 2 files changed, 502 insertions(+) create mode 100644 apps/web/src/lib/ai-gateway/experiments/persist.test.ts create mode 100644 apps/web/src/lib/ai-gateway/experiments/pick-variant.test.ts 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..456a516548e --- /dev/null +++ b/apps/web/src/lib/ai-gateway/experiments/persist.test.ts @@ -0,0 +1,177 @@ +import { describe, expect, it } from '@jest/globals'; +import { + buildExperimentPromptCapture, + REQUEST_BODY_CAP_BYTES, + SYSTEM_PROMPT_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 { + return { kind: 'chat_completions', body } as GatewayRequest; + } + + it('extracts a leading system message into systemPromptContent', () => { + const cap = buildExperimentPromptCapture( + chatRequest({ + model: 'kilo/preview-foo', + messages: [ + { role: 'system', content: 'you are a helpful assistant' }, + { role: 'user', content: 'hi' }, + ], + }) + ); + expect(cap.systemPromptContent).toBe('you are a helpful assistant'); + expect(cap.requestBodyContent).toContain('"role":"user"'); + expect(cap.requestBodyContent).not.toContain('"role":"system"'); + expect(cap.wasTruncated).toBe(false); + }); + + it('leaves systemPromptContent null when there is no leading system message', () => { + const cap = buildExperimentPromptCapture( + chatRequest({ + model: 'kilo/preview-foo', + messages: [{ role: 'user', content: 'hi' }], + }) + ); + expect(cap.systemPromptContent).toBeNull(); + expect(cap.requestBodyContent).toContain('"role":"user"'); + }); + + it('does not extract a system message if it is not the first', () => { + // Mid-conversation system messages are kept in the body. + const cap = buildExperimentPromptCapture( + chatRequest({ + model: 'kilo/preview-foo', + messages: [ + { role: 'user', content: 'hi' }, + { role: 'system', content: 'note' }, + ], + }) + ); + expect(cap.systemPromptContent).toBeNull(); + expect(cap.requestBodyContent).toContain('"role":"system"'); + }); + + it('truncates the system prompt by UTF-8 bytes and marks wasTruncated', () => { + const huge = 'a'.repeat(SYSTEM_PROMPT_CAP_BYTES + 1024); + const cap = buildExperimentPromptCapture( + chatRequest({ + model: 'kilo/preview-foo', + messages: [ + { role: 'system', content: huge }, + { role: 'user', content: 'hi' }, + ], + }) + ); + expect(cap.wasTruncated).toBe(true); + expect(cap.systemPromptContent).not.toBeNull(); + expect(Buffer.byteLength(cap.systemPromptContent ?? '', 'utf8')).toBeLessThanOrEqual( + SYSTEM_PROMPT_CAP_BYTES + ); + }); + + it('truncates the request body remainder by UTF-8 bytes', () => { + 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('counts UTF-8 bytes (not UTF-16 units) when checking the cap', () => { + // Build a string whose UTF-16 length is well under the cap but whose + // UTF-8 byte length exceeds it. CJK chars are 1 UTF-16 unit but 3 + // UTF-8 bytes; with cap=10 bytes, 5 CJK chars (5 units, 15 bytes) + // must trigger truncation. + const cap = buildExperimentPromptCapture({ + kind: 'chat_completions', + body: { + model: 'kilo/preview-foo', + messages: [ + { role: 'system', content: '日'.repeat(5) }, + { role: 'user', content: 'x' }, + ], + }, + } as GatewayRequest); + // We can't trigger this against the real 4 MB cap in a unit test, but + // we can assert the UTF-8 byte count of the captured content directly. + expect(cap.systemPromptContent).toBe('日'.repeat(5)); + expect(Buffer.byteLength(cap.systemPromptContent ?? '', 'utf8')).toBe(15); + expect(cap.systemPromptContent?.length).toBe(5); // UTF-16 unit count for sanity. + }); + + it('serializes non-string system content', () => { + const cap = buildExperimentPromptCapture( + chatRequest({ + model: 'kilo/preview-foo', + messages: [ + { + role: 'system', + content: [{ type: 'text', text: 'hello' }], + }, + { role: 'user', content: 'hi' }, + ], + }) + ); + expect(cap.systemPromptContent).toContain('"text":"hello"'); + }); +}); 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..4e3eaa2ab2f --- /dev/null +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.test.ts @@ -0,0 +1,325 @@ +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, pickModelExperimentVariant } from './pick-variant'; +import { redisDel } from '@/lib/redis'; +import { EXPERIMENTED_PUBLIC_IDS_REDIS_KEY, modelExperimentRedisKey } 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', +}; + +beforeEach(async () => { + await cleanupDbForTest(); + admin = await insertTestUser({ + google_user_email: `admin-${Math.random()}@admin.example.com`, + is_admin: true, + }); +}); + +async function clearRoutingCaches(publicId: string) { + // Tests share the dev Redis instance across runs; flush both keys so + // each test sees a fresh load path. + await redisDel(EXPERIMENTED_PUBLIC_IDS_REDIS_KEY); + await redisDel(modelExperimentRedisKey(publicId)); +} + +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 }); + await clearRoutingCaches(publicId); + return { experimentId: exp.id, variantA: a.id, variantB: b.id }; +} + +describe('isPublicIdExperimented', () => { + it('returns false for an unknown public id', async () => { + await clearRoutingCaches('kilo/preview-not-experimented'); + expect(await isPublicIdExperimented('kilo/preview-not-experimented')).toBe(false); + }); + + it('returns true when the public id has an active experiment', async () => { + await makeActiveExperiment({ publicId: 'kilo/preview-iset-active' }); + expect(await isPublicIdExperimented('kilo/preview-iset-active')).toBe(true); + }); + + it('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 }); + await clearRoutingCaches('kilo/preview-iset-paused'); + 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('kilo/preview-pick-none'); + 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('kilo/preview-paused'); + 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('kilo/preview-hotswap'); + + 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).toBeLessThan(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('kilo/preview-completed'); + 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'); + }); +}); From 33b3baf657fcc351c0510d5c422713f441e367a6 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 19 May 2026 14:06:37 +0200 Subject: [PATCH 15/49] fix(experiments): honor request-level data-collection opt-out Follow-up to review fix #2. The previous commit only refused experiment routing on the org-level setting (`settings.data_collection === 'deny'`). A caller can also opt a single request out of training/data-retention via OpenRouter's request-body fields `provider.data_collection: 'deny'` or `provider.zdr: true`. Direct experiment upstreams ignore those fields entirely (we don't go through OpenRouter), but the prompt-capture path was still uploading to R2 for partner evaluation, violating the caller's stated intent. Refuses experiment routing with `dataCollectionRequiredResponse()` when `!isFreePromptTrainingAllowed(body.provider)` regardless of org/anon/BYOK status. The check runs unconditionally right after `getProvider` returns so it covers anonymous and non-org authenticated users too. Adds a unit test for `isFreePromptTrainingAllowed` covering all opt-out shapes (`data_collection: 'deny'`, `zdr: true`, `zdr: false`, both, neither, undefined) so a regression cannot silently allow capture for users who explicitly opted out at the request scope. Also tightens the `GatewayRequest` cast in `persist.test.ts` per the test-file as-cast convention. --- .../src/app/api/openrouter/[...path]/route.ts | 13 +++++++ .../ai-gateway/experiments/persist.test.ts | 7 +++- .../providers/openrouter/types.test.ts | 37 +++++++++++++++++++ 3 files changed, 55 insertions(+), 2 deletions(-) create mode 100644 apps/web/src/lib/ai-gateway/providers/openrouter/types.test.ts diff --git a/apps/web/src/app/api/openrouter/[...path]/route.ts b/apps/web/src/app/api/openrouter/[...path]/route.ts index c4f968a0d69..6072c00f5a3 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.ts @@ -407,6 +407,19 @@ export async function POST(request: NextRequest): Promise { describe('buildExperimentPromptCapture', () => { function chatRequest(body: Record): GatewayRequest { - return { kind: 'chat_completions', body } as 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('extracts a leading system message into systemPromptContent', () => { @@ -151,7 +154,7 @@ describe('buildExperimentPromptCapture', () => { { role: 'user', content: 'x' }, ], }, - } as GatewayRequest); + } as unknown as GatewayRequest); // We can't trigger this against the real 4 MB cap in a unit test, but // we can assert the UTF-8 byte count of the captured content directly. expect(cap.systemPromptContent).toBe('日'.repeat(5)); diff --git a/apps/web/src/lib/ai-gateway/providers/openrouter/types.test.ts b/apps/web/src/lib/ai-gateway/providers/openrouter/types.test.ts new file mode 100644 index 00000000000..ee7825b4234 --- /dev/null +++ b/apps/web/src/lib/ai-gateway/providers/openrouter/types.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, it } from '@jest/globals'; +import { isFreePromptTrainingAllowed } from './types'; + +// Used by both the kilo-exclusive free-model data-collection gate and the +// model-experiment routing refusal. Locks the per-request opt-out shapes +// (`provider.data_collection: 'deny'` and `provider.zdr: true`) so a +// regression in this predicate cannot silently allow prompt capture for +// callers who explicitly opted out at the request scope. +describe('isFreePromptTrainingAllowed', () => { + it('allows when the provider config is undefined', () => { + expect(isFreePromptTrainingAllowed(undefined)).toBe(true); + }); + + it('allows when no opt-out fields are set', () => { + expect(isFreePromptTrainingAllowed({})).toBe(true); + }); + + it('allows when data_collection is explicitly "allow"', () => { + expect(isFreePromptTrainingAllowed({ data_collection: 'allow' })).toBe(true); + }); + + it('refuses when data_collection is "deny"', () => { + expect(isFreePromptTrainingAllowed({ data_collection: 'deny' })).toBe(false); + }); + + it('refuses when zdr is true', () => { + expect(isFreePromptTrainingAllowed({ zdr: true })).toBe(false); + }); + + it('refuses when both opt-out fields are set', () => { + expect(isFreePromptTrainingAllowed({ data_collection: 'deny', zdr: true })).toBe(false); + }); + + it('allows when zdr is explicitly false and no other opt-out is set', () => { + expect(isFreePromptTrainingAllowed({ zdr: false })).toBe(true); + }); +}); From b6b1ebc3ce9c297675db78dcc5be711bba6d5460 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 19 May 2026 14:27:55 +0200 Subject: [PATCH 16/49] chore(experiments): drop plan references from code comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Plans are ephemeral and get removed when the work is done. Comments that pointed readers at `.plans/experimental-models-1.md` or referenced specific plan sections now stand on their own — the relevant context (sentinel meanings, why the api key column is separate, scope of preview routing) is inlined where it matters instead of forwarded to a doc that won't exist long-term. Touched: - apps/web/src/lib/r2/experiment-prompts.ts (header) - apps/web/src/lib/ai-gateway/experiments/persist.ts (sentinel docs) - apps/web/src/lib/ai-gateway/experiments/upstream-schema.ts (api_key) - apps/web/src/app/admin/model-experiments/ModelExperimentsContent.tsx (admin UI subtitle: dropped trailing 'See .plans/...' sentence) - packages/db/src/schema.ts (three model_experiment* table headers) --- .../ModelExperimentsContent.tsx | 1 - .../src/lib/ai-gateway/experiments/persist.ts | 10 +++++++-- .../ai-gateway/experiments/upstream-schema.ts | 4 ++-- apps/web/src/lib/r2/experiment-prompts.ts | 8 +++---- packages/db/src/schema.ts | 22 +++++++++++-------- 5 files changed, 27 insertions(+), 18 deletions(-) diff --git a/apps/web/src/app/admin/model-experiments/ModelExperimentsContent.tsx b/apps/web/src/app/admin/model-experiments/ModelExperimentsContent.tsx index 0ca0b7ca0b6..24da8948898 100644 --- a/apps/web/src/app/admin/model-experiments/ModelExperimentsContent.tsx +++ b/apps/web/src/app/admin/model-experiments/ModelExperimentsContent.tsx @@ -165,7 +165,6 @@ function ExperimentList({ onSelect }: { onSelect: (id: string) => void }) { public_model_id {' '} and swaps the upstream checkpoint behind it. Clients keep sending the same public model id. - See .plans/experimental-models-1.md for the full design.

{isLoading ? ( diff --git a/apps/web/src/lib/ai-gateway/experiments/persist.ts b/apps/web/src/lib/ai-gateway/experiments/persist.ts index 51901ac8cd3..d22a91c83ba 100644 --- a/apps/web/src/lib/ai-gateway/experiments/persist.ts +++ b/apps/web/src/lib/ai-gateway/experiments/persist.ts @@ -7,8 +7,14 @@ import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types /** * Sentinel values written to `model_experiment_request.system_prompt_sha256` / - * `request_body_sha256` instead of a real 64-char hex digest. See the schema - * CHECK constraints and `.plans/experimental-models-1.md` "Prompt Storage". + * `request_body_sha256` instead of a real 64-char hex digest. The schema + * CHECK constraints accept either a 64-char lowercase hex digest or one + * of these sentinels. + * + * - `__absent__` (system only): the request had no leading system message. + * - `__failed__`: R2 storage failed; the attribution row still lands. + * - `__deleted__`: prompt content intentionally wiped while retaining + * experiment attribution. */ export const PROMPT_HASH_ABSENT = '__absent__'; export const PROMPT_HASH_FAILED = '__failed__'; diff --git a/apps/web/src/lib/ai-gateway/experiments/upstream-schema.ts b/apps/web/src/lib/ai-gateway/experiments/upstream-schema.ts index 5ec5902c087..b7eb958e7c4 100644 --- a/apps/web/src/lib/ai-gateway/experiments/upstream-schema.ts +++ b/apps/web/src/lib/ai-gateway/experiments/upstream-schema.ts @@ -10,8 +10,8 @@ import { * * - **Does not include `api_key`.** The encrypted key lives in the sibling * `model_experiment_variant_version.encrypted_api_key` JSONB column so the - * never-read invariant is enforceable at the column level. See - * `.plans/experimental-models-1.md` §"API Keys". + * never-read invariant is enforceable at the column level — admin tRPC + * responses simply omit the column from their selects. * - **Does not include `extra_headers`.** Partner checkpoint routing should * use `api_key` + `base_url` + `internal_id` + adapter settings. If a * non-secret header is required later, add an explicit allowlisted field diff --git a/apps/web/src/lib/r2/experiment-prompts.ts b/apps/web/src/lib/r2/experiment-prompts.ts index ced4592e5ce..6f7f3805a11 100644 --- a/apps/web/src/lib/r2/experiment-prompts.ts +++ b/apps/web/src/lib/r2/experiment-prompts.ts @@ -13,10 +13,10 @@ import { r2Client, r2ExperimentPromptsBucketName } from './client'; * * - Bucket: `R2_EXPERIMENT_PROMPTS_BUCKET_NAME` (one per environment). * - Trust boundary: same R2 credentials as cli-sessions and - * cloud-agent-attachments; prompts collected under explicit experiment - * opt-in are not subject to the default microdollar_usage_metadata - * soft-delete policy. See `.plans/experimental-models-1.md` "Prompt - * Storage > GDPR and consent". + * cloud-agent-attachments. Prompts collected under explicit experiment + * opt-in use a dedicated retention policy and are NOT subject to the + * default `microdollar_usage_metadata` soft-delete; `softDeleteUser` + * does not touch these rows or their R2 objects. * * Failures here MUST NOT roll back the microdollar usage write or the * experiment attribution row. Callers store `__failed__` for the affected diff --git a/packages/db/src/schema.ts b/packages/db/src/schema.ts index 484d5b1f96e..5c323507e9e 100644 --- a/packages/db/src/schema.ts +++ b/packages/db/src/schema.ts @@ -6084,9 +6084,11 @@ export type NewSecurityAdvisorScan = typeof security_advisor_scans.$inferInsert; // --------------------------------------------------------------------------- // Model experiments (preview/experimental A/B testing) -// See `.plans/experimental-models-1.md` and (eventually) `.specs/model-experiments.md`. +// // Scope: opt-in dedicated preview public model ids only. Never used for -// production/general traffic. See plan for full design rationale. +// production/general traffic. Users only reach this routing path by +// explicitly selecting a dedicated preview public id (e.g. +// `kilo/preview-experiment-foo`). // --------------------------------------------------------------------------- export const model_experiment = pgTable( @@ -6155,11 +6157,10 @@ export type ModelExperimentVariant = typeof model_experiment_variant.$inferSelec export type NewModelExperimentVariant = typeof model_experiment_variant.$inferInsert; // Immutable per-variant version. New RC = new row. Never UPDATEd. -// `upstream` is validated by ExperimentUpstreamSchema in app code (see -// experimental-models-1.md Phase 1). The api key is stored separately in -// `encrypted_api_key` (same shape as byok_api_keys.encrypted_api_key) so the -// JSONB blob never holds the secret and reporting/admin views can simply omit -// the column. +// `upstream` is validated by ExperimentUpstreamSchema in app code. The +// api key is stored separately in `encrypted_api_key` (same shape as +// `byok_api_keys.encrypted_api_key`) so the JSONB blob never holds the +// secret and reporting/admin views can simply omit the column. export const model_experiment_variant_version = pgTable( 'model_experiment_variant_version', { @@ -6185,8 +6186,11 @@ export type ModelExperimentVariantVersion = typeof model_experiment_variant_vers export type NewModelExperimentVariantVersion = typeof model_experiment_variant_version.$inferInsert; // One row per experimented request, keyed on usage_id (1:1 with microdollar_usage). -// Stores attribution + R2 prompt hashes (or reserved sentinel values). -// See experimental-models-1.md "Prompt Storage (R2)" for sentinel values. +// Stores attribution + R2 prompt hashes. The two sha256 columns hold either +// a 64-char lowercase hex digest pointing at an R2 object, or one of the +// reserved sentinels: `__absent__` (system only, no leading system message), +// `__failed__` (R2 storage failed), `__deleted__` (prompt content wiped +// while retaining attribution). export const model_experiment_request = pgTable( 'model_experiment_request', { From 810c76ceb30a48c57b5e8cfb5ed0a8a814b8f79f Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 19 May 2026 14:35:10 +0200 Subject: [PATCH 17/49] chore(experiments): leave shared-with-main files for separate chore PR The previous commit (b6b1ebc3c) cleaned plan references from five files, three of which already exist on main from earlier merged PRs (Phase 1 schema, Phase 5 admin UI). Touching cross-package files unrelated to the gateway work blew up this PR's diff and risked merge churn against PR #3329, which carries those same three edits independently. Restore the on-main copies of: - packages/db/src/schema.ts - apps/web/src/lib/ai-gateway/experiments/upstream-schema.ts - apps/web/src/app/admin/model-experiments/ModelExperimentsContent.tsx The branch-only files (persist.ts, experiment-prompts.ts) keep their cleanups since they were authored on this branch and land/ship together with the gateway work. --- .../ModelExperimentsContent.tsx | 1 + .../ai-gateway/experiments/upstream-schema.ts | 4 ++-- packages/db/src/schema.ts | 22 ++++++++----------- 3 files changed, 12 insertions(+), 15 deletions(-) diff --git a/apps/web/src/app/admin/model-experiments/ModelExperimentsContent.tsx b/apps/web/src/app/admin/model-experiments/ModelExperimentsContent.tsx index 24da8948898..0ca0b7ca0b6 100644 --- a/apps/web/src/app/admin/model-experiments/ModelExperimentsContent.tsx +++ b/apps/web/src/app/admin/model-experiments/ModelExperimentsContent.tsx @@ -165,6 +165,7 @@ function ExperimentList({ onSelect }: { onSelect: (id: string) => void }) { public_model_id {' '} and swaps the upstream checkpoint behind it. Clients keep sending the same public model id. + See .plans/experimental-models-1.md for the full design.

{isLoading ? ( diff --git a/apps/web/src/lib/ai-gateway/experiments/upstream-schema.ts b/apps/web/src/lib/ai-gateway/experiments/upstream-schema.ts index b7eb958e7c4..5ec5902c087 100644 --- a/apps/web/src/lib/ai-gateway/experiments/upstream-schema.ts +++ b/apps/web/src/lib/ai-gateway/experiments/upstream-schema.ts @@ -10,8 +10,8 @@ import { * * - **Does not include `api_key`.** The encrypted key lives in the sibling * `model_experiment_variant_version.encrypted_api_key` JSONB column so the - * never-read invariant is enforceable at the column level — admin tRPC - * responses simply omit the column from their selects. + * never-read invariant is enforceable at the column level. See + * `.plans/experimental-models-1.md` §"API Keys". * - **Does not include `extra_headers`.** Partner checkpoint routing should * use `api_key` + `base_url` + `internal_id` + adapter settings. If a * non-secret header is required later, add an explicit allowlisted field diff --git a/packages/db/src/schema.ts b/packages/db/src/schema.ts index 5c323507e9e..484d5b1f96e 100644 --- a/packages/db/src/schema.ts +++ b/packages/db/src/schema.ts @@ -6084,11 +6084,9 @@ export type NewSecurityAdvisorScan = typeof security_advisor_scans.$inferInsert; // --------------------------------------------------------------------------- // Model experiments (preview/experimental A/B testing) -// +// See `.plans/experimental-models-1.md` and (eventually) `.specs/model-experiments.md`. // Scope: opt-in dedicated preview public model ids only. Never used for -// production/general traffic. Users only reach this routing path by -// explicitly selecting a dedicated preview public id (e.g. -// `kilo/preview-experiment-foo`). +// production/general traffic. See plan for full design rationale. // --------------------------------------------------------------------------- export const model_experiment = pgTable( @@ -6157,10 +6155,11 @@ export type ModelExperimentVariant = typeof model_experiment_variant.$inferSelec export type NewModelExperimentVariant = typeof model_experiment_variant.$inferInsert; // Immutable per-variant version. New RC = new row. Never UPDATEd. -// `upstream` is validated by ExperimentUpstreamSchema in app code. The -// api key is stored separately in `encrypted_api_key` (same shape as -// `byok_api_keys.encrypted_api_key`) so the JSONB blob never holds the -// secret and reporting/admin views can simply omit the column. +// `upstream` is validated by ExperimentUpstreamSchema in app code (see +// experimental-models-1.md Phase 1). The api key is stored separately in +// `encrypted_api_key` (same shape as byok_api_keys.encrypted_api_key) so the +// JSONB blob never holds the secret and reporting/admin views can simply omit +// the column. export const model_experiment_variant_version = pgTable( 'model_experiment_variant_version', { @@ -6186,11 +6185,8 @@ export type ModelExperimentVariantVersion = typeof model_experiment_variant_vers export type NewModelExperimentVariantVersion = typeof model_experiment_variant_version.$inferInsert; // One row per experimented request, keyed on usage_id (1:1 with microdollar_usage). -// Stores attribution + R2 prompt hashes. The two sha256 columns hold either -// a 64-char lowercase hex digest pointing at an R2 object, or one of the -// reserved sentinels: `__absent__` (system only, no leading system message), -// `__failed__` (R2 storage failed), `__deleted__` (prompt content wiped -// while retaining attribution). +// Stores attribution + R2 prompt hashes (or reserved sentinel values). +// See experimental-models-1.md "Prompt Storage (R2)" for sentinel values. export const model_experiment_request = pgTable( 'model_experiment_request', { From da9799e6a3c4a1fc72cc73aadadc6fd5f2c16e5d Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 19 May 2026 15:39:16 +0200 Subject: [PATCH 18/49] revert(gateway): drop kilo-auto/free experimented-id safeguard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reverts the `filterOutExperimentedPublicIds` filter added in `72dd51171` for `kilo-auto/free` candidate selection. The safeguard added a per-candidate `isPublicIdExperimented` Redis check on every `kilo-auto/free` request — material extra latency on a hot path — to defend against an invariant that holds by construction: `autoFreeModels` and the frontier preset list are hand-curated and never include dedicated preview experiment ids. The explicit-opt-in property of preview models is preserved without any runtime guard. `getProvider`'s experiment-membership check is unchanged — explicit requests to a preview public id still resolve through the experiment routing path. Plan updated to reflect the no-guard decision and rationale. --- .plans/experimental-models-1.md | 6 +++--- .../lib/ai-gateway/auto-model/resolution.ts | 19 +------------------ 2 files changed, 4 insertions(+), 21 deletions(-) diff --git a/.plans/experimental-models-1.md b/.plans/experimental-models-1.md index c4075b53dc5..80369e83be8 100644 --- a/.plans/experimental-models-1.md +++ b/.plans/experimental-models-1.md @@ -45,7 +45,7 @@ - `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` — filters experimented public ids out of `kilo-auto/free` candidate sets via `filterOutExperimentedPublicIds`. Dedicated preview ids are never silently selected by auto-routing. +- `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):** @@ -374,7 +374,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' }`. @@ -455,7 +455,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. 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 debf3cc3f66..87a4a408b1f 100644 --- a/apps/web/src/lib/ai-gateway/auto-model/resolution.ts +++ b/apps/web/src/lib/ai-gateway/auto-model/resolution.ts @@ -33,7 +33,6 @@ import { import { getOpenRouterModels } from '@/lib/ai-gateway/providers/gateway-models-cache'; import PROVIDERS from '@/lib/ai-gateway/providers/provider-definitions'; import type { ProviderId } from '@/lib/ai-gateway/providers/types'; -import { isPublicIdExperimented } from '@/lib/ai-gateway/experiments/pick-variant'; type ResolveAutoModelParams = { model: string; @@ -77,20 +76,6 @@ export function getAutoFreeCandidates( return [...candidates].toSorted(); } -/** - * Removes any public ids that are currently routed by an active or paused - * model experiment. Dedicated preview public ids must never be selected by - * `kilo-auto` — users only ever opt in to them explicitly. - */ -async function filterOutExperimentedPublicIds( - candidates: ReadonlyArray -): Promise> { - const checks = await Promise.all( - candidates.map(async id => ({ id, experimented: await isPublicIdExperimented(id) })) - ); - return checks.filter(c => c.experimented !== true).map(c => c.id); -} - function gatewaySupportsApiKind( gateway: ProviderId, apiKind: GatewayRequest['kind'] | null @@ -112,9 +97,7 @@ export async function resolveAutoModel( const { model, modeHeader, featureHeader, sessionId, apiKind, clientIp } = params; if (model === KILO_AUTO_FREE_MODEL.id) { const openRouterModels = await getOpenRouterModels(); - const candidates = await filterOutExperimentedPublicIds( - getAutoFreeCandidates(openRouterModels, apiKind) - ); + const candidates = getAutoFreeCandidates(openRouterModels, apiKind); if (candidates.length === 0) { return { kind: 'no_free_models_available' }; } From 0afa45328e56eda08320e1cfad4436d3646018b8 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 19 May 2026 16:57:46 +0200 Subject: [PATCH 19/49] chore(experiments): drop unused barrel index.ts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every consumer in the repo deep-imports the specific file (`./pick-variant`, `./persist`, `./upstream-schema`, `./build-direct-provider`); nothing imported from the barrel. Most `index.ts` files under `ai-gateway/` are real module implementations, not pure re-exports — drop this one to match. --- .../src/lib/ai-gateway/experiments/index.ts | 24 ------------------- 1 file changed, 24 deletions(-) delete mode 100644 apps/web/src/lib/ai-gateway/experiments/index.ts diff --git a/apps/web/src/lib/ai-gateway/experiments/index.ts b/apps/web/src/lib/ai-gateway/experiments/index.ts deleted file mode 100644 index d75d1aced94..00000000000 --- a/apps/web/src/lib/ai-gateway/experiments/index.ts +++ /dev/null @@ -1,24 +0,0 @@ -export { - buildDirectProvider, - inferSupportedChatApis, - type ResolvedExperimentUpstream, - type DirectProviderInput, -} from './build-direct-provider'; -export { - isPublicIdExperimented, - getRoutingExperimentForPublicId, - pickModelExperimentVariant, - type AllocationSubject, - type PickVariantInput, - type PickVariantResult, -} from './pick-variant'; -export { ExperimentUpstreamSchema, type ExperimentUpstream } from './upstream-schema'; -export { - buildExperimentPromptCapture, - persistExperimentAttribution, - PROMPT_HASH_ABSENT, - PROMPT_HASH_FAILED, - PROMPT_HASH_DELETED, - SYSTEM_PROMPT_CAP_BYTES, - REQUEST_BODY_CAP_BYTES, -} from './persist'; From 88de7c08fa3a91a6be75caa384861af4b02cd505 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Wed, 20 May 2026 11:01:55 +0200 Subject: [PATCH 20/49] refactor(experiments): drop system-prompt split, plumb request_kind through capture Adapts the gateway path to the simpler model_experiment_request schema landed in #3358. Addresses comments #2 + #3 on #3325. - ExperimentPromptCapture: remove `systemPromptContent`, add `requestKind`. The full post-`transformRequest` body is now stored as a single content-addressed blob; the upstream API shape is recorded alongside it for downstream reporting. - buildExperimentPromptCapture: collapse the chat_completions-only system-extraction branch; one body capture path for all kinds. - persistExperimentAttribution: one R2 PUT instead of two, one hash written, no `__absent__` sentinel. `PROMPT_HASH_ABSENT` removed; `PROMPT_HASH_FAILED` and `PROMPT_HASH_DELETED` retained. - Tests: drop the system-extraction cases (no longer applicable), add coverage for the three request kinds and a determinism check (content addressing relies on byte-stable serialization). The dedup-on-shared-system-prompt micro-optimization had no measured benefit. If we ever want it back, content-addressing the full body is the better starting point anyway. --- .../ai-gateway/experiments/persist.test.ts | 110 ++++---------- .../src/lib/ai-gateway/experiments/persist.ts | 143 +++++------------- .../src/lib/ai-gateway/processUsage.types.ts | 12 +- 3 files changed, 77 insertions(+), 188 deletions(-) diff --git a/apps/web/src/lib/ai-gateway/experiments/persist.test.ts b/apps/web/src/lib/ai-gateway/experiments/persist.test.ts index fa8ef2d2c60..5f2335b2233 100644 --- a/apps/web/src/lib/ai-gateway/experiments/persist.test.ts +++ b/apps/web/src/lib/ai-gateway/experiments/persist.test.ts @@ -2,7 +2,6 @@ import { describe, expect, it } from '@jest/globals'; import { buildExperimentPromptCapture, REQUEST_BODY_CAP_BYTES, - SYSTEM_PROMPT_CAP_BYTES, truncateToUtf8Bytes, } from './persist'; import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types'; @@ -66,7 +65,7 @@ describe('buildExperimentPromptCapture', () => { return { kind: 'chat_completions', body } as unknown as GatewayRequest; } - it('extracts a leading system message into systemPromptContent', () => { + it('captures the full serialized body and records the request kind', () => { const cap = buildExperimentPromptCapture( chatRequest({ model: 'kilo/preview-foo', @@ -76,57 +75,35 @@ describe('buildExperimentPromptCapture', () => { ], }) ); - expect(cap.systemPromptContent).toBe('you are a helpful assistant'); + expect(cap.requestKind).toBe('chat_completions'); + expect(cap.requestBodyContent).toContain('"role":"system"'); expect(cap.requestBodyContent).toContain('"role":"user"'); - expect(cap.requestBodyContent).not.toContain('"role":"system"'); expect(cap.wasTruncated).toBe(false); }); - it('leaves systemPromptContent null when there is no leading system message', () => { - const cap = buildExperimentPromptCapture( - chatRequest({ + 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' }], - }) - ); - expect(cap.systemPromptContent).toBeNull(); - expect(cap.requestBodyContent).toContain('"role":"user"'); - }); - - it('does not extract a system message if it is not the first', () => { - // Mid-conversation system messages are kept in the body. - const cap = buildExperimentPromptCapture( - chatRequest({ - model: 'kilo/preview-foo', - messages: [ - { role: 'user', content: 'hi' }, - { role: 'system', content: 'note' }, - ], - }) - ); - expect(cap.systemPromptContent).toBeNull(); - expect(cap.requestBodyContent).toContain('"role":"system"'); + }, + } as unknown as GatewayRequest); + expect(cap.requestKind).toBe('messages'); + expect(cap.requestBodyContent).toContain('"system":"you are a helpful assistant"'); }); - it('truncates the system prompt by UTF-8 bytes and marks wasTruncated', () => { - const huge = 'a'.repeat(SYSTEM_PROMPT_CAP_BYTES + 1024); - const cap = buildExperimentPromptCapture( - chatRequest({ - model: 'kilo/preview-foo', - messages: [ - { role: 'system', content: huge }, - { role: 'user', content: 'hi' }, - ], - }) - ); - expect(cap.wasTruncated).toBe(true); - expect(cap.systemPromptContent).not.toBeNull(); - expect(Buffer.byteLength(cap.systemPromptContent ?? '', 'utf8')).toBeLessThanOrEqual( - SYSTEM_PROMPT_CAP_BYTES - ); + 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 request body remainder by UTF-8 bytes', () => { + it('truncates the body by UTF-8 bytes and marks wasTruncated', () => { const huge = 'a'.repeat(REQUEST_BODY_CAP_BYTES + 1024); const cap = buildExperimentPromptCapture( chatRequest({ @@ -140,41 +117,16 @@ describe('buildExperimentPromptCapture', () => { ); }); - it('counts UTF-8 bytes (not UTF-16 units) when checking the cap', () => { - // Build a string whose UTF-16 length is well under the cap but whose - // UTF-8 byte length exceeds it. CJK chars are 1 UTF-16 unit but 3 - // UTF-8 bytes; with cap=10 bytes, 5 CJK chars (5 units, 15 bytes) - // must trigger truncation. - const cap = buildExperimentPromptCapture({ - kind: 'chat_completions', - body: { - model: 'kilo/preview-foo', - messages: [ - { role: 'system', content: '日'.repeat(5) }, - { role: 'user', content: 'x' }, - ], - }, - } as unknown as GatewayRequest); - // We can't trigger this against the real 4 MB cap in a unit test, but - // we can assert the UTF-8 byte count of the captured content directly. - expect(cap.systemPromptContent).toBe('日'.repeat(5)); - expect(Buffer.byteLength(cap.systemPromptContent ?? '', 'utf8')).toBe(15); - expect(cap.systemPromptContent?.length).toBe(5); // UTF-16 unit count for sanity. - }); - - it('serializes non-string system content', () => { - const cap = buildExperimentPromptCapture( - chatRequest({ - model: 'kilo/preview-foo', - messages: [ - { - role: 'system', - content: [{ type: 'text', text: 'hello' }], - }, - { role: 'user', content: 'hi' }, - ], - }) - ); - expect(cap.systemPromptContent).toContain('"text":"hello"'); + 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 index d22a91c83ba..20012d23aee 100644 --- a/apps/web/src/lib/ai-gateway/experiments/persist.ts +++ b/apps/web/src/lib/ai-gateway/experiments/persist.ts @@ -6,23 +6,19 @@ import type { ExperimentPromptCapture } from '@/lib/ai-gateway/processUsage.type import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types'; /** - * Sentinel values written to `model_experiment_request.system_prompt_sha256` / - * `request_body_sha256` instead of a real 64-char hex digest. The schema - * CHECK constraints accept either a 64-char lowercase hex digest or one - * of these sentinels. + * 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. * - * - `__absent__` (system only): the request had no leading system message. * - `__failed__`: R2 storage failed; the attribution row still lands. * - `__deleted__`: prompt content intentionally wiped while retaining * experiment attribution. */ -export const PROMPT_HASH_ABSENT = '__absent__'; 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 SYSTEM_PROMPT_CAP_BYTES = 4 * 1024 * 1024; export const REQUEST_BODY_CAP_BYTES = 4 * 1024 * 1024; /** @@ -48,64 +44,25 @@ export function truncateToUtf8Bytes(s: string, maxBytes: number): string { /** * Build a bounded capture of the canonical post-`transformRequest` upstream - * request body, split into the leading system message and the remainder. + * 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. * - * - chat_completions: extracts a single leading `role === 'system'` message - * into `systemPromptContent` and serializes the remainder - * (`{...body, messages: nonSystem}`) into `requestBodyContent`. - * - messages / responses: leaves the full body in `requestBodyContent` and - * leaves `systemPromptContent` null. The provider-specific system field - * (e.g. anthropic `body.system`) is captured as part of the body blob. - * - * The 4 MB caps ensure we never retain unbounded data through the async - * `after()` path. If either side exceeds its cap, it is tail-truncated + * 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 wasTruncated = false; - - if (request.kind === 'chat_completions') { - const messages = Array.isArray(request.body.messages) ? request.body.messages : []; - let systemContent: string | null = null; - let nonSystemMessages = messages; - if (messages.length > 0 && messages[0]?.role === 'system') { - systemContent = serialize(messages[0].content ?? ''); - nonSystemMessages = messages.slice(1); - } - const remainder = { ...request.body, messages: nonSystemMessages }; - let bodyContent = serialize(remainder); - if ( - systemContent !== null && - Buffer.byteLength(systemContent, 'utf8') > SYSTEM_PROMPT_CAP_BYTES - ) { - systemContent = truncateToUtf8Bytes(systemContent, SYSTEM_PROMPT_CAP_BYTES); - wasTruncated = true; - } - if (Buffer.byteLength(bodyContent, 'utf8') > REQUEST_BODY_CAP_BYTES) { - bodyContent = truncateToUtf8Bytes(bodyContent, REQUEST_BODY_CAP_BYTES); - wasTruncated = true; - } - if (wasTruncated) { - noteTruncation(request.body.model); - } - return { - systemPromptContent: systemContent, - requestBodyContent: bodyContent, - wasTruncated, - }; - } - - // messages / responses: fold the entire serialized body into the body - // side; system extraction for those API kinds is left for later if it - // turns out to dedup meaningfully. let bodyContent = serialize(request.body); + let wasTruncated = false; if (Buffer.byteLength(bodyContent, 'utf8') > REQUEST_BODY_CAP_BYTES) { bodyContent = truncateToUtf8Bytes(bodyContent, REQUEST_BODY_CAP_BYTES); wasTruncated = true; noteTruncation(request.body.model); } return { - systemPromptContent: null, + requestKind: request.kind, requestBodyContent: bodyContent, wasTruncated, }; @@ -137,8 +94,8 @@ export type PersistExperimentAttributionInput = { /** * Insert one row into `model_experiment_request` for an experimented - * request. Best-effort R2 puts run in parallel; only the failed side - * gets `__failed__`. + * 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. @@ -147,17 +104,15 @@ export async function persistExperimentAttribution( input: PersistExperimentAttributionInput ): Promise { try { - const { systemHash, bodyHash } = await storePrompts(input.capture); - 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, - system_prompt_sha256: systemHash, - request_body_sha256: bodyHash, - was_truncated: input.capture?.wasTruncated ?? false, - created_at: input.createdAt, - }); + if (!input.capture) { + // Should not happen in production: an experimented request always + // builds a capture in the route handler. Record `__failed__` so the + // attribution row still lands and the gap is observable. + await insertRow(input, 'chat_completions', PROMPT_HASH_FAILED, false); + return; + } + const bodyHash = await putPromptSafely(input.capture.requestBodyContent); + await insertRow(input, input.capture.requestKind, bodyHash, input.capture.wasTruncated); } catch (err) { captureException(err, { tags: { source: 'model-experiments', operation: 'persistExperimentAttribution' }, @@ -166,49 +121,31 @@ export async function persistExperimentAttribution( } } -async function storePrompts( - capture: ExperimentPromptCapture | null -): Promise<{ systemHash: string; bodyHash: string }> { - if (!capture) { - return { systemHash: PROMPT_HASH_ABSENT, bodyHash: PROMPT_HASH_FAILED }; - } - const systemPromise = - capture.systemPromptContent === null - ? Promise.resolve({ kind: 'absent' }) - : putPromptSafely(capture.systemPromptContent); - const bodyPromise = putPromptSafely(capture.requestBodyContent); - const [systemResult, bodyResult] = await Promise.allSettled([systemPromise, bodyPromise]); - return { - systemHash: resolveHash(systemResult, 'system'), - bodyHash: resolveHash(bodyResult, 'body'), - }; +async function insertRow( + input: PersistExperimentAttributionInput, + requestKind: 'chat_completions' | 'messages' | 'responses', + bodyHash: string, + wasTruncated: boolean +): Promise { + 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: bodyHash, + was_truncated: wasTruncated, + created_at: input.createdAt, + }); } -type PromptSlotResult = { kind: 'absent' } | { kind: 'sha'; sha: string } | { kind: 'failed' }; - -async function putPromptSafely(content: string): Promise { +async function putPromptSafely(content: string): Promise { try { - return { kind: 'sha', sha: await putPromptIfAbsent(content) }; + return await putPromptIfAbsent(content); } catch (err) { captureException(err, { tags: { source: 'model-experiments', operation: 'putPromptIfAbsent' }, }); - return { kind: 'failed' }; - } -} - -function resolveHash( - settled: PromiseSettledResult, - side: 'system' | 'body' -): string { - if (settled.status === 'rejected') { - captureException(settled.reason, { - tags: { source: 'model-experiments', operation: 'storePrompts', side }, - }); return PROMPT_HASH_FAILED; } - const result = settled.value; - if (result.kind === 'absent') return PROMPT_HASH_ABSENT; - if (result.kind === 'sha') return result.sha; - return PROMPT_HASH_FAILED; } diff --git a/apps/web/src/lib/ai-gateway/processUsage.types.ts b/apps/web/src/lib/ai-gateway/processUsage.types.ts index b16e5b1bf50..69628a8d299 100644 --- a/apps/web/src/lib/ai-gateway/processUsage.types.ts +++ b/apps/web/src/lib/ai-gateway/processUsage.types.ts @@ -166,19 +166,19 @@ export type MicrodollarUsageContext = { modelExperimentAllocationSubject?: 'user' | 'machine' | 'ip'; /** * Bounded (size-capped) capture of the canonical post-`transformRequest` - * upstream request body, split into the leading system message and the - * remainder. Set only for experimented requests. Consumed by - * `persistExperimentAttribution` after the microdollar write. + * 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. Each side - * is either a string (already truncated to fit the cap) or null when absent. + * 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 = { - systemPromptContent: string | null; + requestKind: 'chat_completions' | 'messages' | 'responses'; requestBodyContent: string; wasTruncated: boolean; }; From 454afb853f6c6b37d5a8c839a024821690e79a37 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Wed, 20 May 2026 11:02:06 +0200 Subject: [PATCH 21/49] refactor(gateway): make balance <= 0 the primary condition on the usage-limit guard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reorders the conjuncts on the usage-limit guard so the salient check reads first and `skipBalanceCheck` is clearly an override flag, not the gating condition. Pure reorder — boolean evaluation is unchanged. Addresses comment #1 on #3325. --- apps/web/src/app/api/openrouter/[...path]/route.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/web/src/app/api/openrouter/[...path]/route.ts b/apps/web/src/app/api/openrouter/[...path]/route.ts index 19ab99b6080..17f269e2d29 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.ts @@ -496,7 +496,7 @@ export async function POST(request: NextRequest): Promise Date: Wed, 20 May 2026 11:02:59 +0200 Subject: [PATCH 22/49] fix(experiments): backfill the experimented-public-ids redis cache on miss MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously a Redis miss / corrupt cache / Redis error in `isPublicIdExperimented` fell through to Postgres for a single-id lookup and returned the answer without populating Redis. Result: every request on a cold cache hammered Postgres until something else happened to write the key (admin mutation, etc.). On the cache-miss path we now rebuild the full routing-eligible public id set from `model_experiment` (matching what `recomputeExperimentedPublicIds` writes from the admin router) and SET it back to Redis with the same TTL. Subsequent requests hit Redis again. The backfill is fire-and-forget — if Redis is genuinely down we still return the correct answer for this request. Also drops the per-id-only DB query in favor of the full-set query because the cache contract is the complete list, not per-id presence. Addresses comment #4 on #3325. --- .../ai-gateway/experiments/pick-variant.ts | 42 ++++++++++++++----- 1 file changed, 31 insertions(+), 11 deletions(-) diff --git a/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts index 25968b590aa..84f9750a4e8 100644 --- a/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts @@ -116,18 +116,18 @@ export async function isPublicIdExperimented( // captureException already invoked by the redis helper } + // Cache miss / corrupt / Redis error: rebuild the full set from DB and + // backfill the cache so the next request hits Redis. We deliberately + // query the full set rather than just `publicId` — the cache contract + // is "the complete list of routing-eligible public ids", and a single-id + // lookup can't satisfy it. Falling through without a backfill (the + // previous behavior) sent every miss straight to Postgres. try { - const rows = await db - .select({ id: model_experiment.id }) - .from(model_experiment) - .where( - and( - eq(model_experiment.public_model_id, publicId), - inArray(model_experiment.status, ['active', 'paused']) - ) - ) - .limit(1); - return rows.length > 0; + const ids = await loadExperimentedPublicIdsFromDb(); + // Best-effort backfill; if Redis is genuinely down the helper has + // already reported it and we still return the right answer below. + void backfillExperimentedPublicIdsCache(ids); + return ids.includes(publicId); } catch (err) { captureException(err, { tags: { source: 'model-experiments', operation: 'isPublicIdExperimented' }, @@ -137,6 +137,26 @@ export async function isPublicIdExperimented( } } +async function loadExperimentedPublicIdsFromDb(): Promise { + const rows = await db + .select({ public_model_id: model_experiment.public_model_id }) + .from(model_experiment) + .where(inArray(model_experiment.status, ['active', 'paused'])); + return Array.from(new Set(rows.map(r => r.public_model_id))).sort(); +} + +async function backfillExperimentedPublicIdsCache(ids: string[]): Promise { + try { + await redisSet( + EXPERIMENTED_PUBLIC_IDS_REDIS_KEY, + JSON.stringify(ids), + PER_PUBLIC_ID_CACHE_TTL_SECONDS + ); + } catch { + // captureException already invoked by the redis helper + } +} + function parseStringArray(raw: string): string[] | null { try { const parsed = JSON.parse(raw); From 0498dda64f92ff3161d69a413c314f8ecabccd70 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Wed, 20 May 2026 11:03:42 +0200 Subject: [PATCH 23/49] =?UTF-8?q?refactor(experiments):=20rename=20subject?= =?UTF-8?q?Pick=20=E2=86=92=20allocationSubject,=20fix=20lying=20'unreacha?= =?UTF-8?q?ble'=20comment?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two small clarity fixes from #3325 review: - Rename the local variable bound from `pickAllocationSubject(input)` from `subjectPick` to `allocationSubject`. Matches the field name on the result and answers the reviewer's 'is this the user?' question inline. Pure rename. - Replace the 'Unreachable: getRandomNumber returns [0, totalWeight). Fail closed.' comment with an honest one. The branch is defensive fail-closed insurance for the case where `getRandomNumber`'s contract breaks (NaN, negative, weight rounding), not actually unreachable code. The captureMessage + 'unavailable' return is the right behavior; only the comment was wrong. Addresses comments #5 and #6 on #3325. --- .../lib/ai-gateway/experiments/pick-variant.ts | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts index 84f9750a4e8..861bcfdd261 100644 --- a/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts @@ -420,8 +420,8 @@ export async function pickModelExperimentVariant( return { status: 'not-found' }; } - const subjectPick = pickAllocationSubject(input); - if (!subjectPick) { + const allocationSubject = pickAllocationSubject(input); + if (!allocationSubject) { captureMessage('Experiment request missing all allocation subjects', { level: 'error', tags: { source: 'model-experiments' }, @@ -440,7 +440,7 @@ export async function pickModelExperimentVariant( return { status: 'unavailable' }; } - const seed = `model_exp_${exp.experimentId}_${subjectPick.subject}_${subjectPick.value}`; + 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 @@ -480,12 +480,17 @@ export async function pickModelExperimentVariant( variantId: v.variantId, variantVersionId: v.variantVersionId, upstream: { ...v.upstream, api_key: apiKey }, - allocationSubject: subjectPick.subject, + allocationSubject: allocationSubject.subject, }; } } - // Unreachable: getRandomNumber returns [0, totalWeight). Fail closed. + // 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' }, From e0fa12267a30416750712a69753885c9b6280554 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Wed, 20 May 2026 12:11:08 +0200 Subject: [PATCH 24/49] fix(gateway): return usage identity for zero-cost inserts --- .../src/lib/ai-gateway/processUsage.test.ts | 8 +- apps/web/src/lib/ai-gateway/processUsage.ts | 127 +++++++++++------- 2 files changed, 87 insertions(+), 48 deletions(-) diff --git a/apps/web/src/lib/ai-gateway/processUsage.test.ts b/apps/web/src/lib/ai-gateway/processUsage.test.ts index 4ec09499432..2e11f70a0fb 100644 --- a/apps/web/src/lib/ai-gateway/processUsage.test.ts +++ b/apps/web/src/lib/ai-gateway/processUsage.test.ts @@ -471,6 +471,7 @@ describe('logMicrodollarUsage', () => { messageId: 'test-msg-456', hasError: true, cost_mUsd: 0, // Zero cost + market_cost: 500, model: 'openai/gpt-4.1', }; @@ -486,7 +487,7 @@ describe('logMicrodollarUsage', () => { }, }; - await logMicrodollarUsage(usageStats, usageContext); + const usageIdentity = await logMicrodollarUsage(usageStats, usageContext); // Verify user microdollars were NOT incremented const updatedUser = await findUserById('test-log-user-2'); @@ -501,8 +502,13 @@ describe('logMicrodollarUsage', () => { where: eq(microdollar_usage.id, metadataRecord!.id), }); expect(usageRecord).toBeTruthy(); + expect(usageIdentity).toEqual({ + usageId: usageRecord?.id, + createdAt: usageRecord?.created_at, + }); expect(usageRecord?.kilo_user_id).toBe('test-log-user-2'); expect(usageRecord?.cost).toBe(0); + expect(metadataRecord?.market_cost).toBe(500); expect(usageRecord?.has_error).toBe(true); expect(usageRecord?.model).toBe('openai/gpt-4.1'); expect(metadataRecord?.has_middle_out_transform).toBe(false); diff --git a/apps/web/src/lib/ai-gateway/processUsage.ts b/apps/web/src/lib/ai-gateway/processUsage.ts index fea64509aaa..50b05de78e7 100644 --- a/apps/web/src/lib/ai-gateway/processUsage.ts +++ b/apps/web/src/lib/ai-gateway/processUsage.ts @@ -109,6 +109,12 @@ const extractMessageTextContent = (m: Message) => export type UsageContextInfo = ReturnType; +export type UsageRecordInsertResult = { + usageId: string; + createdAt: string; + newMicrodollarsUsed: number | null; +}; + export function extractUsageContextInfo(usageContext: MicrodollarUsageContext) { return { kilo_user_id: usageContext.kiloUserId, @@ -260,7 +266,7 @@ export async function logMicrodollarUsage( // `insertUsageRecord` swallows DB errors and returns null; surface that // failure to callers so dependent FK writes don't dangle on a row that // was never persisted. - return inserted ? { usageId: core.id, createdAt: core.created_at } : null; + return inserted ? { usageId: inserted.usageId, createdAt: inserted.createdAt } : null; } async function saveUsageRelatedData( @@ -268,21 +274,24 @@ async function saveUsageRelatedData( metadataFields: UsageMetaData, prior_microdollar_usage: number, posthog_distinct_id: string | null -): Promise { +): Promise { const isFirst = await isFirstUsage(coreUsageFields, prior_microdollar_usage); if (isFirst && posthog_distinct_id) await sendFirstUsageEvent(coreUsageFields, posthog_distinct_id); - const balanceUpdateResult = await insertUsageRecord(coreUsageFields, metadataFields); + const inserted = await insertUsageRecord(coreUsageFields, metadataFields); + if (!inserted) return null; if (posthog_distinct_id) { await sendFirstMicrodollarUsageEventIfNeeded( - balanceUpdateResult, + inserted.newMicrodollarsUsed === null + ? null + : { newMicrodollarsUsed: inserted.newMicrodollarsUsed }, coreUsageFields, posthog_distinct_id, isFirst ); } await ingestOrganizationTokenUsage(coreUsageFields); - return balanceUpdateResult; + return inserted; } async function isFirstUsage( @@ -389,7 +398,7 @@ ${metaDataKindName}_cte AS ( export async function insertUsageRecord( coreUsageFields: MicrodollarUsage, metadataFields: UsageMetaData -): Promise { +): Promise { try { const result = await startSpan( { @@ -428,7 +437,7 @@ export async function insertUsageRecord( async function insertUsageAndMetadataWithBalanceUpdate( coreUsageFields: MicrodollarUsage, metadataFields: UsageMetaData -): Promise { +): Promise { // Pick the matching partial unique index for the daily-rollup upsert. The // microdollar_usage_daily table has two partial unique indexes; the upsert // must target the one corresponding to this row's scope. @@ -440,7 +449,9 @@ async function insertUsageAndMetadataWithBalanceUpdate( // Use a single SQL statement with CTEs to insert usage, upsert all lookup values, metadata, and update user balance in one roundtrip // This ensures atomicity: microdollar_usage insert and kilocode_users.microdollars_used update happen together const result = await db.execute<{ - new_microdollars_used: number; + usage_id: string; + usage_created_at: string; + new_microdollars_used: number | null; kilo_pass_threshold: number | null; }>(sql` WITH microdollar_usage_ins AS ( @@ -468,6 +479,7 @@ async function insertUsageAndMetadataWithBalanceUpdate( ${coreUsageFields.inference_provider}, ${coreUsageFields.project_id} ) + RETURNING id, created_at ) , ${createUpsertCTE(sql`http_user_agent`, metadataFields.http_user_agent)} , ${createUpsertCTE(sql`http_ip`, metadataFields.http_x_forwarded_for)} @@ -574,55 +586,76 @@ async function insertUsageAndMetadataWithBalanceUpdate( microdollar_usage_daily.total_cost_microdollars + EXCLUDED.total_cost_microdollars, updated_at = NOW() ) - UPDATE kilocode_users - SET microdollars_used = microdollars_used + ${coreUsageFields.cost} - WHERE id = ${coreUsageFields.kilo_user_id} - AND ${coreUsageFields.organization_id}::uuid IS NULL - AND ${coreUsageFields.cost} > 0 - RETURNING microdollars_used AS new_microdollars_used, kilo_pass_threshold + , balance_update AS ( + UPDATE kilocode_users + SET microdollars_used = microdollars_used + ${coreUsageFields.cost} + WHERE id = ${coreUsageFields.kilo_user_id} + AND ${coreUsageFields.organization_id}::uuid IS NULL + AND ${coreUsageFields.cost} > 0 + RETURNING microdollars_used AS new_microdollars_used, kilo_pass_threshold + ) + SELECT + microdollar_usage_ins.id AS usage_id, + microdollar_usage_ins.created_at AS usage_created_at, + balance_update.new_microdollars_used, + balance_update.kilo_pass_threshold + FROM microdollar_usage_ins + LEFT JOIN balance_update ON true `); - // No rows returned means either: org usage (no user balance update), zero cost, or missing user - if (!result.rows[0]) { - // Only log error if we expected an update (non-org, positive cost) - if (!coreUsageFields.organization_id && coreUsageFields.cost && coreUsageFields.cost > 0) { - captureMessage('impossible: missing user', { - level: 'fatal', - tags: { source: 'insertUsageAndUpdateBalance' }, - extra: { coreUsageFields }, - }); - } - return null; + const inserted = result.rows[0]; + if (!inserted) { + throw new Error('microdollar_usage insert returned no identity'); + } + + // Missing balance update is expected for org usage and zero-cost rows, but + // suspicious for positive-cost personal usage. + if ( + inserted.new_microdollars_used === null && + !coreUsageFields.organization_id && + coreUsageFields.cost > 0 + ) { + captureMessage('impossible: missing user', { + level: 'fatal', + tags: { source: 'insertUsageAndUpdateBalance' }, + extra: { coreUsageFields }, + }); } - // Convert BigInt to number (microdollars_used is a bigint column) - const newMicrodollarsUsed = Number(result.rows[0].new_microdollars_used); + const newMicrodollarsUsed = + inserted.new_microdollars_used === null ? null : Number(inserted.new_microdollars_used); const kiloPassThreshold = - result.rows[0].kilo_pass_threshold == null ? null : Number(result.rows[0].kilo_pass_threshold); - - const effectiveKiloPassThreshold = getEffectiveKiloPassThreshold(kiloPassThreshold); - - if (effectiveKiloPassThreshold !== null && newMicrodollarsUsed >= effectiveKiloPassThreshold) { - // Trigger this async to avoid blocking - void maybeIssueKiloPassBonusFromUsageThreshold({ - kiloUserId: coreUsageFields.kilo_user_id, - nowIso: coreUsageFields.created_at, - }).catch(async error => { - const errorMessage = error instanceof Error ? error.message : String(error); - await appendKiloPassAuditLog(db, { - action: KiloPassAuditLogAction.BonusCreditsIssued, - result: KiloPassAuditLogResult.Failed, + inserted.kilo_pass_threshold == null ? null : Number(inserted.kilo_pass_threshold); + + if (newMicrodollarsUsed !== null) { + const effectiveKiloPassThreshold = getEffectiveKiloPassThreshold(kiloPassThreshold); + + if (effectiveKiloPassThreshold !== null && newMicrodollarsUsed >= effectiveKiloPassThreshold) { + // Trigger this async to avoid blocking + void maybeIssueKiloPassBonusFromUsageThreshold({ kiloUserId: coreUsageFields.kilo_user_id, - payload: { - source: 'usage_threshold', - error: errorMessage, - }, + nowIso: coreUsageFields.created_at, + }).catch(async error => { + const errorMessage = error instanceof Error ? error.message : String(error); + await appendKiloPassAuditLog(db, { + action: KiloPassAuditLogAction.BonusCreditsIssued, + result: KiloPassAuditLogResult.Failed, + kiloUserId: coreUsageFields.kilo_user_id, + payload: { + source: 'usage_threshold', + error: errorMessage, + }, + }); }); - }); + } } - return { newMicrodollarsUsed }; + return { + usageId: inserted.usage_id, + createdAt: inserted.usage_created_at, + newMicrodollarsUsed, + }; } export function countAndStoreUsage( From 0a54c026e150fe7ca53c88d33520c12f077622ed Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Wed, 20 May 2026 14:45:57 +0200 Subject: [PATCH 25/49] fix(gateway): use JS-side identity values to avoid Postgres timestamp format --- apps/web/src/lib/ai-gateway/processUsage.ts | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/apps/web/src/lib/ai-gateway/processUsage.ts b/apps/web/src/lib/ai-gateway/processUsage.ts index 50b05de78e7..de0c222e65b 100644 --- a/apps/web/src/lib/ai-gateway/processUsage.ts +++ b/apps/web/src/lib/ai-gateway/processUsage.ts @@ -266,7 +266,12 @@ export async function logMicrodollarUsage( // `insertUsageRecord` swallows DB errors and returns null; surface that // failure to callers so dependent FK writes don't dangle on a row that // was never persisted. - return inserted ? { usageId: inserted.usageId, createdAt: inserted.createdAt } : null; + // Use the JS-side identity values we constructed in toInsertableDbUsageRecord + // rather than the DB-returned ones. The DB round-trip for created_at returns a + // Postgres timestamp string (e.g. "2026-04-29 01:16:12.945+00") which is not + // strict ISO 8601 and will fail downstream datetime validators. core.created_at + // is always new Date().toISOString() so the format is guaranteed. + return inserted ? { usageId: core.id, createdAt: core.created_at } : null; } async function saveUsageRelatedData( From 3f9e18c2f5697d2ee637ddeb2ecb311b424becb3 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Wed, 20 May 2026 15:53:54 +0200 Subject: [PATCH 26/49] refactor(gateway): extract experiment variant types --- .../ai-gateway/experiments/pick-variant.ts | 88 ++++--------------- .../experiments/pick-variant.types.ts | 68 ++++++++++++++ 2 files changed, 86 insertions(+), 70 deletions(-) create mode 100644 apps/web/src/lib/ai-gateway/experiments/pick-variant.types.ts diff --git a/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts index 861bcfdd261..432c35e360f 100644 --- a/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts @@ -4,7 +4,6 @@ import { model_experiment, model_experiment_variant, model_experiment_variant_version, - type ModelExperiment, } from '@kilocode/db/schema'; import { db } from '@/lib/drizzle'; import { redisGet, redisSet } from '@/lib/redis'; @@ -12,11 +11,18 @@ import { EXPERIMENTED_PUBLIC_IDS_REDIS_KEY, modelExperimentRedisKey } from '@/li 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, - type ExperimentUpstream, -} from '@/lib/ai-gateway/experiments/upstream-schema'; -import type { ResolvedExperimentUpstream } from '@/lib/ai-gateway/experiments/build-direct-provider'; +import { ExperimentUpstreamSchema } from '@/lib/ai-gateway/experiments/upstream-schema'; +import type { + AllocationSubject, + CachedExperiment, + CachedVariant, + ExperimentStatus, + PickVariantInput, + PickVariantResult, + ResolveResult, + SerializedCacheEntry, + SerializedCachedExperiment, +} from '@/lib/ai-gateway/experiments/pick-variant.types'; /** TTL for the per-public-id experiment cache. The cache holds CIPHERTEXT * (encrypted api keys + the upstream blob); decryption happens per-pick @@ -25,58 +31,6 @@ import type { ResolvedExperimentUpstream } from '@/lib/ai-gateway/experiments/bu * how long plaintext keys can live in Redis. */ const PER_PUBLIC_ID_CACHE_TTL_SECONDS = 600; -type ExperimentStatus = 'active' | 'paused'; - -/** - * Per-variant payload cached in Redis. - * - * The api key is stored ENCRYPTED — the cache never holds plaintext. - * Decryption happens once per request, only on the chosen variant, inside - * `pickModelExperimentVariant`. This keeps `BYOK_ENCRYPTION_KEY` rotation - * effective immediately (no 10-minute "plaintext lives in Redis" window) - * and limits the blast radius of a Redis dump. - */ -type CachedVariant = { - variantId: string; - weight: number; - variantVersionId: string; - upstream: ExperimentUpstream; - encryptedApiKey: EncryptedData; -}; - -type CachedExperiment = { - experimentId: string; - publicModelId: string; - status: ExperimentStatus; - variants: CachedVariant[]; -}; - -type ResolveResult = - | { kind: 'experiment'; experiment: CachedExperiment } - | { 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' }; - let warnedMissingEncryptionKey = false; function warnMissingEncryptionKeyOnce() { @@ -225,16 +179,6 @@ export async function getRoutingExperimentForPublicId(publicId: string): Promise return resolved; } -type SerializedCachedExperiment = { - kind: 'experiment'; - experimentId: string; - publicModelId: string; - status: ExperimentStatus; - variants: CachedVariant[]; -}; - -type SerializedCacheEntry = SerializedCachedExperiment | { kind: 'none' }; - function serializeCachedExperiment(exp: CachedExperiment): SerializedCachedExperiment { return { kind: 'experiment', ...exp }; } @@ -508,5 +452,9 @@ function pickAllocationSubject( return null; } -// Re-export for callers that need the raw experiment row type. -export type { ModelExperiment }; +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..da02327db1e --- /dev/null +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.types.ts @@ -0,0 +1,68 @@ +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'; + +/** + * Per-variant payload cached in Redis. + * + * The api key is stored ENCRYPTED — the cache never holds plaintext. + * Decryption happens once per request, only on the chosen variant, inside + * `pickModelExperimentVariant`. This keeps `BYOK_ENCRYPTION_KEY` rotation + * effective immediately (no 10-minute "plaintext lives in Redis" window) + * and limits the blast radius of a Redis dump. + */ +export type CachedVariant = { + variantId: string; + weight: number; + variantVersionId: string; + upstream: ExperimentUpstream; + encryptedApiKey: EncryptedData; +}; + +export type CachedExperiment = { + experimentId: string; + publicModelId: string; + status: ExperimentStatus; + variants: CachedVariant[]; +}; + +export type ResolveResult = + | { kind: 'experiment'; experiment: CachedExperiment } + | { 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 SerializedCachedExperiment = { + kind: 'experiment'; + experimentId: string; + publicModelId: string; + status: ExperimentStatus; + variants: CachedVariant[]; +}; + +export type SerializedCacheEntry = SerializedCachedExperiment | { kind: 'none' }; + +export type { ModelExperiment }; From 92012a2240a14a296827c02afb6537f93e0fd4ed Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Wed, 20 May 2026 16:09:15 +0200 Subject: [PATCH 27/49] refactor(gateway): narrow experiment attribution catch --- .../src/lib/ai-gateway/experiments/persist.ts | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/apps/web/src/lib/ai-gateway/experiments/persist.ts b/apps/web/src/lib/ai-gateway/experiments/persist.ts index 20012d23aee..a9cb0a0464c 100644 --- a/apps/web/src/lib/ai-gateway/experiments/persist.ts +++ b/apps/web/src/lib/ai-gateway/experiments/persist.ts @@ -103,16 +103,14 @@ export type PersistExperimentAttributionInput = { export async function persistExperimentAttribution( input: PersistExperimentAttributionInput ): Promise { + const requestKind = input.capture?.requestKind ?? 'chat_completions'; + const wasTruncated = input.capture?.wasTruncated ?? false; + const bodyHash = input.capture + ? await putPromptOrFailedHash(input.capture.requestBodyContent) + : PROMPT_HASH_FAILED; + try { - if (!input.capture) { - // Should not happen in production: an experimented request always - // builds a capture in the route handler. Record `__failed__` so the - // attribution row still lands and the gap is observable. - await insertRow(input, 'chat_completions', PROMPT_HASH_FAILED, false); - return; - } - const bodyHash = await putPromptSafely(input.capture.requestBodyContent); - await insertRow(input, input.capture.requestKind, bodyHash, input.capture.wasTruncated); + await insertRow(input, requestKind, bodyHash, wasTruncated); } catch (err) { captureException(err, { tags: { source: 'model-experiments', operation: 'persistExperimentAttribution' }, @@ -139,7 +137,7 @@ async function insertRow( }); } -async function putPromptSafely(content: string): Promise { +async function putPromptOrFailedHash(content: string): Promise { try { return await putPromptIfAbsent(content); } catch (err) { From e5ec33afcceddc90b2341e0b2a0f05ff22847767 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Wed, 20 May 2026 16:18:59 +0200 Subject: [PATCH 28/49] refactor(gateway): move prompt storage fallback --- .../web/src/lib/ai-gateway/experiments/persist.ts | 15 ++------------- apps/web/src/lib/r2/experiment-prompts.ts | 11 +++++++++++ 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/apps/web/src/lib/ai-gateway/experiments/persist.ts b/apps/web/src/lib/ai-gateway/experiments/persist.ts index a9cb0a0464c..e759757b143 100644 --- a/apps/web/src/lib/ai-gateway/experiments/persist.ts +++ b/apps/web/src/lib/ai-gateway/experiments/persist.ts @@ -1,7 +1,7 @@ import { captureException, captureMessage } from '@sentry/nextjs'; import { db } from '@/lib/drizzle'; import { model_experiment_request } from '@kilocode/db/schema'; -import { putPromptIfAbsent } from '@/lib/r2/experiment-prompts'; +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'; @@ -106,7 +106,7 @@ export async function persistExperimentAttribution( const requestKind = input.capture?.requestKind ?? 'chat_completions'; const wasTruncated = input.capture?.wasTruncated ?? false; const bodyHash = input.capture - ? await putPromptOrFailedHash(input.capture.requestBodyContent) + ? ((await putPromptOrNull(input.capture.requestBodyContent)) ?? PROMPT_HASH_FAILED) : PROMPT_HASH_FAILED; try { @@ -136,14 +136,3 @@ async function insertRow( created_at: input.createdAt, }); } - -async function putPromptOrFailedHash(content: string): Promise { - try { - return await putPromptIfAbsent(content); - } catch (err) { - captureException(err, { - tags: { source: 'model-experiments', operation: 'putPromptIfAbsent' }, - }); - return PROMPT_HASH_FAILED; - } -} diff --git a/apps/web/src/lib/r2/experiment-prompts.ts b/apps/web/src/lib/r2/experiment-prompts.ts index 6f7f3805a11..4532f523377 100644 --- a/apps/web/src/lib/r2/experiment-prompts.ts +++ b/apps/web/src/lib/r2/experiment-prompts.ts @@ -76,6 +76,17 @@ export async function putPromptIfAbsent(content: string): Promise { 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__`, From dd083836cc3ac70caec90d3e3fbeaa502f18e9c3 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Wed, 20 May 2026 16:22:00 +0200 Subject: [PATCH 29/49] refactor(gateway): simplify prompt hash fallback --- apps/web/src/lib/ai-gateway/experiments/persist.ts | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/apps/web/src/lib/ai-gateway/experiments/persist.ts b/apps/web/src/lib/ai-gateway/experiments/persist.ts index e759757b143..c0f2dce8a06 100644 --- a/apps/web/src/lib/ai-gateway/experiments/persist.ts +++ b/apps/web/src/lib/ai-gateway/experiments/persist.ts @@ -105,12 +105,11 @@ export async function persistExperimentAttribution( ): Promise { const requestKind = input.capture?.requestKind ?? 'chat_completions'; const wasTruncated = input.capture?.wasTruncated ?? false; - const bodyHash = input.capture - ? ((await putPromptOrNull(input.capture.requestBodyContent)) ?? PROMPT_HASH_FAILED) - : PROMPT_HASH_FAILED; + const storedPromptHash = + input.capture && (await putPromptOrNull(input.capture.requestBodyContent)); try { - await insertRow(input, requestKind, bodyHash, wasTruncated); + await insertRow(input, requestKind, storedPromptHash ?? PROMPT_HASH_FAILED, wasTruncated); } catch (err) { captureException(err, { tags: { source: 'model-experiments', operation: 'persistExperimentAttribution' }, From 8fb0a45e9407a6d8bc36914a60848b508c27cd02 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Wed, 20 May 2026 16:23:59 +0200 Subject: [PATCH 30/49] refactor(gateway): inline experiment attribution insert --- .../src/lib/ai-gateway/experiments/persist.ts | 29 +++++++------------ 1 file changed, 10 insertions(+), 19 deletions(-) diff --git a/apps/web/src/lib/ai-gateway/experiments/persist.ts b/apps/web/src/lib/ai-gateway/experiments/persist.ts index c0f2dce8a06..171be7d0d17 100644 --- a/apps/web/src/lib/ai-gateway/experiments/persist.ts +++ b/apps/web/src/lib/ai-gateway/experiments/persist.ts @@ -109,7 +109,16 @@ export async function persistExperimentAttribution( input.capture && (await putPromptOrNull(input.capture.requestBodyContent)); try { - await insertRow(input, requestKind, storedPromptHash ?? PROMPT_HASH_FAILED, wasTruncated); + 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' }, @@ -117,21 +126,3 @@ export async function persistExperimentAttribution( }); } } - -async function insertRow( - input: PersistExperimentAttributionInput, - requestKind: 'chat_completions' | 'messages' | 'responses', - bodyHash: string, - wasTruncated: boolean -): Promise { - 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: bodyHash, - was_truncated: wasTruncated, - created_at: input.createdAt, - }); -} From 5f3ba98ef40b60f8c18cf13d7ed76d895bb53eae Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Wed, 20 May 2026 16:26:16 +0200 Subject: [PATCH 31/49] refactor(gateway): remove prompt truncation breadcrumb --- apps/web/src/lib/ai-gateway/experiments/persist.ts | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/apps/web/src/lib/ai-gateway/experiments/persist.ts b/apps/web/src/lib/ai-gateway/experiments/persist.ts index 171be7d0d17..b01ced18842 100644 --- a/apps/web/src/lib/ai-gateway/experiments/persist.ts +++ b/apps/web/src/lib/ai-gateway/experiments/persist.ts @@ -1,4 +1,4 @@ -import { captureException, captureMessage } from '@sentry/nextjs'; +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'; @@ -59,7 +59,6 @@ export function buildExperimentPromptCapture(request: GatewayRequest): Experimen if (Buffer.byteLength(bodyContent, 'utf8') > REQUEST_BODY_CAP_BYTES) { bodyContent = truncateToUtf8Bytes(bodyContent, REQUEST_BODY_CAP_BYTES); wasTruncated = true; - noteTruncation(request.body.model); } return { requestKind: request.kind, @@ -73,16 +72,6 @@ function serialize(value: unknown): string { return JSON.stringify(value); } -function noteTruncation(model: string | undefined) { - // Sentry breadcrumb only (no payload content) — we want to know if this - // ever fires at non-trivial rates so we can re-evaluate caps. - captureMessage('Experiment prompt capture truncated', { - level: 'info', - tags: { source: 'model-experiments', operation: 'buildExperimentPromptCapture' }, - extra: { model: model ?? null }, - }); -} - export type PersistExperimentAttributionInput = { usageId: string; createdAt: string; From 09d6f771d3f31f2ce9cd8c49811f7214821538e1 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Wed, 20 May 2026 16:29:45 +0200 Subject: [PATCH 32/49] refactor(gateway): inline prompt capture serialization --- apps/web/src/lib/ai-gateway/experiments/persist.ts | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/apps/web/src/lib/ai-gateway/experiments/persist.ts b/apps/web/src/lib/ai-gateway/experiments/persist.ts index b01ced18842..c057594a7e7 100644 --- a/apps/web/src/lib/ai-gateway/experiments/persist.ts +++ b/apps/web/src/lib/ai-gateway/experiments/persist.ts @@ -54,7 +54,7 @@ export function truncateToUtf8Bytes(s: string, maxBytes: number): string { * deterministically and `wasTruncated` is set to true. */ export function buildExperimentPromptCapture(request: GatewayRequest): ExperimentPromptCapture { - let bodyContent = serialize(request.body); + 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); @@ -67,11 +67,6 @@ export function buildExperimentPromptCapture(request: GatewayRequest): Experimen }; } -function serialize(value: unknown): string { - if (typeof value === 'string') return value; - return JSON.stringify(value); -} - export type PersistExperimentAttributionInput = { usageId: string; createdAt: string; From 7107be80072a427f6fc2441294d78f888e27518c Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Wed, 20 May 2026 16:34:33 +0200 Subject: [PATCH 33/49] refactor(gateway): simplify experiment attribution guard --- .../src/lib/ai-gateway/llm-proxy-helpers.ts | 23 +++++++++---------- 1 file changed, 11 insertions(+), 12 deletions(-) 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 7278a8d3c4d..96461b82d89 100644 --- a/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts +++ b/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts @@ -318,20 +318,19 @@ export function accountForMicrodollarUsage( // the time we reach here. `persistExperimentAttribution` // swallows errors internally. if ( - !usageIdentity || - !usageContext.modelExperimentVariantVersionId || - !usageContext.modelExperimentAllocationSubject + usageIdentity && + usageContext.modelExperimentVariantVersionId && + usageContext.modelExperimentAllocationSubject ) { - return; + await persistExperimentAttribution({ + usageId: usageIdentity.usageId, + createdAt: usageIdentity.createdAt, + variantVersionId: usageContext.modelExperimentVariantVersionId, + allocationSubject: usageContext.modelExperimentAllocationSubject, + clientRequestId: usageContext.clientRequestId ?? null, + capture: usageContext.experimentPromptCapture ?? null, + }); } - await persistExperimentAttribution({ - usageId: usageIdentity.usageId, - createdAt: usageIdentity.createdAt, - variantVersionId: usageContext.modelExperimentVariantVersionId, - allocationSubject: usageContext.modelExperimentAllocationSubject, - clientRequestId: usageContext.clientRequestId ?? null, - capture: usageContext.experimentPromptCapture ?? null, - }); } ) ); From 0ea4021f964b68ad8dd87ed0e1ebe9a67fd24e12 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Wed, 20 May 2026 16:54:52 +0200 Subject: [PATCH 34/49] fix(web): require BYOK encryption key --- apps/web/src/lib/config.server.ts | 7 +++++-- apps/web/src/lib/constants.ts | 12 ++---------- apps/web/src/lib/dotenvx.ts | 9 +++++++++ .../src/routers/admin/model-experiments-router.ts | 11 ++--------- apps/web/src/routers/byok-router.ts | 14 -------------- 5 files changed, 18 insertions(+), 35 deletions(-) diff --git a/apps/web/src/lib/config.server.ts b/apps/web/src/lib/config.server.ts index b2bc5e9a755..2ddf4dd5bc7 100644 --- a/apps/web/src/lib/config.server.ts +++ b/apps/web/src/lib/config.server.ts @@ -1,5 +1,5 @@ import { APP_URL } from '@/lib/constants'; -import { getEnvVariable } from '@/lib/dotenvx'; +import { getEnvVariable, requireEnv } from '@/lib/dotenvx'; import 'server-only'; export const IS_IN_AUTOMATED_TEST = !!getEnvVariable('IS_IN_AUTOMATED_TEST'); @@ -101,7 +101,10 @@ export const USER_DEPLOYMENTS_GIT_TOKEN_ENCRYPTION_KEY = getEnvVariable( * Must be a base64-encoded 32-byte (256-bit) key. * Generate with: node -e "console.log(require('crypto').randomBytes(32).toString('base64'))" */ -export const BYOK_ENCRYPTION_KEY = getEnvVariable('BYOK_ENCRYPTION_KEY') || ''; +export const BYOK_ENCRYPTION_KEY = requireEnv( + 'BYOK_ENCRYPTION_KEY', + getEnvVariable('BYOK_ENCRYPTION_KEY') +); // Artificial Analysis API export const ARTIFICIAL_ANALYSIS_API_KEY = getEnvVariable('ARTIFICIAL_ANALYSIS_API_KEY'); diff --git a/apps/web/src/lib/constants.ts b/apps/web/src/lib/constants.ts index 40bedc6b360..20be917978b 100644 --- a/apps/web/src/lib/constants.ts +++ b/apps/web/src/lib/constants.ts @@ -1,3 +1,5 @@ +import { requireEnv } from '@/lib/dotenvx'; + export const FIRST_TOPUP_BONUS_AMOUNT: number = 0; export const REFERRAL_BONUS_AMOUNT = 10; @@ -43,16 +45,6 @@ export const CLOUD_AGENT_NEXT_WS_URL = process.env.NEXT_PUBLIC_CLOUD_AGENT_NEXT_ // Used by the CLI live transport for real-time event streaming export const SESSION_INGEST_WS_URL = process.env.NEXT_PUBLIC_SESSION_INGEST_WS_URL ?? ''; -// Next.js inlines NEXT_PUBLIC_* at build time. Fail loudly when required -// public URLs are missing so misconfiguration surfaces at startup instead of -// producing relative URLs against the app origin. -function requireEnv(name: string, value: string | undefined): string { - if (!value) { - throw new Error(`Missing required environment variable ${name}`); - } - return value; -} - // Gastown worker URL (client-side, inlined at build time) // The browser talks directly to the gastown Cloudflare Worker for tRPC + WS. // Must use NEXT_PUBLIC_ prefix so Next.js exposes it to the browser bundle. diff --git a/apps/web/src/lib/dotenvx.ts b/apps/web/src/lib/dotenvx.ts index 4a10f3f5fb7..71202bbb1dd 100644 --- a/apps/web/src/lib/dotenvx.ts +++ b/apps/web/src/lib/dotenvx.ts @@ -6,3 +6,12 @@ export function getEnvVariable(key: string): string { return process.env[key] || ''; } + +// Next.js inlines NEXT_PUBLIC_* at build time. Fail loudly when required +// environment variables are missing so misconfiguration surfaces at startup. +export function requireEnv(name: string, value: string | undefined): string { + if (!value) { + throw new Error(`Missing required environment variable ${name}`); + } + return value; +} diff --git a/apps/web/src/routers/admin/model-experiments-router.ts b/apps/web/src/routers/admin/model-experiments-router.ts index 3c2b4e37632..56e51c577ae 100644 --- a/apps/web/src/routers/admin/model-experiments-router.ts +++ b/apps/web/src/routers/admin/model-experiments-router.ts @@ -38,13 +38,6 @@ const weightSchema = z.number().int().positive(); const apiKeySchema = z.string().min(1); -function ensureEncryptionKey(): string { - if (!BYOK_ENCRYPTION_KEY) { - trpcThrow('INTERNAL_SERVER_ERROR', 'BYOK encryption is not configured'); - } - return BYOK_ENCRYPTION_KEY; -} - async function loadExperimentOrThrow(id: string) { const row = await db.query.model_experiment.findFirst({ where: eq(model_experiment.id, id), @@ -451,7 +444,7 @@ export const adminModelExperimentsRouter = createTRPCRouter({ // copy and the key is required. let encrypted_api_key; if (input.apiKey !== undefined) { - encrypted_api_key = encryptApiKey(input.apiKey, ensureEncryptionKey()); + encrypted_api_key = encryptApiKey(input.apiKey, BYOK_ENCRYPTION_KEY); } else { const previous = await db .select({ encrypted_api_key: model_experiment_variant_version.encrypted_api_key }) @@ -482,7 +475,7 @@ export const adminModelExperimentsRouter = createTRPCRouter({ }), rotateApiKey: adminProcedure.input(RotateApiKeySchema).mutation(async ({ input, ctx }) => { - const key = ensureEncryptionKey(); + const key = BYOK_ENCRYPTION_KEY; const variant = await loadVariantOrThrow(input.variantId); const experiment = await loadExperimentOrThrow(variant.experiment_id); assertNonTerminal(experiment.status as Status, 'Rotating api key'); diff --git a/apps/web/src/routers/byok-router.ts b/apps/web/src/routers/byok-router.ts index ae58a59f63c..c9b37e0e995 100644 --- a/apps/web/src/routers/byok-router.ts +++ b/apps/web/src/routers/byok-router.ts @@ -185,13 +185,6 @@ export const byokRouter = createTRPCRouter({ await ensureOrganizationAccess(ctx, organizationId, ['owner', 'billing_manager']); } - if (!BYOK_ENCRYPTION_KEY) { - throw new TRPCError({ - code: 'INTERNAL_SERVER_ERROR', - message: 'BYOK encryption is not configured', - }); - } - // Encrypt the API key const encrypted = encryptApiKey(api_key, BYOK_ENCRYPTION_KEY); @@ -243,13 +236,6 @@ export const byokRouter = createTRPCRouter({ await ensureOrganizationAccess(ctx, organizationId, ['owner', 'billing_manager']); } - if (!BYOK_ENCRYPTION_KEY) { - throw new TRPCError({ - code: 'INTERNAL_SERVER_ERROR', - message: 'BYOK encryption is not configured', - }); - } - // Verify key exists and belongs to the organization or user const [existingKey] = await db .select({ From 131a6a2b711916ff9e64c1664b0d3dfd63738d8d Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Wed, 20 May 2026 17:05:00 +0200 Subject: [PATCH 35/49] refactor(gateway): remove missing BYOK key branch --- .../ai-gateway/experiments/pick-variant.ts | 33 ++++--------------- 1 file changed, 6 insertions(+), 27 deletions(-) diff --git a/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts index 432c35e360f..eee282546ff 100644 --- a/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts @@ -31,20 +31,6 @@ import type { * how long plaintext keys can live in Redis. */ const PER_PUBLIC_ID_CACHE_TTL_SECONDS = 600; -let warnedMissingEncryptionKey = false; - -function warnMissingEncryptionKeyOnce() { - if (warnedMissingEncryptionKey) return; - warnedMissingEncryptionKey = true; - captureMessage('BYOK_ENCRYPTION_KEY is unset; model experiment routing returning unavailable', { - level: 'error', - tags: { source: 'model-experiments' }, - }); - console.error( - '[model-experiments] BYOK_ENCRYPTION_KEY is unset; experiment routing will return unavailable' - ); -} - /** * Fast pre-check: does any active|paused experiment target this public id? * @@ -292,11 +278,9 @@ async function loadExperimentFromDb(publicId: string): Promise { } } - // We deliberately DO NOT call decryptApiKey here, even when - // BYOK_ENCRYPTION_KEY is missing. The encrypted blob is harmless on its - // own; decryption + the missing-key check happen per-pick in - // `pickModelExperimentVariant` so a warm cache cannot serve plaintext - // keys after key rotation/removal. + // We deliberately DO NOT call decryptApiKey here. The cache should only + // hold encrypted blobs; decryption happens per-pick so key rotation takes + // effect on the next request. const variants: CachedVariant[] = []; for (const v of variantRows) { const ver = versionByVariantId.get(v.id); @@ -396,14 +380,9 @@ export async function pickModelExperimentVariant( // 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 - // BYOK_ENCRYPTION_KEY missing-key fail-closed effective on warm - // caches, (b) lets key rotation invalidate access immediately - // (next request after the key flips fails closed), and (c) means - // we never decrypt N-1 keys we won't use. - if (!BYOK_ENCRYPTION_KEY) { - warnMissingEncryptionKeyOnce(); - return { status: 'unavailable' }; - } + // 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); From caf2002eccb4ac929e12f0f1bf3b81fee5921171 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Wed, 20 May 2026 17:26:10 +0200 Subject: [PATCH 36/49] refactor(gateway): cache experimented ids from Redis only --- .../ai-gateway/experiments/pick-variant.ts | 100 ++++++------------ .../lib/ai-gateway/providers/get-provider.ts | 3 - 2 files changed, 32 insertions(+), 71 deletions(-) diff --git a/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts index eee282546ff..c77a665df02 100644 --- a/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts @@ -8,6 +8,7 @@ import { import { db } from '@/lib/drizzle'; import { redisGet, redisSet } from '@/lib/redis'; import { EXPERIMENTED_PUBLIC_IDS_REDIS_KEY, modelExperimentRedisKey } from '@/lib/redis-keys'; +import { createCachedFetch } from '@/lib/cached-fetch'; 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'; @@ -24,77 +25,36 @@ import type { SerializedCachedExperiment, } from '@/lib/ai-gateway/experiments/pick-variant.types'; -/** TTL for the per-public-id experiment cache. The cache holds CIPHERTEXT - * (encrypted api keys + the upstream blob); decryption happens per-pick - * inside `pickModelExperimentVariant`, so this TTL only governs how stale - * the variant set / current-version-per-variant resolution can be, not - * how long plaintext keys can live in Redis. */ -const PER_PUBLIC_ID_CACHE_TTL_SECONDS = 600; +/** Shared Redis cache TTL. Cached experiment rows hold ciphertext only; + * plaintext keys are decrypted per-pick and never stored in Redis. */ +const EXPERIMENT_REDIS_CACHE_TTL_SECONDS = 600; +const EXPERIMENTED_PUBLIC_IDS_LOCAL_CACHE_TTL_MS = 60_000; + +const getExperimentedPublicIds = createCachedFetch( + 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 JSON-encoded array in Redis. On Redis error, falls back to a - * single Postgres query for that public id. If both fail, returns - * `unavailable` so the caller can fail closed instead of routing as - * unexperimented traffic. + * Reads from a short-lived in-process cache backed by Redis. Admin writes are + * responsible for maintaining the Redis membership set; if Redis is empty, + * corrupt, or unavailable, treat it as "no experimented public ids". */ -export async function isPublicIdExperimented( - publicId: string -): Promise { - try { - const cached = await redisGet(EXPERIMENTED_PUBLIC_IDS_REDIS_KEY); - if (cached !== null) { - const ids = parseStringArray(cached); - if (ids === null) { - // Corrupt cache; fall through to DB. (Already captured by Sentry below.) - } else { - return ids.includes(publicId); - } - } - } catch { - // captureException already invoked by the redis helper - } - - // Cache miss / corrupt / Redis error: rebuild the full set from DB and - // backfill the cache so the next request hits Redis. We deliberately - // query the full set rather than just `publicId` — the cache contract - // is "the complete list of routing-eligible public ids", and a single-id - // lookup can't satisfy it. Falling through without a backfill (the - // previous behavior) sent every miss straight to Postgres. - try { - const ids = await loadExperimentedPublicIdsFromDb(); - // Best-effort backfill; if Redis is genuinely down the helper has - // already reported it and we still return the right answer below. - void backfillExperimentedPublicIdsCache(ids); - return ids.includes(publicId); - } catch (err) { - captureException(err, { - tags: { source: 'model-experiments', operation: 'isPublicIdExperimented' }, - extra: { publicId }, - }); - return 'unavailable'; - } -} - -async function loadExperimentedPublicIdsFromDb(): Promise { - const rows = await db - .select({ public_model_id: model_experiment.public_model_id }) - .from(model_experiment) - .where(inArray(model_experiment.status, ['active', 'paused'])); - return Array.from(new Set(rows.map(r => r.public_model_id))).sort(); -} - -async function backfillExperimentedPublicIdsCache(ids: string[]): Promise { - try { - await redisSet( - EXPERIMENTED_PUBLIC_IDS_REDIS_KEY, - JSON.stringify(ids), - PER_PUBLIC_ID_CACHE_TTL_SECONDS - ); - } catch { - // captureException already invoked by the redis helper - } +export async function isPublicIdExperimented(publicId: string): Promise { + const ids = await getExperimentedPublicIds(); + return ids.includes(publicId); } function parseStringArray(raw: string): string[] | null { @@ -153,10 +113,14 @@ export async function getRoutingExperimentForPublicId(publicId: string): Promise await redisSet( cacheKey, JSON.stringify(serializeCachedExperiment(resolved.experiment)), - PER_PUBLIC_ID_CACHE_TTL_SECONDS + EXPERIMENT_REDIS_CACHE_TTL_SECONDS ); } else if (resolved.kind === 'none') { - await redisSet(cacheKey, JSON.stringify({ kind: 'none' }), PER_PUBLIC_ID_CACHE_TTL_SECONDS); + await redisSet( + cacheKey, + JSON.stringify({ kind: 'none' }), + EXPERIMENT_REDIS_CACHE_TTL_SECONDS + ); } } catch { // already captured 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 15fb31cb05b..dcf35ecaf87 100644 --- a/apps/web/src/lib/ai-gateway/providers/get-provider.ts +++ b/apps/web/src/lib/ai-gateway/providers/get-provider.ts @@ -198,9 +198,6 @@ export async function getProvider(input: GetProviderInput): Promise Date: Thu, 21 May 2026 10:13:22 +0200 Subject: [PATCH 37/49] refactor(gateway): load experiment routing from Postgres --- .../experiments/build-direct-provider.ts | 9 +- .../experiments/pick-variant.test.ts | 29 +++--- .../ai-gateway/experiments/pick-variant.ts | 93 +++---------------- .../experiments/pick-variant.types.ts | 27 +----- apps/web/src/lib/redis-keys.ts | 13 +-- .../routers/admin/model-experiments-router.ts | 30 +++--- 6 files changed, 47 insertions(+), 154 deletions(-) diff --git a/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts b/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts index b0836b52fe7..3f794397360 100644 --- a/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts +++ b/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts @@ -38,11 +38,10 @@ export function inferSupportedChatApis( * Plain in-memory shape: an `ExperimentUpstream` (no key) merged with the * decrypted partner-issued api key. * - * The cache loader (`getRoutingExperimentForPublicId`) decrypts - * `model_experiment_variant_version.encrypted_api_key` once and stores the - * resulting plaintext alongside the rest of the upstream blob in this shape - * for hot-path use. The plaintext NEVER touches Postgres or any tRPC - * response — only this in-memory record and the per-public-id Redis cache. + * `pickModelExperimentVariant` decrypts the chosen + * `model_experiment_variant_version.encrypted_api_key` and merges the + * plaintext with the upstream blob for the outbound provider request. The + * plaintext NEVER touches Postgres, Redis, or any tRPC response. */ export type ResolvedExperimentUpstream = ExperimentUpstream & { api_key: string }; 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 index 4e3eaa2ab2f..93665a860f4 100644 --- a/apps/web/src/lib/ai-gateway/experiments/pick-variant.test.ts +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.test.ts @@ -5,8 +5,8 @@ import { createCallerForUser } from '@/routers/test-utils'; import { model_experiment, model_experiment_variant_version } from '@kilocode/db/schema'; import { eq } from 'drizzle-orm'; import { isPublicIdExperimented, pickModelExperimentVariant } from './pick-variant'; -import { redisDel } from '@/lib/redis'; -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 type { User } from '@kilocode/db/schema'; let admin: User; @@ -28,11 +28,14 @@ beforeEach(async () => { }); }); -async function clearRoutingCaches(publicId: string) { - // Tests share the dev Redis instance across runs; flush both keys so - // each test sees a fresh load path. +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); - await redisDel(modelExperimentRedisKey(publicId)); +} + +async function seedExperimentedPublicIds(ids: string[]): Promise { + return await redisSet(EXPERIMENTED_PUBLIC_IDS_REDIS_KEY, JSON.stringify(ids)); } afterEach(async () => { @@ -71,18 +74,18 @@ async function makeActiveExperiment(opts: { apiKey: apiKeys[1], }); await caller.admin.modelExperiments.activate({ id: exp.id }); - await clearRoutingCaches(publicId); return { experimentId: exp.id, variantA: a.id, variantB: b.id }; } describe('isPublicIdExperimented', () => { it('returns false for an unknown public id', async () => { - await clearRoutingCaches('kilo/preview-not-experimented'); + await clearRoutingCaches(); expect(await isPublicIdExperimented('kilo/preview-not-experimented')).toBe(false); }); it('returns true when the public id has an active experiment', async () => { await makeActiveExperiment({ publicId: 'kilo/preview-iset-active' }); + if (!(await seedExperimentedPublicIds(['kilo/preview-iset-active']))) return; expect(await isPublicIdExperimented('kilo/preview-iset-active')).toBe(true); }); @@ -90,14 +93,14 @@ describe('isPublicIdExperimented', () => { const { experimentId } = await makeActiveExperiment({ publicId: 'kilo/preview-iset-paused' }); const caller = await createCallerForUser(admin.id); await caller.admin.modelExperiments.pause({ id: experimentId }); - await clearRoutingCaches('kilo/preview-iset-paused'); + if (!(await seedExperimentedPublicIds(['kilo/preview-iset-paused']))) return; 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('kilo/preview-pick-none'); + await clearRoutingCaches(); const result = await pickModelExperimentVariant({ publicModelId: 'kilo/preview-pick-none', userId: 'user-1', @@ -191,7 +194,7 @@ describe('pickModelExperimentVariant', () => { const { experimentId } = await makeActiveExperiment({ publicId: 'kilo/preview-paused' }); const caller = await createCallerForUser(admin.id); await caller.admin.modelExperiments.pause({ id: experimentId }); - await clearRoutingCaches('kilo/preview-paused'); + await clearRoutingCaches(); const result = await pickModelExperimentVariant({ publicModelId: 'kilo/preview-paused', userId: 'user-q', @@ -221,7 +224,7 @@ describe('pickModelExperimentVariant', () => { upstream: { ...upstreamA, internal_id: 'partner-checkpoint-rc-next' }, apiKey: 'sk-rc-next', }); - await clearRoutingCaches('kilo/preview-hotswap'); + await clearRoutingCaches(); const after = await pickModelExperimentVariant({ publicModelId: 'kilo/preview-hotswap', @@ -307,7 +310,7 @@ describe('pickModelExperimentVariant', () => { }); const caller = await createCallerForUser(admin.id); await caller.admin.modelExperiments.complete({ id: experimentId }); - await clearRoutingCaches('kilo/preview-completed'); + await clearRoutingCaches(); const result = await pickModelExperimentVariant({ publicModelId: 'kilo/preview-completed', userId: 'user-c', diff --git a/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts index c77a665df02..58e6245b48f 100644 --- a/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts @@ -6,8 +6,8 @@ import { model_experiment_variant_version, } from '@kilocode/db/schema'; import { db } from '@/lib/drizzle'; -import { redisGet, redisSet } from '@/lib/redis'; -import { EXPERIMENTED_PUBLIC_IDS_REDIS_KEY, modelExperimentRedisKey } from '@/lib/redis-keys'; +import { redisGet } from '@/lib/redis'; +import { EXPERIMENTED_PUBLIC_IDS_REDIS_KEY } from '@/lib/redis-keys'; import { createCachedFetch } from '@/lib/cached-fetch'; import { decryptApiKey, type EncryptedData } from '@/lib/ai-gateway/byok/encryption'; import { BYOK_ENCRYPTION_KEY } from '@/lib/config.server'; @@ -15,20 +15,14 @@ import { getRandomNumber } from '@/lib/ai-gateway/getRandomNumber'; import { ExperimentUpstreamSchema } from '@/lib/ai-gateway/experiments/upstream-schema'; import type { AllocationSubject, - CachedExperiment, - CachedVariant, ExperimentStatus, PickVariantInput, PickVariantResult, ResolveResult, - SerializedCacheEntry, - SerializedCachedExperiment, + RoutingVariant, } from '@/lib/ai-gateway/experiments/pick-variant.types'; -/** Shared Redis cache TTL. Cached experiment rows hold ciphertext only; - * plaintext keys are decrypted per-pick and never stored in Redis. */ -const EXPERIMENT_REDIS_CACHE_TTL_SECONDS = 600; -const EXPERIMENTED_PUBLIC_IDS_LOCAL_CACHE_TTL_MS = 60_000; +const EXPERIMENTED_PUBLIC_IDS_LOCAL_CACHE_TTL_MS = process.env.NODE_ENV === 'test' ? 0 : 60_000; const getExperimentedPublicIds = createCachedFetch( async () => { @@ -71,34 +65,13 @@ function parseStringArray(raw: string): string[] | null { /** * Returns the routing-relevant experiment for `publicId` (status active or * paused) with all variants resolved to their current - * `model_experiment_variant_version`. The cached payload holds CIPHERTEXT - * for partner-issued api keys; decryption is deferred to - * `pickModelExperimentVariant`. Returns `none` only after Postgres has - * proven there is no routing-relevant experiment for this id. - * - * Cached per-publicId for ~10 minutes. The cache is invalidated by every - * admin mutation that affects routing. + * `model_experiment_variant_version`. This only runs after the Redis-backed + * membership pre-check 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 { - const cacheKey = modelExperimentRedisKey(publicId); - try { - const cached = await redisGet(cacheKey); - if (cached !== null) { - const parsed = tryParseCachedExperiment(cached); - if (parsed === 'none' || parsed === null) { - // null = corrupt; fall through to DB rebuild. - if (parsed === 'none') return { kind: 'none' }; - } else { - return { kind: 'experiment', experiment: parsed }; - } - } - } catch { - // captured by redis helper; fall through to DB - } - - let resolved: ResolveResult; try { - resolved = await loadExperimentFromDb(publicId); + return await loadExperimentFromDb(publicId); } catch (err) { captureException(err, { tags: { source: 'model-experiments', operation: 'getRoutingExperimentForPublicId' }, @@ -106,49 +79,6 @@ export async function getRoutingExperimentForPublicId(publicId: string): Promise }); return { kind: 'unavailable' }; } - - // Best-effort cache write. A failure here does not affect routing. - try { - if (resolved.kind === 'experiment') { - await redisSet( - cacheKey, - JSON.stringify(serializeCachedExperiment(resolved.experiment)), - EXPERIMENT_REDIS_CACHE_TTL_SECONDS - ); - } else if (resolved.kind === 'none') { - await redisSet( - cacheKey, - JSON.stringify({ kind: 'none' }), - EXPERIMENT_REDIS_CACHE_TTL_SECONDS - ); - } - } catch { - // already captured - } - - return resolved; -} - -function serializeCachedExperiment(exp: CachedExperiment): SerializedCachedExperiment { - return { kind: 'experiment', ...exp }; -} - -function tryParseCachedExperiment(raw: string): CachedExperiment | 'none' | null { - try { - const parsed = JSON.parse(raw) as SerializedCacheEntry; - if (parsed.kind === 'none') return 'none'; - if (parsed.kind === 'experiment') { - return { - experimentId: parsed.experimentId, - publicModelId: parsed.publicModelId, - status: parsed.status, - variants: parsed.variants, - }; - } - return null; - } catch { - return null; - } } async function loadExperimentFromDb(publicId: string): Promise { @@ -242,10 +172,9 @@ async function loadExperimentFromDb(publicId: string): Promise { } } - // We deliberately DO NOT call decryptApiKey here. The cache should only - // hold encrypted blobs; decryption happens per-pick so key rotation takes - // effect on the next request. - const variants: CachedVariant[] = []; + // 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) { 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 index da02327db1e..1e4c669d37d 100644 --- a/apps/web/src/lib/ai-gateway/experiments/pick-variant.types.ts +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.types.ts @@ -5,16 +5,7 @@ import type { ExperimentUpstream } from '@/lib/ai-gateway/experiments/upstream-s export type ExperimentStatus = 'active' | 'paused'; -/** - * Per-variant payload cached in Redis. - * - * The api key is stored ENCRYPTED — the cache never holds plaintext. - * Decryption happens once per request, only on the chosen variant, inside - * `pickModelExperimentVariant`. This keeps `BYOK_ENCRYPTION_KEY` rotation - * effective immediately (no 10-minute "plaintext lives in Redis" window) - * and limits the blast radius of a Redis dump. - */ -export type CachedVariant = { +export type RoutingVariant = { variantId: string; weight: number; variantVersionId: string; @@ -22,15 +13,15 @@ export type CachedVariant = { encryptedApiKey: EncryptedData; }; -export type CachedExperiment = { +export type RoutingExperiment = { experimentId: string; publicModelId: string; status: ExperimentStatus; - variants: CachedVariant[]; + variants: RoutingVariant[]; }; export type ResolveResult = - | { kind: 'experiment'; experiment: CachedExperiment } + | { kind: 'experiment'; experiment: RoutingExperiment } | { kind: 'none' } | { kind: 'unavailable' }; @@ -55,14 +46,4 @@ export type PickVariantResult = | { status: 'not-found' } | { status: 'unavailable' }; -export type SerializedCachedExperiment = { - kind: 'experiment'; - experimentId: string; - publicModelId: string; - status: ExperimentStatus; - variants: CachedVariant[]; -}; - -export type SerializedCacheEntry = SerializedCachedExperiment | { kind: 'none' }; - export type { ModelExperiment }; diff --git a/apps/web/src/lib/redis-keys.ts b/apps/web/src/lib/redis-keys.ts index 65e72a22b55..0f770cde1e0 100644 --- a/apps/web/src/lib/redis-keys.ts +++ b/apps/web/src/lib/redis-keys.ts @@ -39,7 +39,7 @@ export const botIdentityRedisKey = (platform: string, teamId: string, userId: st /** * Set of public_model_ids that have a routing-relevant model_experiment row * (status IN 'active' | 'paused'). Used by `getProvider` as a fast pre-check - * before fetching the per-public-id experiment payload. + * before loading the experiment routing payload from Postgres. * * Stored as a JSON array string. Recomputed and rewritten on every status * transition into or out of (active, paused). @@ -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 56e51c577ae..b4c6240fa34 100644 --- a/apps/web/src/routers/admin/model-experiments-router.ts +++ b/apps/web/src/routers/admin/model-experiments-router.ts @@ -8,8 +8,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'; @@ -109,7 +109,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)) { @@ -122,15 +122,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 { @@ -141,7 +134,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, @@ -322,11 +315,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; }), @@ -470,7 +462,7 @@ export const adminModelExperimentsRouter = createTRPCRouter({ created_by: ctx.user.id, }) .returning(variantVersionPublicColumns); - await invalidateExperimentCaches(experiment.public_model_id); + await refreshExperimentedPublicIdsCache(); return inserted; }), @@ -511,7 +503,7 @@ export const adminModelExperimentsRouter = createTRPCRouter({ created_by: ctx.user.id, }) .returning(variantVersionPublicColumns); - await invalidateExperimentCaches(experiment.public_model_id); + await refreshExperimentedPublicIdsCache(); return inserted; }), }); From c6c5e087ba91197def4cf99104bf81533c027b9d Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Thu, 21 May 2026 10:45:32 +0200 Subject: [PATCH 38/49] test(gateway): avoid silent Redis membership skips --- .../lib/ai-gateway/experiments/pick-variant.test.ts | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) 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 index 93665a860f4..79b90a60727 100644 --- a/apps/web/src/lib/ai-gateway/experiments/pick-variant.test.ts +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.test.ts @@ -19,6 +19,7 @@ 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(); @@ -83,17 +84,17 @@ describe('isPublicIdExperimented', () => { expect(await isPublicIdExperimented('kilo/preview-not-experimented')).toBe(false); }); - it('returns true when the public id has an active experiment', async () => { + redisIt('returns true when the public id has an active experiment', async () => { await makeActiveExperiment({ publicId: 'kilo/preview-iset-active' }); - if (!(await seedExperimentedPublicIds(['kilo/preview-iset-active']))) return; + expect(await seedExperimentedPublicIds(['kilo/preview-iset-active'])).toBe(true); expect(await isPublicIdExperimented('kilo/preview-iset-active')).toBe(true); }); - it('returns true when the public id has only a paused experiment', async () => { + 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 }); - if (!(await seedExperimentedPublicIds(['kilo/preview-iset-paused']))) return; + expect(await seedExperimentedPublicIds(['kilo/preview-iset-paused'])).toBe(true); expect(await isPublicIdExperimented('kilo/preview-iset-paused')).toBe(true); }); }); @@ -265,7 +266,7 @@ describe('pickModelExperimentVariant', () => { // Loose bounds: 1:3 ≈ 25/75. Allow ±10pp for n=200. const controlPct = counts.control / 200; expect(controlPct).toBeGreaterThan(0.15); - expect(controlPct).toBeLessThan(0.35); + expect(controlPct).toBeLessThanOrEqual(0.35); }); it('historical attribution survives hot-swap: old variant_version_id still resolves to old upstream via DB', async () => { From 705bafa6d6081dee763d7668e81af98870f1a54c Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Thu, 21 May 2026 12:09:47 +0200 Subject: [PATCH 39/49] feat(gateway): expose active experiment models --- .../src/app/api/openrouter/models/route.ts | 6 +- .../list-available-experiment-models.test.ts | 68 +++++++++++++++++++ .../list-available-experiment-models.ts | 52 ++++++++++++++ .../lib/organizations/organization-models.ts | 2 + 4 files changed, 127 insertions(+), 1 deletion(-) create mode 100644 apps/web/src/lib/ai-gateway/experiments/list-available-experiment-models.test.ts create mode 100644 apps/web/src/lib/ai-gateway/experiments/list-available-experiment-models.ts diff --git a/apps/web/src/app/api/openrouter/models/route.ts b/apps/web/src/app/api/openrouter/models/route.ts index c9af1304bba..e09714ddd54 100644 --- a/apps/web/src/app/api/openrouter/models/route.ts +++ b/apps/web/src/app/api/openrouter/models/route.ts @@ -8,6 +8,7 @@ import { getDirectByokModelsForUser } from '@/lib/ai-gateway/providers/direct-by import { getAvailableModelsForOrganization } from '@/lib/organizations/organization-models'; import { FEATURE_HEADER, validateFeatureHeader } from '@/lib/feature-detection'; import { filterByFeature } from '@/lib/ai-gateway/models'; +import { listAvailableExperimentModels } from '@/lib/ai-gateway/experiments/list-available-experiment-models'; async function tryGetUserFromAuth() { try { @@ -43,7 +44,10 @@ export async function GET( return NextResponse.json(data); } const byokModels = auth?.user ? await getDirectByokModelsForUser(auth.user.id) : []; - return NextResponse.json({ data: filterByFeature(data.data.concat(byokModels), feature) }); + const experimentModels = await listAvailableExperimentModels(); + return NextResponse.json({ + data: filterByFeature(data.data.concat(byokModels, experimentModels), feature), + }); } catch (error) { captureException(error, { tags: { endpoint: 'openrouter/models' }, diff --git a/apps/web/src/lib/ai-gateway/experiments/list-available-experiment-models.test.ts b/apps/web/src/lib/ai-gateway/experiments/list-available-experiment-models.test.ts new file mode 100644 index 00000000000..6c39570dfbd --- /dev/null +++ b/apps/web/src/lib/ai-gateway/experiments/list-available-experiment-models.test.ts @@ -0,0 +1,68 @@ +import { afterEach, describe, expect, it } from '@jest/globals'; +import { db } from '@/lib/drizzle'; +import { model_experiment } from '@kilocode/db/schema'; +import { listAvailableExperimentModels } from './list-available-experiment-models'; +import { inArray } from 'drizzle-orm'; + +const testPublicIds = [ + 'kilo/preview-active-a', + 'kilo/preview-active-b', + 'kilo/preview-paused', + 'kilo/preview-draft', +]; + +describe('listAvailableExperimentModels', () => { + afterEach(async () => { + await db + .delete(model_experiment) + .where(inArray(model_experiment.public_model_id, testPublicIds)); + }); + + it('returns only active experiment public ids as selectable models', async () => { + await db + .delete(model_experiment) + .where(inArray(model_experiment.public_model_id, testPublicIds)); + await db.insert(model_experiment).values([ + { + public_model_id: 'kilo/preview-active-a', + name: 'Preview Active A', + description: 'First active preview', + status: 'active', + }, + { + public_model_id: 'kilo/preview-active-b', + name: 'Preview Active B', + description: null, + status: 'active', + }, + { + public_model_id: 'kilo/preview-paused', + name: 'Preview Paused', + description: 'Paused preview', + status: 'paused', + }, + { + public_model_id: 'kilo/preview-draft', + name: 'Preview Draft', + description: 'Draft preview', + status: 'draft', + }, + ]); + + const models = await listAvailableExperimentModels(); + + expect(models.map(model => model.id)).toEqual([ + 'kilo/preview-active-a', + 'kilo/preview-active-b', + ]); + expect(models[0]).toEqual( + expect.objectContaining({ + id: 'kilo/preview-active-a', + name: 'Preview Active A', + description: 'First active preview', + context_length: 200000, + }) + ); + expect(models[1].description).toBe('Preview Active B'); + }); +}); diff --git a/apps/web/src/lib/ai-gateway/experiments/list-available-experiment-models.ts b/apps/web/src/lib/ai-gateway/experiments/list-available-experiment-models.ts new file mode 100644 index 00000000000..00f11530268 --- /dev/null +++ b/apps/web/src/lib/ai-gateway/experiments/list-available-experiment-models.ts @@ -0,0 +1,52 @@ +import { asc, eq } from 'drizzle-orm'; +import { model_experiment } from '@kilocode/db/schema'; +import { readDb } from '@/lib/drizzle'; +import type { OpenRouterModel } from '@/lib/organizations/organization-types'; + +const DEFAULT_CONTEXT_LENGTH = 200_000; + +function toCreatedSeconds(createdAt: string): number { + const ms = new Date(createdAt).getTime(); + return Number.isFinite(ms) ? Math.floor(ms / 1000) : 0; +} + +function convertExperimentModel(row: typeof model_experiment.$inferSelect): OpenRouterModel { + return { + id: row.public_model_id, + name: row.name, + created: toCreatedSeconds(row.created_at), + description: row.description ?? row.name, + architecture: { + input_modalities: ['text'], + output_modalities: ['text'], + tokenizer: 'Other', + }, + top_provider: { + is_moderated: false, + context_length: DEFAULT_CONTEXT_LENGTH, + max_completion_tokens: null, + }, + pricing: { + prompt: '0.0000000', + completion: '0.0000000', + request: '0', + image: '0', + web_search: '0', + internal_reasoning: '0', + input_cache_read: '0.00000000', + input_cache_write: '0.00000000', + }, + context_length: DEFAULT_CONTEXT_LENGTH, + per_request_limits: null, + supported_parameters: ['max_tokens', 'temperature', 'tools', 'reasoning', 'include_reasoning'], + }; +} + +export async function listAvailableExperimentModels(): Promise { + const rows = await readDb + .select() + .from(model_experiment) + .where(eq(model_experiment.status, 'active')) + .orderBy(asc(model_experiment.public_model_id)); + return rows.map(convertExperimentModel); +} diff --git a/apps/web/src/lib/organizations/organization-models.ts b/apps/web/src/lib/organizations/organization-models.ts index 749bc6fe27f..6e04f65cfae 100644 --- a/apps/web/src/lib/organizations/organization-models.ts +++ b/apps/web/src/lib/organizations/organization-models.ts @@ -9,6 +9,7 @@ import { listAvailableCustomLlms } from '@/lib/ai-gateway/custom-llm/listAvailab import { getDirectByokModelsForOrganization } from '@/lib/ai-gateway/providers/direct-byok'; import { getOrganizationById } from '@/lib/organizations/organizations'; import { getEffectiveModelRestrictions } from '@/lib/organizations/model-restrictions'; +import { listAvailableExperimentModels } from '@/lib/ai-gateway/experiments/list-available-experiment-models'; export async function getAvailableModelsForOrganization( organizationId: string @@ -40,6 +41,7 @@ export async function getAvailableModelsForOrganization( filteredModels.push(...(await getDirectByokModelsForOrganization(organizationId))); filteredModels.push(...(await listAvailableCustomLlms(organizationId))); + filteredModels.push(...(await listAvailableExperimentModels())); return { ...responseData, From 8a6bfc6b8538e644e2df8d3cbc85f846c56945d5 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Thu, 21 May 2026 12:49:00 +0200 Subject: [PATCH 40/49] feat(admin): show experiment request log --- .../app/admin/api/model-experiments/hooks.ts | 5 + apps/web/src/app/admin/gateway/page.tsx | 15 +- .../ModelExperimentRequestsContent.tsx | 144 ++++++++++++++++++ .../admin/model-experiments-router.test.ts | 64 ++++++++ .../routers/admin/model-experiments-router.ts | 49 ++++++ 5 files changed, 276 insertions(+), 1 deletion(-) create mode 100644 apps/web/src/app/admin/model-experiments/ModelExperimentRequestsContent.tsx diff --git a/apps/web/src/app/admin/api/model-experiments/hooks.ts b/apps/web/src/app/admin/api/model-experiments/hooks.ts index f3b132bf65a..984bbf77f6f 100644 --- a/apps/web/src/app/admin/api/model-experiments/hooks.ts +++ b/apps/web/src/app/admin/api/model-experiments/hooks.ts @@ -16,6 +16,11 @@ export function useModelExperiment(id: string | null) { }); } +export function useModelExperimentRequests() { + const trpc = useTRPC(); + return useQuery(trpc.admin.modelExperiments.listRequests.queryOptions()); +} + function useInvalidate() { const trpc = useTRPC(); const qc = useQueryClient(); diff --git a/apps/web/src/app/admin/gateway/page.tsx b/apps/web/src/app/admin/gateway/page.tsx index e3b6d49d792..622d1eff0fe 100644 --- a/apps/web/src/app/admin/gateway/page.tsx +++ b/apps/web/src/app/admin/gateway/page.tsx @@ -9,14 +9,21 @@ import { SyncProvidersContent } from '@/app/admin/sync-providers/SyncProvidersCo import { CustomLlmsContent } from '@/app/admin/custom-llms/CustomLlmsContent'; import { RoutingContent } from '@/app/admin/gateway/RoutingContent'; import { ModelExperimentsContent } from '@/app/admin/model-experiments/ModelExperimentsContent'; +import { ModelExperimentRequestsContent } from '@/app/admin/model-experiments/ModelExperimentRequestsContent'; const VALID_TABS: readonly string[] = [ 'sync-providers', 'custom-llms', 'routing', 'model-experiments', + 'experiment-requests', ]; -type Tab = 'sync-providers' | 'custom-llms' | 'routing' | 'model-experiments'; +type Tab = + | 'sync-providers' + | 'custom-llms' + | 'routing' + | 'model-experiments' + | 'experiment-requests'; const isValidTab = (value: string | null): value is Tab => value !== null && VALID_TABS.includes(value); @@ -76,6 +83,9 @@ export default function AdminGatewayPage() { Model Experiments + + Experiment Requests + @@ -89,6 +99,9 @@ export default function AdminGatewayPage() { + + + diff --git a/apps/web/src/app/admin/model-experiments/ModelExperimentRequestsContent.tsx b/apps/web/src/app/admin/model-experiments/ModelExperimentRequestsContent.tsx new file mode 100644 index 00000000000..f7f573f1ac5 --- /dev/null +++ b/apps/web/src/app/admin/model-experiments/ModelExperimentRequestsContent.tsx @@ -0,0 +1,144 @@ +'use client'; + +import { Badge } from '@/components/ui/badge'; +import { + Table, + TableBody, + TableCell, + TableHead, + TableHeader, + TableRow, +} from '@/components/ui/table'; +import { useModelExperimentRequests } from '@/app/admin/api/model-experiments/hooks'; + +function formatDate(value: string): string { + return new Date(value).toLocaleString(); +} + +function formatNullable(value: string | null): string { + return value && value.length > 0 ? value : '—'; +} + +function formatTokens(value: number): string { + return value.toLocaleString(); +} + +function formatMicrodollars(value: number | null): string { + if (value === null) return '—'; + return value.toLocaleString(); +} + +function HashCell({ value }: { value: string }) { + const display = value.length > 18 ? `${value.slice(0, 18)}…` : value; + return ( + + {display} + + ); +} + +export function ModelExperimentRequestsContent() { + const { data, isLoading } = useModelExperimentRequests(); + const rows = data?.items ?? []; + + return ( +
+
+

Experiment Requests

+

+ Last 100 requests attributed to model experiments. Prompt bodies stay in R2; this table + shows attribution, routing, and usage metadata only. +

+
+ + {isLoading ? ( +
Loading…
+ ) : rows.length === 0 ? ( +
+ No experiment requests recorded yet. +
+ ) : ( +
+ + + + Created + Experiment + Variant + Request + Usage + Prompt hash + User / org + + + + {rows.map(row => ( + + + {formatDate(row.createdAt)} + + +
{row.experimentName}
+
+ {row.publicModelId} +
+
+ {row.experimentId} +
+ + +
{row.variantLabel}
+
+ {row.variantVersionId} +
+ + +
+ {row.requestKind} + {row.allocationSubject} + {row.wasTruncated && truncated} +
+
+ Client: {formatNullable(row.clientRequestId)} +
+
{row.usageId}
+ + +
+ Tokens: {formatTokens(row.inputTokens)} in / {formatTokens(row.outputTokens)}{' '} + out +
+
+ Cache: {formatTokens(row.cacheWriteTokens)} write /{' '} + {formatTokens(row.cacheHitTokens)} hit +
+
+ Cost: {formatMicrodollars(row.costMicrodollars)} µ$ · discount:{' '} + {formatMicrodollars(row.cacheDiscountMicrodollars)} µ$ +
+ {row.hasError && error} + + + + + +
{row.userId}
+
+ {formatNullable(row.organizationId)} +
+
+ {formatNullable(row.provider)} / {formatNullable(row.inferenceProvider)} +
+
+ {formatNullable(row.requestedModel)} → {formatNullable(row.upstreamModel)} +
+ + + ))} + +
+
+ )} +
+ ); +} diff --git a/apps/web/src/routers/admin/model-experiments-router.test.ts b/apps/web/src/routers/admin/model-experiments-router.test.ts index 960cfe0a793..14367c9213f 100644 --- a/apps/web/src/routers/admin/model-experiments-router.test.ts +++ b/apps/web/src/routers/admin/model-experiments-router.test.ts @@ -3,7 +3,9 @@ import { cleanupDbForTest, db } from '@/lib/drizzle'; import { insertTestUser } from '@/tests/helpers/user.helper'; import { createCallerForUser } from '@/routers/test-utils'; import { + microdollar_usage, model_experiment, + model_experiment_request, model_experiment_variant, model_experiment_variant_version, } from '@kilocode/db/schema'; @@ -11,6 +13,7 @@ import { decryptApiKey } from '@/lib/ai-gateway/byok/encryption'; import { BYOK_ENCRYPTION_KEY } from '@/lib/config.server'; import { eq } from 'drizzle-orm'; import type { User } from '@kilocode/db/schema'; +import { randomUUID } from 'crypto'; let admin: User; @@ -93,6 +96,67 @@ describe('admin.modelExperiments — basic CRUD', () => { expect(withArchived.items.map(i => i.id).sort()).toEqual([a.id, b.id].sort()); }); + it('lists the last experiment attribution requests with usage and variant metadata', async () => { + const { caller, experimentId, variantA } = + await makeDraftWithTwoVariants('kilo/preview-requests'); + await caller.admin.modelExperiments.activate({ id: experimentId }); + const [version] = await db + .select({ id: model_experiment_variant_version.id }) + .from(model_experiment_variant_version) + .where(eq(model_experiment_variant_version.variant_id, variantA.id)) + .limit(1); + expect(version).toBeDefined(); + + const usageId = randomUUID(); + await db.insert(microdollar_usage).values({ + id: usageId, + kilo_user_id: 'request-user', + cost: 0, + input_tokens: 123, + output_tokens: 45, + cache_write_tokens: 6, + cache_hit_tokens: 7, + provider: 'custom', + model: 'partner-checkpoint-rc1', + requested_model: 'kilo/preview-requests', + inference_provider: 'partner', + organization_id: null, + }); + await db.insert(model_experiment_request).values({ + usage_id: usageId, + variant_version_id: version.id, + allocation_subject: 'user', + client_request_id: 'client-request-123', + request_kind: 'chat_completions', + request_body_sha256: '__failed__', + was_truncated: true, + }); + + const requests = await caller.admin.modelExperiments.listRequests(); + + expect(requests.items[0]).toEqual( + expect.objectContaining({ + usageId, + experimentId, + experimentName: 'exp-kilo/preview-requests', + publicModelId: 'kilo/preview-requests', + variantId: variantA.id, + variantLabel: 'control', + variantVersionId: version.id, + allocationSubject: 'user', + clientRequestId: 'client-request-123', + requestKind: 'chat_completions', + requestBodySha256: '__failed__', + wasTruncated: true, + userId: 'request-user', + requestedModel: 'kilo/preview-requests', + upstreamModel: 'partner-checkpoint-rc1', + inputTokens: 123, + outputTokens: 45, + }) + ); + }); + it('rejects delete unless status is draft', async () => { const { caller, experimentId } = await makeDraftWithTwoVariants('kilo/preview-delete'); await caller.admin.modelExperiments.activate({ id: experimentId }); diff --git a/apps/web/src/routers/admin/model-experiments-router.ts b/apps/web/src/routers/admin/model-experiments-router.ts index b4c6240fa34..8d027674b42 100644 --- a/apps/web/src/routers/admin/model-experiments-router.ts +++ b/apps/web/src/routers/admin/model-experiments-router.ts @@ -1,7 +1,9 @@ import { adminProcedure, createTRPCRouter } from '@/lib/trpc/init'; import { db } from '@/lib/drizzle'; import { + microdollar_usage, model_experiment, + model_experiment_request, model_experiment_variant, model_experiment_variant_version, } from '@kilocode/db/schema'; @@ -280,6 +282,53 @@ export const adminModelExperimentsRouter = createTRPCRouter({ return { items: rows }; }), + listRequests: adminProcedure.query(async () => { + const rows = await db + .select({ + usageId: model_experiment_request.usage_id, + createdAt: model_experiment_request.created_at, + experimentId: model_experiment.id, + experimentName: model_experiment.name, + publicModelId: model_experiment.public_model_id, + variantId: model_experiment_variant.id, + variantLabel: model_experiment_variant.label, + variantVersionId: model_experiment_variant_version.id, + allocationSubject: model_experiment_request.allocation_subject, + clientRequestId: model_experiment_request.client_request_id, + requestKind: model_experiment_request.request_kind, + requestBodySha256: model_experiment_request.request_body_sha256, + wasTruncated: model_experiment_request.was_truncated, + userId: microdollar_usage.kilo_user_id, + organizationId: microdollar_usage.organization_id, + requestedModel: microdollar_usage.requested_model, + upstreamModel: microdollar_usage.model, + provider: microdollar_usage.provider, + inferenceProvider: microdollar_usage.inference_provider, + inputTokens: microdollar_usage.input_tokens, + outputTokens: microdollar_usage.output_tokens, + cacheWriteTokens: microdollar_usage.cache_write_tokens, + cacheHitTokens: microdollar_usage.cache_hit_tokens, + costMicrodollars: microdollar_usage.cost, + cacheDiscountMicrodollars: microdollar_usage.cache_discount, + hasError: microdollar_usage.has_error, + }) + .from(model_experiment_request) + .innerJoin(microdollar_usage, eq(model_experiment_request.usage_id, microdollar_usage.id)) + .innerJoin( + model_experiment_variant_version, + eq(model_experiment_request.variant_version_id, model_experiment_variant_version.id) + ) + .innerJoin( + model_experiment_variant, + eq(model_experiment_variant_version.variant_id, model_experiment_variant.id) + ) + .innerJoin(model_experiment, eq(model_experiment_variant.experiment_id, model_experiment.id)) + .orderBy(desc(model_experiment_request.created_at)) + .limit(100); + + return { items: rows }; + }), + get: adminProcedure.input(idSchema).query(async ({ input }) => { const experiment = await loadExperimentOrThrow(input.id); const variants = await listVariantsWithCurrentVersion(input.id); From 5d728dcacde81c2a216ae8c546f6b54b494aa927 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Thu, 21 May 2026 14:05:33 +0200 Subject: [PATCH 41/49] fix(gateway): allow anonymous experiment models --- apps/web/src/app/api/openrouter/[...path]/route.ts | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/apps/web/src/app/api/openrouter/[...path]/route.ts b/apps/web/src/app/api/openrouter/[...path]/route.ts index 17f269e2d29..f7395d1457f 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.ts @@ -19,6 +19,7 @@ import type { 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/pick-variant'; import { upstreamRequest } from '@/lib/ai-gateway/providers/upstream-request'; import { debugSaveProxyRequest } from '@/lib/debugUtils'; import { setTag, startInactiveSpan } from '@sentry/nextjs'; @@ -316,7 +317,10 @@ export async function POST(request: NextRequest): Promise Date: Thu, 21 May 2026 14:26:13 +0200 Subject: [PATCH 42/49] refactor(gateway): treat experiment models as free models Resolve isPublicIdExperimented once per request and OR the result into the existing free-model gates: rate limit, anonymous-access, and balance check. Drop the now-redundant skipBalanceCheck flag from GetProviderProviderResult. --- .../src/app/api/openrouter/[...path]/route.ts | 34 ++++++++++++------- .../lib/ai-gateway/providers/get-provider.ts | 4 --- 2 files changed, 22 insertions(+), 16 deletions(-) diff --git a/apps/web/src/app/api/openrouter/[...path]/route.ts b/apps/web/src/app/api/openrouter/[...path]/route.ts index f7395d1457f..93aaa233a29 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.ts @@ -273,12 +273,19 @@ export async function POST(request: NextRequest): Promise Date: Thu, 21 May 2026 15:50:58 +0200 Subject: [PATCH 43/49] chore(gateway): remove unrelated review changes --- .../providers/openrouter/types.test.ts | 37 ------------------- apps/web/src/lib/redis-keys.ts | 2 +- 2 files changed, 1 insertion(+), 38 deletions(-) delete mode 100644 apps/web/src/lib/ai-gateway/providers/openrouter/types.test.ts diff --git a/apps/web/src/lib/ai-gateway/providers/openrouter/types.test.ts b/apps/web/src/lib/ai-gateway/providers/openrouter/types.test.ts deleted file mode 100644 index ee7825b4234..00000000000 --- a/apps/web/src/lib/ai-gateway/providers/openrouter/types.test.ts +++ /dev/null @@ -1,37 +0,0 @@ -import { describe, expect, it } from '@jest/globals'; -import { isFreePromptTrainingAllowed } from './types'; - -// Used by both the kilo-exclusive free-model data-collection gate and the -// model-experiment routing refusal. Locks the per-request opt-out shapes -// (`provider.data_collection: 'deny'` and `provider.zdr: true`) so a -// regression in this predicate cannot silently allow prompt capture for -// callers who explicitly opted out at the request scope. -describe('isFreePromptTrainingAllowed', () => { - it('allows when the provider config is undefined', () => { - expect(isFreePromptTrainingAllowed(undefined)).toBe(true); - }); - - it('allows when no opt-out fields are set', () => { - expect(isFreePromptTrainingAllowed({})).toBe(true); - }); - - it('allows when data_collection is explicitly "allow"', () => { - expect(isFreePromptTrainingAllowed({ data_collection: 'allow' })).toBe(true); - }); - - it('refuses when data_collection is "deny"', () => { - expect(isFreePromptTrainingAllowed({ data_collection: 'deny' })).toBe(false); - }); - - it('refuses when zdr is true', () => { - expect(isFreePromptTrainingAllowed({ zdr: true })).toBe(false); - }); - - it('refuses when both opt-out fields are set', () => { - expect(isFreePromptTrainingAllowed({ data_collection: 'deny', zdr: true })).toBe(false); - }); - - it('allows when zdr is explicitly false and no other opt-out is set', () => { - expect(isFreePromptTrainingAllowed({ zdr: false })).toBe(true); - }); -}); diff --git a/apps/web/src/lib/redis-keys.ts b/apps/web/src/lib/redis-keys.ts index 0f770cde1e0..d60790697f1 100644 --- a/apps/web/src/lib/redis-keys.ts +++ b/apps/web/src/lib/redis-keys.ts @@ -39,7 +39,7 @@ export const botIdentityRedisKey = (platform: string, teamId: string, userId: st /** * Set of public_model_ids that have a routing-relevant model_experiment row * (status IN 'active' | 'paused'). Used by `getProvider` as a fast pre-check - * before loading the experiment routing payload from Postgres. + * before fetching the per-public-id experiment payload. * * Stored as a JSON array string. Recomputed and rewritten on every status * transition into or out of (active, paused). From b0e2971032ddfc7668e451275a5cbf2f48583b2b Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Thu, 21 May 2026 15:57:02 +0200 Subject: [PATCH 44/49] refactor(gateway): fold experiment check into isFreeModel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Now that isFreeModel is async (PR #3393), let it answer for experimented public ids directly. Drops the route's local isExperimentedModel boolean from the auth and balance gates and removes the provider_funded flag — cost zeroing is handled by the isFreeModel check in processUsage.ts. --- .../src/app/api/openrouter/[...path]/route.ts | 32 +++++++------------ apps/web/src/lib/ai-gateway/models.ts | 4 ++- apps/web/src/lib/ai-gateway/processUsage.ts | 6 +--- .../src/lib/ai-gateway/processUsage.types.ts | 7 ---- 4 files changed, 15 insertions(+), 34 deletions(-) diff --git a/apps/web/src/app/api/openrouter/[...path]/route.ts b/apps/web/src/app/api/openrouter/[...path]/route.ts index fa5c4b7578e..a596a26807c 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.ts @@ -273,19 +273,17 @@ export async function POST(request: NextRequest): 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/processUsage.ts b/apps/web/src/lib/ai-gateway/processUsage.ts index aff376ca61c..ea96a9ff4bb 100644 --- a/apps/web/src/lib/ai-gateway/processUsage.ts +++ b/apps/web/src/lib/ai-gateway/processUsage.ts @@ -1026,11 +1026,7 @@ export async function processTokenData( // Preserve the real cost before zeroing for free/BYOK usageStats.market_cost = usageStats.cost_mUsd; - if ( - (await isFreeModel(usageContext.requested_model)) || - usageContext.user_byok || - usageContext.provider_funded - ) { + if ((await isFreeModel(usageContext.requested_model)) || usageContext.user_byok) { usageStats.cost_mUsd = 0; usageStats.cacheDiscount_mUsd = 0; } diff --git a/apps/web/src/lib/ai-gateway/processUsage.types.ts b/apps/web/src/lib/ai-gateway/processUsage.types.ts index 69628a8d299..249a30b4466 100644 --- a/apps/web/src/lib/ai-gateway/processUsage.types.ts +++ b/apps/web/src/lib/ai-gateway/processUsage.types.ts @@ -129,13 +129,6 @@ export type MicrodollarUsageContext = { machine_id: string | null; /** True if user/org is using their own API key - cost should be zeroed out */ user_byok: boolean; - /** - * True for traffic where Kilo (or a partner) is footing the bill, so the - * user's billable cost must be zeroed out (preserving `market_cost` for - * reporting). Currently set for model-experiment routing — preview - * checkpoints are provider-funded for v1. - */ - provider_funded?: boolean; has_tools: boolean; botId?: string; tokenSource?: string; From 7098d88c2e66fed75ee0d0653084eb585e8b0106 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Thu, 21 May 2026 16:10:41 +0200 Subject: [PATCH 45/49] chore(gateway): revert unnecessary comment edits --- .../src/app/api/openrouter/[...path]/route.ts | 22 +++++++------------ .../src/lib/ai-gateway/processUsage.types.ts | 7 +++--- 2 files changed, 12 insertions(+), 17 deletions(-) diff --git a/apps/web/src/app/api/openrouter/[...path]/route.ts b/apps/web/src/app/api/openrouter/[...path]/route.ts index a596a26807c..815182a11f5 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.ts @@ -225,9 +225,8 @@ export async function POST(request: NextRequest): Promise Date: Thu, 21 May 2026 16:22:07 +0200 Subject: [PATCH 46/49] fix(gateway): refuse experiments under enterprise allow-list Enterprise provider allow-list is enforced via OpenRouter's body.provider.only field, which doesn't reach a direct partner upstream. Refuse experiment routing rather than silently bypassing the org's allow-list. Note Option 3 (unify direct-upstream routing abstraction) in the plan as the proper followup. --- .plans/experimental-models-1.md | 32 +++++++++++++++++++ .../src/app/api/openrouter/[...path]/route.ts | 9 ++++++ 2 files changed, 41 insertions(+) diff --git a/.plans/experimental-models-1.md b/.plans/experimental-models-1.md index 80369e83be8..8ee27a1ca30 100644 --- a/.plans/experimental-models-1.md +++ b/.plans/experimental-models-1.md @@ -600,6 +600,38 @@ These constraints exist because of how the gateway is built today. The spec must - 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: 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 815182a11f5..4d8070415b7 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.ts @@ -44,6 +44,7 @@ import { malformedJsonResponse, makeErrorReadable, modelDoesNotExistResponse, + modelNotAllowedResponse, extractHeaderAndLimitLength, noFreeModelsAvailableResponse, temporarilyUnavailableResponse, @@ -519,6 +520,14 @@ export async function POST(request: NextRequest): Promise Date: Thu, 21 May 2026 16:28:37 +0200 Subject: [PATCH 47/49] fix(gateway): remove duplicate unfiltered experiment models Experiment models were already added to restrictionCandidates and folded into the org-restrictions filter, then appended again unfiltered at the end. The duplicate append created two issues: catalog entries appeared twice, and Enterprise model_deny_list entries were silently re-exposed. Drop the duplicate append; the filter at line 32-42 covers experiments and matches the gateway's behavior (experiments do not set bypassAccessCheck and go through checkOrganizationModelRestrictions on the request path). --- apps/web/src/lib/organizations/organization-models.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/apps/web/src/lib/organizations/organization-models.ts b/apps/web/src/lib/organizations/organization-models.ts index 87358a776b9..b41711c30ec 100644 --- a/apps/web/src/lib/organizations/organization-models.ts +++ b/apps/web/src/lib/organizations/organization-models.ts @@ -43,7 +43,6 @@ export async function getAvailableModelsForOrganization( filteredModels.push(...(await getDirectByokModelsForOrganization(organizationId))); filteredModels.push(...(await listAvailableCustomLlms(organizationId))); - filteredModels.push(...(await listAvailableExperimentModels())); return { ...responseData, From 1eefb73edd72bcb489e48da17af56d2443a5c630 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Thu, 21 May 2026 17:03:01 +0200 Subject: [PATCH 48/49] fix(gateway): split membership and isFreeModel out of models.ts models.ts is reachable from client bundles via constants like PRIMARY_DEFAULT_MODEL and preferredModels. Folding the experiment-membership check into isFreeModel inside models.ts dragged in the Redis client (and via pick-variant.ts the Postgres client), breaking the production build with 'Module not found' errors for net/tls/dns/promises and pg. Split along two seams that already existed: pick-variant.ts now contains only the drizzle-using routing (getRoutingExperimentForPublicId, pickModelExperimentVariant), while a new membership.ts holds the Redis-only isPublicIdExperimented. The async isFreeModel itself moves to a new is-free-model.ts that imports the sync helpers from models.ts and the membership check from membership.ts. Net result: models.ts has no I/O dependencies, the fold is preserved for server callers, and the experiment-membership check still lives once in the codebase. --- apps/web/src/app/api/fim/completions/route.ts | 2 +- .../src/app/api/openrouter/[...path]/route.ts | 4 +- .../app/api/openrouter/embeddings/route.ts | 2 +- .../lib/ai-gateway/experiments/membership.ts | 51 +++++++++++++++++++ .../experiments/pick-variant.test.ts | 3 +- .../ai-gateway/experiments/pick-variant.ts | 48 ++--------------- apps/web/src/lib/ai-gateway/is-free-model.ts | 24 +++++++++ apps/web/src/lib/ai-gateway/models.test.ts | 3 +- apps/web/src/lib/ai-gateway/models.ts | 14 +---- apps/web/src/lib/ai-gateway/processUsage.ts | 3 +- .../lib/ai-gateway/providers/get-provider.ts | 2 +- .../ai-gateway/providers/openrouter/index.ts | 2 +- apps/web/src/lib/bot/agent-runner.ts | 2 +- 13 files changed, 92 insertions(+), 68 deletions(-) create mode 100644 apps/web/src/lib/ai-gateway/experiments/membership.ts create mode 100644 apps/web/src/lib/ai-gateway/is-free-model.ts diff --git a/apps/web/src/app/api/fim/completions/route.ts b/apps/web/src/app/api/fim/completions/route.ts index c0f0ccef14c..156fe908b0b 100644 --- a/apps/web/src/app/api/fim/completions/route.ts +++ b/apps/web/src/app/api/fim/completions/route.ts @@ -5,7 +5,7 @@ import z from 'zod'; import { captureException, setTag, startInactiveSpan } from '@sentry/nextjs'; import type { MicrodollarUsageContext } from '@/lib/ai-gateway/processUsage.types'; import { validateFeatureHeader, FEATURE_HEADER } from '@/lib/feature-detection'; -import { isFreeModel } from '@/lib/ai-gateway/models'; +import { isFreeModel } from '@/lib/ai-gateway/is-free-model'; import { sentryRootSpan } from '@/lib/getRootSpan'; import { getUserFromAuth } from '@/lib/user.server'; import { diff --git a/apps/web/src/app/api/openrouter/[...path]/route.ts b/apps/web/src/app/api/openrouter/[...path]/route.ts index 4d8070415b7..235a8350275 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.ts @@ -19,19 +19,19 @@ import type { 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/pick-variant'; +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'; import { getUserFromAuth } from '@/lib/user.server'; import { sentryRootSpan } from '@/lib/getRootSpan'; import { - isFreeModel, isDeadFreeModel, isExcludedForFeature, isKiloExclusiveFreeModel, isKiloStealthModel, } from '@/lib/ai-gateway/models'; +import { isFreeModel } from '@/lib/ai-gateway/is-free-model'; import { accountForMicrodollarUsage, captureProxyError, diff --git a/apps/web/src/app/api/openrouter/embeddings/route.ts b/apps/web/src/app/api/openrouter/embeddings/route.ts index 799bdca4a57..04675f441e2 100644 --- a/apps/web/src/app/api/openrouter/embeddings/route.ts +++ b/apps/web/src/app/api/openrouter/embeddings/route.ts @@ -8,7 +8,7 @@ import { debugSaveProxyRequest } from '@/lib/debugUtils'; import { captureException, setTag, startInactiveSpan } from '@sentry/nextjs'; import { getUserFromAuth } from '@/lib/user.server'; import { sentryRootSpan } from '@/lib/getRootSpan'; -import { isFreeModel } from '@/lib/ai-gateway/models'; +import { isFreeModel } from '@/lib/ai-gateway/is-free-model'; import { captureProxyError, checkOrganizationModelRestrictions, diff --git a/apps/web/src/lib/ai-gateway/experiments/membership.ts b/apps/web/src/lib/ai-gateway/experiments/membership.ts new file mode 100644 index 00000000000..98403c7568c --- /dev/null +++ b/apps/web/src/lib/ai-gateway/experiments/membership.ts @@ -0,0 +1,51 @@ +import { redisGet } from '@/lib/redis'; +import { EXPERIMENTED_PUBLIC_IDS_REDIS_KEY } from '@/lib/redis-keys'; +import { createCachedFetch } from '@/lib/cached-fetch'; + +/** + * Membership pre-check for model-experiment routing. Lives in its own module + * (separate from `pick-variant.ts`) so it can be imported from + * `lib/ai-gateway/models.ts` — which is reachable from client bundles — + * without dragging in `pick-variant.ts`'s drizzle dependency. + * + * Touches only Redis + an in-process cache. Admin writes are responsible for + * maintaining the Redis membership set; if Redis is empty, corrupt, or + * unavailable, treat it as "no experimented public ids". + */ +const EXPERIMENTED_PUBLIC_IDS_LOCAL_CACHE_TTL_MS = process.env.NODE_ENV === 'test' ? 0 : 60_000; + +const getExperimentedPublicIds = createCachedFetch( + 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/pick-variant.test.ts b/apps/web/src/lib/ai-gateway/experiments/pick-variant.test.ts index 79b90a60727..2ea839860c9 100644 --- a/apps/web/src/lib/ai-gateway/experiments/pick-variant.test.ts +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.test.ts @@ -4,7 +4,8 @@ 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, pickModelExperimentVariant } from './pick-variant'; +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'; diff --git a/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts index 58e6245b48f..bc73ede4fb0 100644 --- a/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts +++ b/apps/web/src/lib/ai-gateway/experiments/pick-variant.ts @@ -6,9 +6,6 @@ import { model_experiment_variant_version, } from '@kilocode/db/schema'; import { db } from '@/lib/drizzle'; -import { redisGet } from '@/lib/redis'; -import { EXPERIMENTED_PUBLIC_IDS_REDIS_KEY } from '@/lib/redis-keys'; -import { createCachedFetch } from '@/lib/cached-fetch'; 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'; @@ -22,52 +19,13 @@ import type { RoutingVariant, } from '@/lib/ai-gateway/experiments/pick-variant.types'; -const EXPERIMENTED_PUBLIC_IDS_LOCAL_CACHE_TTL_MS = process.env.NODE_ENV === 'test' ? 0 : 60_000; - -const getExperimentedPublicIds = createCachedFetch( - 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. Admin writes are - * responsible for maintaining the Redis membership set; if Redis is empty, - * corrupt, or unavailable, treat it as "no experimented public ids". - */ -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; - } -} - /** * 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 says the public id is experiment-routed, so preview - * traffic pays the Postgres lookup while normal models avoid it. + * 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 { diff --git a/apps/web/src/lib/ai-gateway/is-free-model.ts b/apps/web/src/lib/ai-gateway/is-free-model.ts new file mode 100644 index 00000000000..cbd8eab678b --- /dev/null +++ b/apps/web/src/lib/ai-gateway/is-free-model.ts @@ -0,0 +1,24 @@ +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 — including dedicated experimented public ids, which are + * partner/Kilo-funded for v1. + * + * 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 ( + isKiloExclusiveFreeModel(model) || + model === KILO_AUTO_FREE_MODEL.id || + (model ?? '').endsWith(':free') || + model === 'openrouter/free' || + isOpenRouterStealthModel(model ?? '') || + (await isPublicIdExperimented(model ?? '')) + ); +} diff --git a/apps/web/src/lib/ai-gateway/models.test.ts b/apps/web/src/lib/ai-gateway/models.test.ts index 8ddc2928af3..c0df8f0ed53 100644 --- a/apps/web/src/lib/ai-gateway/models.test.ts +++ b/apps/web/src/lib/ai-gateway/models.test.ts @@ -1,5 +1,6 @@ import { describe, test, expect } from '@jest/globals'; -import { autoFreeModels, isFreeModel, kiloExclusiveModels } from './models'; +import { autoFreeModels, kiloExclusiveModels } from './models'; +import { isFreeModel } from './is-free-model'; import { getInferenceProvider } from './providers/kilo-exclusive-model'; describe('isFreeModel', () => { diff --git a/apps/web/src/lib/ai-gateway/models.ts b/apps/web/src/lib/ai-gateway/models.ts index 43be294f332..fed0d76ecd1 100644 --- a/apps/web/src/lib/ai-gateway/models.ts +++ b/apps/web/src/lib/ai-gateway/models.ts @@ -34,7 +34,6 @@ import { import { isClaudeModel } from '@/lib/ai-gateway/providers/anthropic.constants'; import { GPT_CURRENT_MODEL_ID, isOpenAiModel } from '@/lib/ai-gateway/providers/openai'; import { GLM_CURRENT_MODEL_ID } from '@/lib/ai-gateway/providers/zai'; -import { isPublicIdExperimented } from '@/lib/ai-gateway/experiments/pick-variant'; export const PRIMARY_DEFAULT_MODEL = CLAUDE_SONNET_CURRENT_MODEL_ID; @@ -59,17 +58,6 @@ export const preferredModels = [ GLM_CURRENT_MODEL_ID, ]; -export async function isFreeModel(model: string): Promise { - return ( - isKiloExclusiveFreeModel(model) || - model === KILO_AUTO_FREE_MODEL.id || - (model ?? '').endsWith(':free') || - model === 'openrouter/free' || - isOpenRouterStealthModel(model ?? '') || - (await isPublicIdExperimented(model ?? '')) - ); -} - export function isPdfSupportingModel(model: string): boolean { return isClaudeModel(model) || isOpenAiModel(model) || isGrok4Model(model); } @@ -99,7 +87,7 @@ export function isKiloStealthModel(model: string): boolean { return kiloExclusiveModels.some(m => m.public_id === model && m.flags.includes('stealth')); } -function isOpenRouterStealthModel(model: string): boolean { +export function isOpenRouterStealthModel(model: string): boolean { return model.startsWith('openrouter/') && (model.endsWith('-alpha') || model.endsWith('-beta')); } diff --git a/apps/web/src/lib/ai-gateway/processUsage.ts b/apps/web/src/lib/ai-gateway/processUsage.ts index ea96a9ff4bb..54a46a2c214 100644 --- a/apps/web/src/lib/ai-gateway/processUsage.ts +++ b/apps/web/src/lib/ai-gateway/processUsage.ts @@ -21,7 +21,8 @@ import { eq, sql } from 'drizzle-orm'; import { sentryRootSpan } from '../getRootSpan'; import { ingestOrganizationTokenUsage } from '@/lib/organizations/organization-usage'; import type { ProviderId } from '@/lib/ai-gateway/providers/types'; -import { findKiloExclusiveModel, isFreeModel, isKiloStealthModel } from '@/lib/ai-gateway/models'; +import { findKiloExclusiveModel, isKiloStealthModel } from '@/lib/ai-gateway/models'; +import { isFreeModel } from '@/lib/ai-gateway/is-free-model'; import { sentryLogger } from '@/lib/utils.server'; import { maybeIssueKiloPassBonusFromUsageThreshold } from '@/lib/kilo-pass/usage-triggered-bonus'; import { getEffectiveKiloPassThreshold } from '@/lib/kilo-pass/threshold'; 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 b735e8885e0..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,8 +19,8 @@ import { buildDirectProvider, inferSupportedChatApis, } from '@/lib/ai-gateway/experiments/build-direct-provider'; +import { isPublicIdExperimented } from '@/lib/ai-gateway/experiments/membership'; import { - isPublicIdExperimented, pickModelExperimentVariant, type AllocationSubject, } from '@/lib/ai-gateway/experiments/pick-variant'; diff --git a/apps/web/src/lib/ai-gateway/providers/openrouter/index.ts b/apps/web/src/lib/ai-gateway/providers/openrouter/index.ts index 629994f1b60..3c7e39b7728 100644 --- a/apps/web/src/lib/ai-gateway/providers/openrouter/index.ts +++ b/apps/web/src/lib/ai-gateway/providers/openrouter/index.ts @@ -1,9 +1,9 @@ import { - isFreeModel, isPdfSupportingModel, kiloExclusiveModels, preferredModels, } from '@/lib/ai-gateway/models'; +import { isFreeModel } from '@/lib/ai-gateway/is-free-model'; import PROVIDERS from '@/lib/ai-gateway/providers/provider-definitions'; import type { OpenRouterModel } from '@/lib/organizations/organization-types'; import { diff --git a/apps/web/src/lib/bot/agent-runner.ts b/apps/web/src/lib/bot/agent-runner.ts index 58cc3ecbf42..14a3d348e54 100644 --- a/apps/web/src/lib/bot/agent-runner.ts +++ b/apps/web/src/lib/bot/agent-runner.ts @@ -29,7 +29,7 @@ import { formatGitLabRepositoriesForPrompt, getGitLabRepositoryContext, } from '@/lib/slack-bot/gitlab-repository-context'; -import { isFreeModel } from '@/lib/ai-gateway/models'; +import { isFreeModel } from '@/lib/ai-gateway/is-free-model'; import { generateApiToken } from '@/lib/tokens'; import { captureException } from '@sentry/nextjs'; import { createOpenAICompatible } from '@ai-sdk/openai-compatible'; From e4b1ca1d73b864632e3ecb49fa5334b8618d0d89 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Mon, 25 May 2026 09:11:57 +0200 Subject: [PATCH 49/49] docs(experiments): note response model-id leak as v1 followup --- .plans/experimental-models-1.md | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/.plans/experimental-models-1.md b/.plans/experimental-models-1.md index b1485f4ef3a..63a1ff22855 100644 --- a/.plans/experimental-models-1.md +++ b/.plans/experimental-models-1.md @@ -614,9 +614,25 @@ 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: