Repository navigation
feat(gateway): route and persist model experiments - #3325
Merged
Merged
Conversation
added 7 commits
May 19, 2026 12:06
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.
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.
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.
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.
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.
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.
markijbema
force-pushed
the
mark/experimental-models-gateway
branch
from
May 19, 2026 10:50
4864b01 to
ebceefb
Compare
added 7 commits
May 19, 2026 12:53
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.
…raffic 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`.
…llection 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.
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.
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.
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).
Contributor
Author
|
Addressed review feedback in the latest pushes:
Each fix is a separate commit for review. |
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.
Contributor
Author
|
Addressed the request-level data-collection follow-up:
Commit: 33b3baf |
markijbema
marked this pull request as ready for review
May 19, 2026 12:19
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)
![kilo-code-bot[bot]](https://avatars.githubusercontent.com/in/2193792?s=60&v=4){kind=small}
Contributor
Code Review SummaryStatus: 1 Suggestion | Recommendation: Merge Executive SummaryNo new issues introduced by the incremental commits. The docs commit ( Overview
Issue Details (click to expand)SUGGESTION| File | Line | Issue | New commits reviewed (since 1eefb73)
Resolved since last review
Files Reviewed
Fix these issues in Kilo Cloud Reviewed by claude-sonnet-4.6 · 4,174,653 tokens Review guidance: REVIEW.md from base branch |
The previous commit (b6b1ebc) 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.
Contributor
Author
|
Extracted the buildDirectProvider refactor as a standalone prep PR: #3330. Either order of merge works:
Refactor is pure and well-tested by the existing custom_llm coverage (83 tests). No behavioral change to the live kilo-internal/... path. |
added 2 commits
May 19, 2026 15:39
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.
…els-gateway # Conflicts: # apps/web/src/lib/ai-gateway/providers/get-provider.ts
This was referenced May 19, 2026
added 3 commits
May 21, 2026 14:05
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.
RSO
approved these changes
May 21, 2026
added 7 commits
May 21, 2026 15:50
…els-gateway # Conflicts: # apps/web/src/app/api/openrouter/[...path]/route.ts # apps/web/src/lib/ai-gateway/processUsage.ts
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.
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.
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).
…els-gateway # Conflicts: # .plans/experimental-models-1.md
markijbema
force-pushed
the
mark/experimental-models-gateway
branch
from
May 21, 2026 14:52
b2be6a9 to
cb9f60d
Compare
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.
…els-gateway # Conflicts: # apps/web/src/lib/ai-gateway/is-free-model.ts
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements backend model-experiment routing and attribution persistence for preview/provider-funded traffic.
getProviderchecks a Redis-backed experimented-public-ID membership set, then loads the selected experiment, variants, and current variant versions from Postgres for preview traffic only.buildDirectProviderpath. Paused experiments return local model-unavailable behavior instead of falling through to default providers.BYOK_ENCRYPTION_KEYat config load and addsR2_EXPERIMENT_PROMPTS_BUCKET_NAMEfor prompt blob storage.Verification
model_experiment_requestrows land 1:1 withmicrodollar_usage.client_request_idis populated fromx-kilo-request.upstream.base_urlrather than OpenRouter/Vercel.variant_version_idfrom Postgres while older attribution rows still point at the prior version.Visual Changes
N/A
Reviewer Notes
EXPERIMENTED_PUBLIC_IDS_REDIS_KEY; if Redis is missing, corrupt, or unavailable, gateway treats that as “no preview IDs.”persistExperimentAttributionis intentionally best-effort and runs after the microdollar usage write. Billing remains the source of truth.__failed__.route.ts,get-provider.ts); variant selection and fail-closed behavior (pick-variant.ts,pick-variant.types.ts); billing and policy guarantees (processUsage.ts,llm-proxy-helpers.ts); attribution persistence (persist.ts,experiment-prompts.ts); tests (pick-variant.test.ts,persist.test.ts, related provider/request tests).