feat(gateway): route and persist model experiments

apps/web/src/app/api/openrouter/[...path]/route.ts

@@ -18,6 +18,8 @@ import type {
 } from '@/lib/ai-gateway/providers/openrouter/types';
 import { applyProviderSpecificLogic } from '@/lib/ai-gateway/providers/apply-provider-specific-logic';
 import { getProvider } from '@/lib/ai-gateway/providers/get-provider';
+import { buildExperimentPromptCapture } from '@/lib/ai-gateway/experiments/persist';
+import { isPublicIdExperimented } from '@/lib/ai-gateway/experiments/membership';
 import { upstreamRequest } from '@/lib/ai-gateway/providers/upstream-request';
 import { debugSaveProxyRequest } from '@/lib/debugUtils';
 import { setTag, startInactiveSpan } from '@sentry/nextjs';
@@ -43,6 +45,7 @@ import {
   malformedJsonResponse,
   makeErrorReadable,
   modelDoesNotExistResponse,
+  modelNotAllowedResponse,
   extractHeaderAndLimitLength,
   noFreeModelsAvailableResponse,
   temporarilyUnavailableResponse,
@@ -276,7 +279,9 @@ export async function POST(request: NextRequest): Promise<NextResponseType<unkno
   // per user when the request comes from Cloudflare IPs (Kilo infrastructure).
   // All other products rate-limit per IP (fast pre-auth path).
   const isRateLimitedFreeModelRequest =
-    isKiloExclusiveFreeModel(originalModelIdLowerCased) || autoModel === KILO_AUTO_FREE_MODEL.id;
+    isKiloExclusiveFreeModel(originalModelIdLowerCased) ||
+    autoModel === KILO_AUTO_FREE_MODEL.id ||
+    (await isPublicIdExperimented(originalModelIdLowerCased));
   if (isRateLimitedFreeModelRequest) {
     const rateLimit = await resolveRateLimit(feature, ipAddress, authPromise);
     if (rateLimit instanceof NextResponse) return rateLimit;
@@ -386,14 +391,38 @@ export async function POST(request: NextRequest): Promise<NextResponseType<unkno
     user,
     organizationId,
     taskId,
+    clientIp: ipAddress ?? null,
+    machineId: machineIdHeader,
   });
   if (providerResult.kind === 'not-found') {
+    // Paused experiment for this public id — return a local model-unavailable
+    // response instead of silently falling through to default routing.
     return modelDoesNotExistResponse();
   }
   if (providerResult.kind === 'unavailable') {
     return temporarilyUnavailableResponse();
   }
-  const { provider, userByok, bypassAccessCheck } = providerResult;
+  const {
+    provider,
+    userByok,
+    bypassAccessCheck,
+    skipProviderPin,
+    skipKiloExclusiveModelSettings,
+    experiment,
+  } = providerResult;
+
+  // Request-level data-collection opt-out: a caller can set
+  // `provider.data_collection: 'deny'` or `provider.zdr: true` on any
+  // request to opt that single request out of training/data-retention.
+  // Direct experiment upstreams ignore those OpenRouter/Vercel flags
+  // (we never reach OpenRouter), but we still capture the prompt to R2
+  // for partner evaluation — which violates the caller's stated
+  // intent. Refuse here regardless of org settings, anon/BYOK status,
+  // or the org-level check below.
+  if (experiment && !isFreePromptTrainingAllowed(requestBodyParsed.body.provider)) {
+    return dataCollectionRequiredResponse();
+  }
+
   if (!provider.supportedChatApis.includes(requestBodyParsed.kind)) {
     return apiKindNotSupportedResponse(
       requestBodyParsed.kind,
@@ -483,11 +512,38 @@ export async function POST(request: NextRequest): Promise<NextResponseType<unkno
     });
     if (modelRestrictionError) return modelRestrictionError;
 
