perf(ai-gateway): run the usage write in the primary's region

AGENTS.md

@@ -66,6 +66,43 @@ package manifests before running repository JavaScript or package scripts. Load
 | UI and product design | `DESIGN.md`, relevant app `AGENTS.md`; for `apps/web`, the `kilo-design-cloud` skill synced from `Kilo-Org/kilo-design` |
 | Contribution and PR workflow | `CONTRIBUTING.md` and relevant Git or PR skill |
 
+## Vercel Function Regions
+
+Our functions only ever run in these regions, regardless of what request metadata
+suggests:
+
+| Vercel project | Function regions |
+|---|---|
+| `kilocode-app` | Frankfurt only |
+| `kilocode-global-app` | Frankfurt and us-west (SFO) |
+
+Treat this table as the source of truth for compute location, database
+round-trip latency, and replica reasoning. If an observed region is not `fra1` or
+`sfo1`, it is a proxy/edge hop, not the function. Confirm against the Vercel
+project's region list rather than inferring from a request or log field.
+
+Region fields are not interchangeable. Per Vercel's Log Drains reference:
+
+| Field | Meaning |
+|---|---|
+| `executionRegion` | Region where the request is executed |
+| `proxy.lambdaRegion` | Region where the function executed |
+| `proxy.region` | Region where the request is **processed** — the proxy/edge hop |
+
+So a `region` field on a log's proxy object is the edge, and `executionRegion` /
+`proxy.lambdaRegion` are the function. Do not quote a bare "region" from a log
+line as the compute location without checking which of these it maps to.
+
+`x-vercel-id` mixes both: Vercel documents it as "a list of Vercel regions your
+request hit, as well as the region the function was executed in". Never read its
+leading token as the compute location, especially for requests that pass through
+a rewrite to another Vercel app, where PoP hops accumulate.
+
+`VERCEL_REGION` is documented as "the ID of the Region where the app is running",
+i.e. the function region. Given the table above it should only ever be `fra1` or
+`sfo1`, which is what makes `isUSRegion` in `apps/web/src/lib/drizzle.ts` behave
+correctly for the SFO half of `kilocode-global-app`.
+
 ## Security Baseline
 
 - Never log tokens, credentials, API keys, authentication headers, cookies, or webhook secrets. Use `redactSensitiveHeaders` when headers must be retained or logged. Do not enable `sendDefaultPii` or `attachRpcInput` in Sentry.

apps/web/src/app/api/internal/usage/record/route.ts

