Repository navigation
perf(ai-gateway): run the usage write in the primary's region - #5034
Merged
Merged
Conversation
Deployment region facts are not discoverable from the repository: both Vercel projects share one vercel.json with no regions key, and there is no per-route region pinning anywhere, so the Frankfurt/SFO split only exists in the Vercel dashboard. Also record which observability fields actually name the compute region. proxy.region is documented as where the request is processed, i.e. the edge hop, while executionRegion and proxy.lambdaRegion name the function. x-vercel-id mixes PoP hops with the execution region, which is easy to misread on requests that pass through the rewrite to global-api.kilo.ai.
The AI gateway runs on kilocode-global-app, whose functions execute in both
Frankfurt and SFO, while the PostgreSQL primary is Frankfurt-only. The usage
write takes row locks on organization_user_usage and organizations (and
kilocode_users for personal usage) and holds them across the remaining
statements of the transaction until COMMIT. From SFO that hold is dominated by
transatlantic round trips rather than by database work.
pg_stat_statements on the primary shows the shape: UPDATE organizations has a
mean of 0.5ms and a max of 100,590ms, and the organization_user_usage upsert a
mean of 0.30ms and a max of 100,242ms. A sub-millisecond statement taking 100
seconds is queueing, not work. Arrival rate on a single hot counter row peaked
at 9 writes/second against a hold of a few hundred milliseconds, so the queue
grows until statements are killed at the database statement_timeout ceiling of
120s. Each waiter pins a pool connection while it waits, and with max: 10 per
instance ten of them take out every route on that instance, which is what
surfaces as "timeout exceeded when trying to connect" on unrelated endpoints.
SFO instances now hand the write to POST /api/internal/usage/record on
APP_URL, which is the Frankfurt-only kilocode-app deployment whose rewrites do
not divert /api/internal/*. Frankfurt instances keep writing directly, since a
Frankfurt-to-Frankfurt hop is pure overhead and a pointless failure mode.
The seam is saveUsageRelatedData: the caller keeps stream parsing, cost
computation and NUL sanitization, and ships the already-normalized
{ core, metadata } from toInsertableDbUsageRecord. That payload carries no
prompt or response bodies.
Correctness details:
- The endpoint is idempotent on core.id, which the sender generates and which is
the microdollar_usage primary key. Without that, a caller retry after a lost
response would collide on the primary key and report a billed request as
unbilled.
- Nullable contract fields use .nullable(), never .optional(), because
JSON.stringify drops undefined keys and null is load-bearing: org usage sets
the prompt prefixes to null precisely so prompt text is never persisted.
- created_at is validated as strict ISO 8601 so a PostgreSQL-shaped timestamp
can never enter the path, per packages/db/AGENTS.md. Covered by fixtures.
- On any non-retryable response or exhausted retries the client returns
unavailable and the caller falls back to the local write. A slow billing
record beats a lost one; the failure is reported to Sentry with the usageId.
No new environment variables: APP_URL and INTERNAL_API_SECRET already exist.
A retry of the Frankfurt usage write can arrive while the first attempt is still open: the client retries its 10s attempt timeout, and after exhausting its retries it falls back to writing locally with the same core.id. The route's pre-flight lookup only catches a redelivery that arrives after the first one committed, so the in-flight case collided on the microdollar_usage primary key and was reported as a failed write even though the usage was billed, which made callers skip dependent writes such as experiment attribution. insertUsageRecord now looks the row up when the write fails and returns its identity when it is there. The row's presence under the same user is the proof rather than the PostgreSQL error code, so a deadlock victim retried into the collision recovers the same way.
`usage-record-client.ts` imported `UsageRecordInsertResult` from `processUsage.ts`, which imports that client to hand the write to Frankfurt. The two modules were therefore mutually dependent, and `dependency-cycle-check` rejects that even though the import is type-only and erased at runtime. Move the type to `processUsage.types.ts`, which is already the shared home for the usage types the write path and the wire contract both need.
… once `isFirstUsage` can only be answered from committed rows, so a redelivery that arrives while the first delivery's transaction is still open also classified the record as the user's first usage and emitted `first_usage` a second time. `recordUsageInPrimaryRegion` produces exactly that: it retries after a 10s attempt timeout while the write it started may still be committing. Row-level idempotency covers the billing row, not this pre-insert side effect. `insertUsageRecord` now reports whether the identity was recovered from an already-committed row, and `saveUsageRelatedDataLocally` emits the first-usage events only after an insert of its own commits. `isFirst` is still computed before the insert, because afterwards the record is its own prior usage. The route maps a recovered write to the existing `duplicate` status and builds the response `result` explicitly, so the new flag does not leak onto a wire contract that does not document it. The regression test lives in its own file because `processUsage` resolves its PostHog client at module scope, so the mock has to be registered before that import.
RSO
force-pushed
the
database-connection-timeo
branch
from
August 6, 2026 09:05
07bd889 to
b9ce59b
Compare
RSO
marked this pull request as ready for review
August 6, 2026 09:05
![kilo-code-bot[bot]](https://avatars.githubusercontent.com/in/2193792?s=60&v=4){kind=small}
Contributor
Code Review SummaryStatus: No Issues Found | Recommendation: Merge Executive SummaryIncremental review of commit Files Reviewed (1 file)
Previous Review Summary (commit b9ce59b)Current summary above is authoritative. Previous snapshots are kept for context only. Previous review (commit b9ce59b)Status: 1 Issue Found | Recommendation: Address before merge Executive SummaryThe new Overview
Issue Details (click to expand)WARNING
Files Reviewed (10 files)
Reviewed by kimi-k3 · Input: 75.2K · Output: 4.5K · Cached: 180.2K Review guidance: REVIEW.md from base branch |
…write A fresh insert reports `RETURNING created_at`, which PostgreSQL renders as "2026-04-29 01:16:12.945+00" — not the strict ISO 8601 the rest of this path guarantees. `core.created_at` is the same instant in the guaranteed format, because that column was inserted from it, and it is already what the `duplicate` branch and `logMicrodollarUsage` report.
chrarnoldus
reviewed
Aug 6, 2026
| abuse_classification: z.nativeEnum(ABUSE_CLASSIFICATION), | ||
| organization_id: z.string().uuid().nullable(), | ||
| inference_provider: z.string().nullable(), | ||
| project_id: z.string().nullable(), |
Contributor
There was a problem hiding this comment.
If you make changes here they should be backwards and forwards compatible I assume?
Contributor
Author
There was a problem hiding this comment.
Yeah, I think so. I was wondering whether a retry would pick up the new schema, but I guess not.
chrarnoldus
approved these changes
Aug 6, 2026
3 tasks
RSO
added a commit
that referenced
this pull request
Aug 6, 2026
* feat(ai-gateway): add usage-record write diagnostics POST /api/internal/usage/record is strongly bimodal: p50 around 5ms against a p95 near 50s. The database is idle throughout — zero blocked sessions, sub-second oldest transaction, and log_lock_waits (on, deadlock_timeout 1s) producing only a handful of entries per hour — and Supavisor reports client_waiting 0 against a pool size of 320. So the time is spent inside the Node process and nothing currently emitted can say where. Two candidates need different fixes, so they have to be separated by measurement rather than argument: - Waiting on the in-process pg pool, capped at 10 connections per instance. /api/cron/db-pool-metrics reports Supavisor's counters, not this pool's, so its saturation is invisible today. waitingCount settles it. - Event-loop starvation, which would inflate every phase including ones that do no I/O. The delay histogram settles that one. Also adds describeDatabaseError, which extracts the driver's code, constraint, table and routine from a drizzle-wrapped failure while carrying none of the statement. A drizzle error stringifies to the interpolated query, which for this write is ~30KB and includes user_prompt_prefix, system_prompt_prefix, the client IP, city and JA4 digest. detail is omitted too, since PostgreSQL puts the conflicting key values there. isUsageRowConflict identifies a delivery colliding with its own id. It is deliberately narrow: microdollar_usage and microdollar_usage_metadata are keyed by the per-delivery id, while the twelve deduplicated lookup tables have a WHERE NOT EXISTS race that retrying genuinely does resolve. Constraint names were checked against pg_constraint — the lookup tables use <table>_pkey and microdollar_usage carries the generated PK_a71b90d910e7882358c3defe8e2. Emission is gated on a 1s threshold plus a 1% sample, and it replaces a log line that serialized ~30KB per failure, so net log volume falls. * fix(ai-gateway): stop the usage-record redelivery storm Moving the usage write to a Frankfurt-local endpoint (#5034) fixed the hot-row convoy — the local write is 5ms at p50 and the primary is now idle — but replaced it with a self-inflicted retry storm. recordUsageInPrimaryRegion used a 10s per-attempt timeout against an endpoint whose p95 is roughly 50s. Every request slower than the timeout was abandoned and re-sent with the same core.id, so around 20% of deliveries became redeliveries colliding on the microdollar_usage primary key: about 4,400 collisions per five minutes, each one adding load that made the endpoint slower still. The failures were invisible because the client only logs after exhausting every attempt, which happened just 167 times. Three changes break the loop: - The attempt budget is 90s, above the observed tail and below the endpoint's maxDuration of 150s. A timeout is now never retried: the endpoint may still be mid-write, so re-sending reproduces exactly this collision. It falls back to the local write instead, which is slower but which the endpoint's conflict handling already recognises. - Retryable statuses narrow from "5xx or 429" to 429, 502 and 503 — refusals that prove the function did not run. A 500 means our handler threw, possibly after commit, and 504 is equally ambiguous; both used to be retried. - insertUsageRecord stops retrying a collision on this record's own id. id is fixed for the delivery, so all three attempts collided and each rebuilt and re-logged the same ~30KB statement before the outer handler recovered the committed identity. Failure logs no longer carry the raw error. Both the retry warning and the final captureException report the redacted description instead, so prompt prefixes, client IP, city and JA4 digest stop reaching Axiom and Sentry. The stack is preserved on the replacement error so the call site stays identifiable. The endpoint now records phase timings, in-process pool gauges and event-loop delay on one structured line, which is what will locate the remaining tail. * fix(ai-gateway): keep redacted usage errors free of the statement Two findings from review on #5098. The redacted Sentry error copied the original stack to preserve the throw site, but V8 renders error.stack as `${name}: ${message}` followed by the frames, so the copy put the interpolated statement — prompt prefixes, client IP, city, JA4 digest — straight back into the captured event. Verified: the assembled stack still contained the prompt text and IP. That defeated the redaction it sat next to. stackFramesUnderHeader carries over only the frames, under a caller-supplied header. Frames are `at fn (chunk:line:col)` and hold no request data, so the throw site survives, which Error.captureStackTrace on the replacement would have discarded. Separately, `await response.json()` sat outside any guard, so a truncated or non-JSON body threw into the outer catch and was treated as a retryable transport failure. The status was already 2xx at that point, so the write may have committed — the same hazard as the schema mismatch immediately below, which deliberately does not retry. Both now fall back instead of re-sending and manufacturing a redelivery. * fix(ai-gateway): strip the error header by exact prefix, not by frame search Follow-up finding on #5098. stackFramesUnderHeader located the message boundary with `/^\s+at\s/m`, a heuristic applied to a string that embeds user_prompt_prefix. A user pasting a stack trace into their prompt — routine for a coding assistant — puts a frame-shaped line inside the message, so the search matched there and the slice carried the remaining parameters into the result. Verified against a prompt containing a pasted trace: the assembled stack still held the client IP and city. The previous fix closed the header leak and left this one open. The header is now removed by exact `${name}: ${message}` prefix, so no part of the message can survive, and the result is filtered to frame-shaped lines as a second barrier against a message mutated after construction. Fails closed: an unrecognised header means the boundary is unknown, so no frames are emitted. describeDatabaseError still carries the actionable code and constraint, and losing a throw site is cheaper than leaking a prompt. Also renames the shadowed `payload` binding in the response-body guard, noted in the same review as currently harmless.
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
The AI gateway runs on
kilocode-global-app, whose functions execute in both Frankfurt and SFO, while the PostgreSQL primary is Frankfurt-only. The usage write takes row locks onorganization_user_usageandorganizations(andkilocode_usersfor personal usage) and holds them across the remaining statements of the transaction untilCOMMIT. From SFO that hold is dominated by transatlantic round trips rather than by database work.pg_stat_statementson the production primary shows the shape:UPDATE organizations SET microdollars_used …INSERT organization_user_usage … ON CONFLICT DO UPDATEmicrodollar_usage_insCTEA sub-millisecond statement taking 100 seconds is queueing, not work, and the CTE's max is clipped exactly at the database
statement_timeoutceiling of 120 s. Measured arrival rate on a single hot counter row peaked at 9 writes/second against a hold of a few hundred milliseconds, so the queue grows without bound. Each waiter pins a pool connection while it waits, and withmax: 10per instance ten of them take out every route on that instance — which is what surfaces in Sentry astimeout exceeded when trying to connecton unrelated endpoints like/api/profileand/api/webhooks/github.SFO instances now hand the write to
POST /api/internal/usage/recordonAPP_URL, the Frankfurt-onlykilocode-appdeployment whose rewrites do not divert/api/internal/*. Frankfurt instances keep writing directly, since a Frankfurt-to-Frankfurt hop is pure overhead and a pointless new failure mode.The seam is
saveUsageRelatedData. The caller keeps stream parsing, cost computation and NUL sanitization, and ships the already-normalized{ core, metadata }fromtoInsertableDbUsageRecord. That payload carries no prompt or response bodies.Correctness details worth knowing:
core.id, which the sender generates and which is themicrodollar_usageprimary key. Without it, a caller retry after a lost response would collide on the primary key and report a billed request as unbilled..nullable(), never.optional().JSON.stringifydropsundefinedkeys, andnullis load-bearing here: org usage sets the prompt prefixes tonullprecisely so prompt text is never persisted.created_atis validated as strict ISO 8601 so a PostgreSQL-shaped timestamp can never enter the path, perpackages/db/AGENTS.md. Covered by fixtures using production-shaped text.unavailableand the caller falls back. A slow billing record beats a lost one; the failure goes to Sentry with theusageIdso a genuine loss is reconcilable.No new environment variables —
APP_URLandINTERNAL_API_SECRETalready exist.The first commit is an unrelated documentation addition recording the Frankfurt/SFO function-region split and which observability fields actually name the compute region. It is separable if you'd rather it went on its own.
Verification
Automated:
tsgo --noEmit -p apps/web/tsconfig.jsonclean;./scripts/lint-all.sh0 warnings 0 errors; 29 new tests passing; regression runs ofsrc/lib/ai-gateway+src/lib/drizzle.test.ts(626 tests, 54 suites),src/lib/kilo-pass-org+organization-usage(127 tests), andspend-writer-audit. This is targeted verification, notpnpm validate.Manual: none. The behaviour that matters only exists on a multi-region production deployment — locally
VERCEL_REGIONis unset, soisUSRegion()is false and the code takes the pre-existing local path. I could not exercise the SFO branch, the HTTP hop, or the dedupe path by hand.INTERNAL_API_SECRETis present onkilocode-global-app. If it is missing, every SFO write silently takes the fallback path and nothing improves.app.kilo.airesolves tokilocode-app. If that domain ever moves to the global project this silently becomes a no-op hop.POST /api/internal/usage/recordagainst a preview deployment, including sending the samecore.idtwice to seestatus: "duplicate".recordUsageInPrimaryRegionfailures in Sentry and themax_exec_timeofUPDATE organizationsbefore/after.Visual Changes
N/A
Reviewer Notes
Please push back on the premise. I have not proven that the waiters are each other. The mean-vs-max evidence is statistical; I never captured
pg_locksor wait events during a burst. A single long-running transaction holding anorganizationsrow — an admin script, a cron, a batch job — would produce identical maxima with no convoy, and in that case this change does not help. Apg_stat_activity+pg_blocking_pids()sample during an incident would settle it and is cheap to add to the existingdb-pool-metricscron. I'd rather that ran before this is relied on.Known gaps in test coverage:
saveUsageRelatedDatais module-private and reaching it means going throughcountAndStoreUsage, which needs aResponse.isUSRegionand the client are tested separately, but the seam between them is not.Other things to look at:
apps/web(see the comments inlib/auto-fix/github/get-fix-config.ts). This is cross-deployment rather than same-instance, which is why the hop buys locality — but it is a fair objection and theisUSRegion()gate exists partly to answer it.maxDuration = 150on the new route is deliberately above the 120 s databasestatement_timeoutso a blocked write is reported rather than truncated into a lost billing row.duplicate.organization_user_usageis one row per (org, user, day) andorganizations.microdollars_usedis shared by every member. Batching in a single place is the natural follow-up, and the endpoint is now that place.