-    if (providerConfig) {
+    // Experiment traffic captures prompts to R2 for partner evaluation, which
+    // is a form of data collection that the gateway-pinned `data_collection`
+    // setting cannot enforce on a direct partner upstream. If the org has
+    // explicitly disabled data collection, refuse the experimented public id
+    // here rather than routing through and silently capturing prompts.
+    if (experiment && settings?.data_collection === 'deny') {
+      return dataCollectionRequiredResponse();
+    }
+
+    // Enterprise `provider_allow_list` is enforced via OpenRouter's
+    // `body.provider.only` field, which doesn't reach a direct partner
+    // upstream. Refuse the experimented public id rather than routing
+    // around the org's allow-list.
+    if (experiment && providerConfig?.only !== undefined) {
+      return modelNotAllowedResponse();
+    }
+
+    // Direct experiment upstreams must not have a Vercel/OpenRouter
+    // provider config pinned onto them — the partner endpoint is selected
+    // by the variant version.
+    if (providerConfig && !skipProviderPin) {
       requestBodyParsed.body.provider = providerConfig;
     }
   }
 
+  if (experiment) {
+    usageContext.modelExperimentVariantVersionId = experiment.variantVersionId;
+    usageContext.modelExperimentAllocationSubject = experiment.allocationSubject;
+    // Cost zeroing for experiment traffic is handled by `isFreeModel`, which
+    // returns true for experimented public ids.
+  }
+
   sentryRootSpan()?.setAttribute(
     'openrouter.time_to_request_start_ms',
     performance.now() - requestStartedAt
@@ -535,9 +591,18 @@ export async function POST(request: NextRequest): Promise<NextResponseType<unkno
     requestBodyParsed,
     extraHeaders,
     userByok,
-    fraudHeaders
+    fraudHeaders,
+    { skipKiloExclusiveModelSettings: skipKiloExclusiveModelSettings === true }
   );
 