@@ -0,0 +1,92 @@
+import { timingSafeEqual } from '@kilocode/encryption';
+import { microdollar_usage } from '@kilocode/db/schema';
+import { eq } from 'drizzle-orm';
+import { NextResponse } from 'next/server';
+import type { NextRequest } from 'next/server';
+
+import { INTERNAL_API_SECRET } from '@/lib/config.server';
+import { db } from '@/lib/drizzle';
+import { saveUsageRelatedDataLocally } from '@/lib/ai-gateway/processUsage';
+import {
+  UsageRecordRequestSchema,
+  type UsageRecordResponse,
+} from '@/lib/ai-gateway/usage-record-contract';
+
+/**
+ * Frankfurt-local sink for AI-gateway usage writes.
+ *
+ * Callers are SFO instances of `kilocode-global-app`, which are a transatlantic
+ * round trip away from the PostgreSQL primary. Executing the write here collapses
+ * the row-lock hold on `kilocode_users` / `organizations` /
+ * `organization_user_usage` from hundreds of milliseconds to single-digit
+ * milliseconds. See `usage-record-client.ts` for the calling side.
+ *
+ * This route only does the right thing on a deployment whose functions run in
+ * Frankfurt. `APP_URL` points at `kilocode-app`, which is Frankfurt-only, and
+ * whose rewrites do not divert `/api/internal/*`.
+ */
+
+// The usage transaction can legitimately block on a contended counter row up to
+// the database `statement_timeout` ceiling of 120s. Allow more than that so a
+// slow write is reported rather than silently truncated into a lost billing row.
+export const maxDuration = 150;
+
+export async function POST(request: NextRequest): Promise<NextResponse> {
+  const secret = request.headers.get('x-internal-api-key');
+  if (!INTERNAL_API_SECRET || !secret || !timingSafeEqual(secret, INTERNAL_API_SECRET)) {
+    return NextResponse.json({ error: 'Unauthorized' }, { status: 401 });
+  }
+
+  const rawBody: unknown = await request.json().catch(() => null);
+  const parsed = UsageRecordRequestSchema.safeParse(rawBody);
+  if (!parsed.success) {
+    // Deliberately a 400: the client must not retry a payload we cannot accept,
+    // and a schema failure here means a lost billing row that needs a human.
+    console.error('usage record request failed validation', {
+      issues: parsed.error.issues.map(issue => ({ path: issue.path.join('.'), code: issue.code })),
+    });
+    return NextResponse.json({ error: 'Invalid body' }, { status: 400 });
+  }
+
+  const { core, metadata, prior_microdollar_usage, posthog_distinct_id } = parsed.data;
+
+  // Idempotency. `core.id` is generated by the sender (`randomUUID()` in
+  // `toInsertableDbUsageRecord`) and is the `microdollar_usage` primary key, so a
+  // redelivery of a call that already committed is detectable. Without this, a
+  // caller-side retry after a lost response would collide on the primary key,
+  // burn the retry budget in `insertUsageRecord`, and report the request as
+  // unbilled even though it was billed.
+  const existing = await db
+    .select({ id: microdollar_usage.id })
+    .from(microdollar_usage)
+    .where(eq(microdollar_usage.id, core.id))
+    .limit(1);
+
+  if (existing.length > 0) {
+    // Post-commit side effects ran on the first delivery, so do not re-run them.
+    // `newMicrodollarsUsed` is not reconstructable here and is only used for
+    // PostHog attribution, which already fired.
+    const duplicate: UsageRecordResponse = {
+      status: 'duplicate',
+      result: { usageId: core.id, createdAt: core.created_at, newMicrodollarsUsed: null },
+    };
+    return NextResponse.json(duplicate);
+  }
+
+  const result = await saveUsageRelatedDataLocally(
+    core,
+    metadata,
+    prior_microdollar_usage,
+    posthog_distinct_id
+  );
+
+  // `saveUsageRelatedDataLocally` swallows database errors and returns null,
+  // matching the pre-existing local behaviour. Report it as a successful HTTP
+  // exchange with a negative outcome so the client does not retry a write that
+  // deliberately gave up.
+  const response: UsageRecordResponse = result
+    ? { status: 'recorded', result }
+    : { status: 'not_recorded', result: null };
+
+  return NextResponse.json(response);
+}

apps/web/src/lib/ai-gateway/processUsage.ts

@@ -1,5 +1,6 @@
 import { randomUUID } from 'crypto';
-import { db } from '../drizzle';
+import { db, isUSRegion } from '../drizzle';
+import { recordUsageInPrimaryRegion } from './usage-record-client';
 import type { MicrodollarUsage } from '@kilocode/db/schema';
 import { microdollar_usage } from '@kilocode/db/schema';
 import { createTimer } from '@/lib/timer';
@@ -297,11 +298,56 @@ export async function logMicrodollarUsage(
   return inserted ? { usageId: core.id, createdAt: core.created_at } : null;
 }
 
+/**
+ * Dispatches the usage write to whichever side of the Atlantic the PostgreSQL
+ * primary is on.
+ *
+ * `kilocode-global-app` executes in both Frankfurt and SFO while the primary is
+ * Frankfurt-only. This write holds row locks on `kilocode_users`,
+ * `organizations` and `organization_user_usage` across several sequential
+ * statements, so from SFO the lock hold is dominated by transatlantic round
+ * trips rather than by database work — which is what turns a contended counter
+ * row into a queue and, downstream, exhausts the connection pool.
+ *
+ * Frankfurt instances keep writing directly: a Frankfurt-to-Frankfurt HTTP hop
+ * would be pure overhead and a pointless new failure mode.
+ */
 async function saveUsageRelatedData(
   coreUsageFields: MicrodollarUsage,
   metadataFields: UsageMetaData,
   prior_microdollar_usage: number,
   posthog_distinct_id: string | null
+): Promise<UsageRecordInsertResult | null> {
+  if (isUSRegion()) {
+    const outcome = await recordUsageInPrimaryRegion({
+      core: coreUsageFields,
+      metadata: metadataFields,
+      prior_microdollar_usage,
+      posthog_distinct_id,
+    });
+    // On `unavailable` fall through to the local write. It is slow from here,
+    // but a slow billing record beats a lost one. `recordUsageInPrimaryRegion`
+    // has already reported the failure.
+    if (outcome.kind === 'ok') return outcome.result;
+  }
+
+  return saveUsageRelatedDataLocally(
+    coreUsageFields,
+    metadataFields,
+    prior_microdollar_usage,
+    posthog_distinct_id
+  );
+}
+
+/**
+ * The write itself, always executed against the primary from wherever it runs.
+ * Exported so `POST /api/internal/usage/record` can invoke it in Frankfurt.
+ */
+export async function saveUsageRelatedDataLocally(
+  coreUsageFields: MicrodollarUsage,
+  metadataFields: UsageMetaData,
+  prior_microdollar_usage: number,
+  posthog_distinct_id: string | null
 ): Promise<UsageRecordInsertResult | null> {
   const isFirst = await isFirstUsage(coreUsageFields, prior_microdollar_usage);
   if (isFirst && posthog_distinct_id)

apps/web/src/lib/ai-gateway/usage-record-client.test.ts

@@ -0,0 +1,179 @@
+// Do not import `jest` from '@jest/globals' in this file. Doing so shadows the
+// global `jest` that the SWC transform needs in order to hoist these
+// `jest.mock` calls, and the mocks then silently fail to apply. The rest of the
+// repository follows the same convention: import the assertion helpers only.
+jest.mock('@/lib/config.server', () => ({
+  ...jest.requireActual<object>('@/lib/config.server'),
+  INTERNAL_API_SECRET: 'test-internal-secret',
+}));
+
+jest.mock('@/lib/constants', () => ({
+  ...jest.requireActual<object>('@/lib/constants'),
+  APP_URL: 'https://app.example.com',
+}));
+
+// The mock function is created inside the factory rather than captured from
+// module scope: jest.mock is hoisted above const declarations, so referencing an
+// outer const here throws "Cannot access before initialization" and silently
+// leaves the real module in place. Retrieve it with jest.requireMock instead.
+jest.mock('@sentry/nextjs', () => ({
+  ...jest.requireActual<object>('@sentry/nextjs'),
+  captureException: jest.fn(),
+}));
+
+import { beforeEach, describe, expect, test } from '@jest/globals';
+import type * as SentryNextjs from '@sentry/nextjs';
+
+import { recordUsageInPrimaryRegion } from './usage-record-client';
+import type { UsageRecordRequest } from './usage-record-contract';
+
+const mockedCaptureException = jest.requireMock<typeof SentryNextjs>('@sentry/nextjs')
+  .captureException as unknown as jest.Mock;
+
+const mockFetch = jest.fn();
+global.fetch = mockFetch as unknown as typeof fetch;
+
+// Deliberately not faking timers or stubbing setTimeout here. The shared Jest
+// worker setup opens a real pg pool, and pg-pool arms its connection timeout with
+// setTimeout — a global stub that invokes callbacks synchronously fires that
+// timeout immediately and closes the client. The client's real backoff is a few
+// hundred milliseconds, which is cheap enough to just wait out.
+
+const payload = {
+  core: { id: 'usage-1' },
+  metadata: {},
+  prior_microdollar_usage: 0,
+  posthog_distinct_id: null,
+} as unknown as UsageRecordRequest;
+
+function jsonResponse(body: unknown, status = 200) {
+  return { ok: status >= 200 && status < 300, status, json: async () => body };
+}
+
+const recorded = {
+  status: 'recorded',
+  result: { usageId: 'usage-1', createdAt: '2026-08-05T10:11:12.945Z', newMicrodollarsUsed: 42 },
+};
+
+beforeEach(() => {
+  mockFetch.mockReset();
+  mockedCaptureException.mockClear();
+});
+
+describe('recordUsageInPrimaryRegion', () => {
+  test('posts to the Frankfurt endpoint with the internal secret', async () => {
+    mockFetch.mockResolvedValue(jsonResponse(recorded));
+
+    const outcome = await recordUsageInPrimaryRegion(payload);
+
+    expect(outcome).toEqual({ kind: 'ok', result: recorded.result });
+    expect(mockFetch).toHaveBeenCalledTimes(1);
+    const [url, init] = mockFetch.mock.calls[0] as [string, RequestInit];
+    expect(url).toBe('https://app.example.com/api/internal/usage/record');
+    expect(init.method).toBe('POST');
+    expect((init.headers as Record<string, string>)['x-internal-api-key']).toBe(
+      'test-internal-secret'
+    );
+  });
+
+  test('treats a duplicate as success so a redelivered write is not reported as lost', async () => {
+    mockFetch.mockResolvedValue(
+      jsonResponse({
+        status: 'duplicate',
+        result: {
+          usageId: 'usage-1',
+          createdAt: '2026-08-05T10:11:12.945Z',
+          newMicrodollarsUsed: null,
+        },
+      })
+    );
+
+    const outcome = await recordUsageInPrimaryRegion(payload);
+
+    expect(outcome.kind).toBe('ok');
+    expect(mockFetch).toHaveBeenCalledTimes(1);
+  });
+
+  test('propagates a deliberate not_recorded outcome without retrying', async () => {
+    mockFetch.mockResolvedValue(jsonResponse({ status: 'not_recorded', result: null }));
+
+    const outcome = await recordUsageInPrimaryRegion(payload);
+
+    expect(outcome).toEqual({ kind: 'ok', result: null });
+    expect(mockFetch).toHaveBeenCalledTimes(1);
+  });
+
+  test('retries a 5xx and succeeds', async () => {
+    mockFetch
+      .mockResolvedValueOnce(jsonResponse({ error: 'boom' }, 503))
+      .mockResolvedValueOnce(jsonResponse(recorded));
+
+    const outcome = await recordUsageInPrimaryRegion(payload);
+
+    expect(outcome.kind).toBe('ok');
+    expect(mockFetch).toHaveBeenCalledTimes(2);
+  });
+
+  test('retries a network failure and succeeds', async () => {
+    mockFetch
+      .mockRejectedValueOnce(new Error('socket hang up'))
+      .mockResolvedValueOnce(jsonResponse(recorded));
+
+    const outcome = await recordUsageInPrimaryRegion(payload);
+
+    expect(outcome.kind).toBe('ok');
+    expect(mockFetch).toHaveBeenCalledTimes(2);
+  });
+
+  test('does not retry a 400, because the payload cannot become valid', async () => {
+    mockFetch.mockResolvedValue(jsonResponse({ error: 'Invalid body' }, 400));
+
+    const outcome = await recordUsageInPrimaryRegion(payload);
+
+    expect(outcome).toEqual({ kind: 'unavailable', reason: 'http_400' });
+    expect(mockFetch).toHaveBeenCalledTimes(1);
+  });
+
+  test('does not retry a 401', async () => {
+    mockFetch.mockResolvedValue(jsonResponse({ error: 'Unauthorized' }, 401));
+
+    const outcome = await recordUsageInPrimaryRegion(payload);
+
+    expect(outcome).toEqual({ kind: 'unavailable', reason: 'http_401' });
+    expect(mockFetch).toHaveBeenCalledTimes(1);
+  });
+
+  test('stops after three attempts and reports unavailable', async () => {
+    mockFetch.mockResolvedValue(jsonResponse({ error: 'boom' }, 500));
+
+    const outcome = await recordUsageInPrimaryRegion(payload);
+
+    expect(outcome.kind).toBe('unavailable');
+    expect(mockFetch).toHaveBeenCalledTimes(3);
+    expect(mockedCaptureException).toHaveBeenCalledTimes(1);
+  });
+
+  // A malformed response may mean the write committed, so retrying risks nothing
+  // (the receiver dedupes) but is pointless; the caller must reconcile instead.
+  test('does not retry a malformed response body', async () => {
+    mockFetch.mockResolvedValue(jsonResponse({ status: 'nonsense' }));
+
+    const outcome = await recordUsageInPrimaryRegion(payload);
+
+    expect(outcome).toEqual({ kind: 'unavailable', reason: 'malformed_response' });
+    expect(mockFetch).toHaveBeenCalledTimes(1);
+  });
+
+  test('reports the failure to Sentry with the usage id for reconciliation', async () => {
+    mockFetch.mockResolvedValue(jsonResponse({ error: 'boom' }, 500));
+
+    await recordUsageInPrimaryRegion(payload);
+
+    const [, options] = mockedCaptureException.mock.calls[0] as [
+      unknown,
+      { tags: Record<string, string>; extra: Record<string, unknown> },
+    ];
+    expect(options.tags.source).toBe('recordUsageInPrimaryRegion');
+    expect(options.extra.usageId).toBe('usage-1');
+  });
+});