+  // Capture the bounded prompt for experimented requests AFTER provider
+  // transforms have produced the canonical upstream body. Stored on the
+  // usage context so the async `after()` hook can persist it without
+  // retaining a reference to the full uncapped body.
+  if (experiment) {
+    usageContext.experimentPromptCapture = buildExperimentPromptCapture(requestBodyParsed);
+  }
+
   const response = await upstreamRequest({
     path,
     search: url.search,

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

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<string[]>(
+  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<boolean> {
+  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;
+  }
+}

apps/web/src/lib/ai-gateway/experiments/persist.test.ts

@@ -0,0 +1,132 @@
+import { describe, expect, it } from '@jest/globals';
+import {
+  buildExperimentPromptCapture,
+  REQUEST_BODY_CAP_BYTES,
+  truncateToUtf8Bytes,
+} from './persist';
+import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types';
+
+describe('truncateToUtf8Bytes', () => {
+  it('returns the input unchanged when already within the cap', () => {
+    expect(truncateToUtf8Bytes('hello', 100)).toBe('hello');
+  });
+
+  it('returns valid UTF-8 with length <= maxBytes for ASCII', () => {
+    const out = truncateToUtf8Bytes('a'.repeat(1000), 16);
+    expect(Buffer.byteLength(out, 'utf8')).toBe(16);
+    expect(out).toBe('a'.repeat(16));
+  });
+
+  it('does not split a multi-byte UTF-8 codepoint', () => {
+    // '日' is 3 bytes in UTF-8 (0xE6 0x97 0xA5). With maxBytes=4 the cut
+    // would otherwise land mid-codepoint and produce invalid UTF-8.
+    const input = '日日日'; // 9 UTF-8 bytes
+    const out = truncateToUtf8Bytes(input, 4);
+    // Result must be valid UTF-8 and <= 4 bytes. The only valid prefixes
+    // here are '' (0 bytes) and '日' (3 bytes).
+    expect(Buffer.byteLength(out, 'utf8')).toBeLessThanOrEqual(4);
+    expect(out).toBe('日');
+  });
+
+  it('handles 4-byte emoji codepoints correctly', () => {
+    // '😀' is 4 bytes in UTF-8 (a surrogate pair in UTF-16).
+    const input = '😀😀😀';
+    const out = truncateToUtf8Bytes(input, 6);
+    // Only '😀' (4 bytes) fits cleanly; cutting at byte 6 would split
+    // the second emoji.
+    expect(out).toBe('😀');
+    expect(Buffer.byteLength(out, 'utf8')).toBe(4);
+  });
+
+  it('produces deterministic output across calls', () => {
+    const input = 'hello 日本語 world 🎉 こんにちは';
+    const a = truncateToUtf8Bytes(input, 12);
+    const b = truncateToUtf8Bytes(input, 12);
+    expect(a).toBe(b);
+    expect(Buffer.byteLength(a, 'utf8')).toBeLessThanOrEqual(12);
+  });
+
+  it('cuts cleanly when the cap exactly aligns with a codepoint boundary', () => {
+    const input = '日日'; // 6 bytes
+    const out = truncateToUtf8Bytes(input, 3);
+    expect(out).toBe('日');
+  });
+
+  it('handles maxBytes of 0', () => {
+    expect(truncateToUtf8Bytes('hello', 0)).toBe('');
+  });
+});
+
+describe('buildExperimentPromptCapture', () => {
+  function chatRequest(body: Record<string, unknown>): GatewayRequest {
+    // Tests may pass synthetic message shapes that don't fully match
+    // the production OpenAI types; double-cast through unknown is the
+    // pragmatic test-only escape hatch.
+    return { kind: 'chat_completions', body } as unknown as GatewayRequest;
+  }
+
+  it('captures the full serialized body and records the request kind', () => {
+    const cap = buildExperimentPromptCapture(
+      chatRequest({
+        model: 'kilo/preview-foo',
+        messages: [
+          { role: 'system', content: 'you are a helpful assistant' },
+          { role: 'user', content: 'hi' },
+        ],
+      })
+    );
+    expect(cap.requestKind).toBe('chat_completions');
+    expect(cap.requestBodyContent).toContain('"role":"system"');
+    expect(cap.requestBodyContent).toContain('"role":"user"');
+    expect(cap.wasTruncated).toBe(false);
+  });
+
+  it('records `messages` for an Anthropic-shape request', () => {
+    const cap = buildExperimentPromptCapture({
+      kind: 'messages',
+      body: {
+        model: 'kilo/preview-foo',
+        system: 'you are a helpful assistant',
+        messages: [{ role: 'user', content: 'hi' }],
+      },
+    } as unknown as GatewayRequest);
+    expect(cap.requestKind).toBe('messages');
+    expect(cap.requestBodyContent).toContain('"system":"you are a helpful assistant"');
+  });
+
+  it('records `responses` for a Responses-API request', () => {
+    const cap = buildExperimentPromptCapture({
+      kind: 'responses',
+      body: { model: 'kilo/preview-foo', input: 'hi' },
+    } as unknown as GatewayRequest);
+    expect(cap.requestKind).toBe('responses');
+    expect(cap.requestBodyContent).toContain('"input":"hi"');
+  });
+
+  it('truncates the body by UTF-8 bytes and marks wasTruncated', () => {
+    const huge = 'a'.repeat(REQUEST_BODY_CAP_BYTES + 1024);
+    const cap = buildExperimentPromptCapture(
+      chatRequest({
+        model: 'kilo/preview-foo',
+        messages: [{ role: 'user', content: huge }],
+      })
+    );
+    expect(cap.wasTruncated).toBe(true);
+    expect(Buffer.byteLength(cap.requestBodyContent, 'utf8')).toBeLessThanOrEqual(
+      REQUEST_BODY_CAP_BYTES
+    );
+  });
+
+  it('produces deterministic output for identical inputs', () => {
+    // Content addressing on top of this capture relies on byte-for-byte
+    // determinism; a JSON.stringify implementation that reordered keys
+    // would silently break dedup.
+    const body = {
+      model: 'kilo/preview-foo',
+      messages: [{ role: 'user', content: 'hi' }],
+    };
+    const a = buildExperimentPromptCapture(chatRequest({ ...body }));
+    const b = buildExperimentPromptCapture(chatRequest({ ...body }));
+    expect(a.requestBodyContent).toBe(b.requestBodyContent);
+  });
+});