From 0d4d682e4c5a710265e1cbdcc9345a07f9ed463b Mon Sep 17 00:00:00 2001 From: Lawrence Chen <54008264+lawrencecchen@users.noreply.github.com> Date: Wed, 23 Sep 2026 05:33:36 -0700 Subject: [PATCH] coderouter: serve stored usage in account status GET /api/coderouter/accounts decrypted a credential and called ChatGPT's usage endpoint for every active Codex account on every request, so status took 1.4-3.9 s. Readings are now stored on coderouter_accounts. A reading under 30 s old is served as is; an older one is served and refreshed after the response by one claiming instance; one over 10 min old, or a missing one, is read before answering. The refresh still sets cooldowns from limit_reached, which is what routing reads. --- .../migration.sql | 5 + web/db/schema.ts | 10 + web/services/coderouter/repository.ts | 74 ++++- web/services/coderouter/usage.ts | 292 ++++++++++++------ web/tests/coderouter-account-usage.test.ts | 111 +++++++ 5 files changed, 405 insertions(+), 87 deletions(-) create mode 100644 web/db/migrations/20260923120000_coderouter_account_usage_cache/migration.sql create mode 100644 web/tests/coderouter-account-usage.test.ts diff --git a/web/db/migrations/20260923120000_coderouter_account_usage_cache/migration.sql b/web/db/migrations/20260923120000_coderouter_account_usage_cache/migration.sql new file mode 100644 index 000000000000..203eefe9f24f --- /dev/null +++ b/web/db/migrations/20260923120000_coderouter_account_usage_cache/migration.sql @@ -0,0 +1,5 @@ +ALTER TABLE "coderouter_accounts" + ADD COLUMN "usage" jsonb, + ADD COLUMN "usage_error" text, + ADD COLUMN "usage_fetched_at" timestamp with time zone, + ADD COLUMN "usage_refresh_claimed_at" timestamp with time zone; diff --git a/web/db/schema.ts b/web/db/schema.ts index dc7579662034..b0c2f20d59e9 100644 --- a/web/db/schema.ts +++ b/web/db/schema.ts @@ -1210,6 +1210,16 @@ export const coderouterAccounts = pgTable( cooldownUntil: timestamp("cooldown_until", { withTimezone: true }), lastUsedAt: timestamp("last_used_at", { withTimezone: true }), lastFailureCode: text("last_failure_code"), + /** + * Last provider quota read for `cr` status. Display data: routing uses + * `cooldown_until`. Served while fresh so status does not call the + * provider once per account per request. + */ + usage: jsonb("usage").$type(), + usageError: text("usage_error"), + usageFetchedAt: timestamp("usage_fetched_at", { withTimezone: true }), + /** One web instance refreshes a stale reading at a time. */ + usageRefreshClaimedAt: timestamp("usage_refresh_claimed_at", { withTimezone: true }), createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(), updatedAt: timestamp("updated_at", { withTimezone: true }).notNull().defaultNow(), }, diff --git a/web/services/coderouter/repository.ts b/web/services/coderouter/repository.ts index 8905dc94f0ef..0da362adc3a3 100644 --- a/web/services/coderouter/repository.ts +++ b/web/services/coderouter/repository.ts @@ -1,6 +1,6 @@ import { grantVmImportedAccount } from "./vmAccountImport"; import { createHash, randomBytes, randomUUID } from "node:crypto"; -import { and, eq, gt, isNotNull, isNull, lt, lte, or, sql } from "drizzle-orm"; +import { and, eq, gt, inArray, isNotNull, isNull, lt, lte, or, sql } from "drizzle-orm"; import { cloudDb } from "../../db/client"; import { runWithCloudDbQuerySignal } from "../../db/queryScope"; import { @@ -642,6 +642,78 @@ export async function listCoderouterTeamIds(): Promise { .then((rows) => rows.map((row) => row.teamId)); } +export type StoredAccountUsage = { + readonly usage: unknown; + readonly usageError: string | null; + readonly fetchedAt: Date; +}; + +/** The last quota reading per account, for accounts that have one. */ +export async function storedAccountUsage( + teamId: string, +): Promise> { + const rows = await cloudDb() + .select({ + id: coderouterAccounts.id, + usage: coderouterAccounts.usage, + usageError: coderouterAccounts.usageError, + fetchedAt: coderouterAccounts.usageFetchedAt, + }) + .from(coderouterAccounts) + .where(and(eq(coderouterAccounts.teamId, teamId), isNotNull(coderouterAccounts.usageFetchedAt))); + return new Map(rows.flatMap((row) => row.fetchedAt + ? [[row.id, { usage: row.usage ?? null, usageError: row.usageError, fetchedAt: row.fetchedAt }] as const] + : [])); +} + +/** + * Claims the stale readings this instance will refresh. An account whose + * refresh another instance claimed after `claimedBefore` is skipped, so + * concurrent status views issue one provider read per account. + */ +export async function claimAccountUsageRefresh( + teamId: string, + accountIds: readonly string[], + now: Date, + claimedBefore: Date, +): Promise { + if (accountIds.length === 0) return []; + const rows = await cloudDb() + .update(coderouterAccounts) + .set({ usageRefreshClaimedAt: now }) + .where(and( + eq(coderouterAccounts.teamId, teamId), + inArray(coderouterAccounts.id, [...accountIds]), + or( + isNull(coderouterAccounts.usageRefreshClaimedAt), + lt(coderouterAccounts.usageRefreshClaimedAt, claimedBefore), + ), + )) + .returning({ id: coderouterAccounts.id }); + return rows.map((row) => row.id); +} + +/** Records a quota reading. An older reading never replaces a newer one. */ +export async function storeAccountUsage( + teamId: string, + accountId: string, + reading: { readonly usage: unknown } | { readonly usageError: string }, + fetchedAt: Date, +): Promise { + await cloudDb() + .update(coderouterAccounts) + .set({ + usage: "usage" in reading ? reading.usage : null, + usageError: "usageError" in reading ? reading.usageError : null, + usageFetchedAt: fetchedAt, + }) + .where(and( + eq(coderouterAccounts.teamId, teamId), + eq(coderouterAccounts.id, accountId), + or(isNull(coderouterAccounts.usageFetchedAt), lt(coderouterAccounts.usageFetchedAt, fetchedAt)), + )); +} + export async function listEncryptedCredentials( teamId: string, ): Promise { diff --git a/web/services/coderouter/usage.ts b/web/services/coderouter/usage.ts index f3102a70c5a6..eb8be9ca776e 100644 --- a/web/services/coderouter/usage.ts +++ b/web/services/coderouter/usage.ts @@ -1,106 +1,226 @@ +import { after } from "next/server"; import type { CoderouterAccountAccess } from "./accountAccess"; +import type { EncryptedCredential } from "./encryption"; import { + claimAccountUsageRefresh, listAccounts, listEncryptedCredentials, markAccountCooldown, + storeAccountUsage, + storedAccountUsage, + type StoredAccountUsage, } from "./repository"; import { freshCredential } from "./refresh"; import { fetchProviderRead } from "./providerFetch"; import { addCoderouterBreadcrumb, reportCoderouterFailure } from "./observability"; const CODEX_USAGE_URL = "https://chatgpt.com/backend-api/wham/usage"; -const usageRequests = new Map< - string, - Promise>> ->(); - -export async function accountsWithUsage(teamId: string, access?: CoderouterAccountAccess) { - const key = JSON.stringify([teamId, access]); - const pending = usageRequests.get(key); - if (pending) return await pending; - - // Provider reads fan out in parallel. Coalesce only requests that are - // concurrently in flight; completed quota data is never served from cache. - const request = loadAccountsWithUsage(teamId, access); - usageRequests.set(key, request); - try { - return await request; - } finally { - usageRequests.delete(key); - } -} -async function loadAccountsWithUsage(teamId: string, access?: CoderouterAccountAccess) { - const startedAt = performance.now(); - addCoderouterBreadcrumb("status", "Loading account usage"); - // Account metadata and encrypted envelopes are independent RDS reads. - const rdsStartedAt = performance.now(); - const [accounts, credentials] = await Promise.all([ - listAccounts(teamId, access), - listEncryptedCredentials(teamId), - ]); - const rdsMs = performance.now() - rdsStartedAt; - const credentialsByAccount = new Map( - credentials.map((credential) => [credential.accountId, credential]), - ); - const providerStartedAt = performance.now(); - const withUsage = await Promise.all(accounts.map(async (account) => { - if (account.provider !== "codex" || account.state !== "active") { - return account; - } +/** + * Status serves each account's last quota reading instead of calling the + * provider once per account per request, which made `cr` wait on the slowest + * of up to ~100 ChatGPT reads. A reading younger than `USAGE_FRESH_MS` is + * served as is. An older one is served and refreshed after the response. One + * older than `USAGE_MAX_STALE_MS`, or a missing one, is read before answering, + * so an occasional `cr` never shows a long-stale quota. Routing does not read + * these numbers; it uses `cooldown_until`, which a refresh still sets. + */ +export const USAGE_FRESH_MS = 30_000; +export const USAGE_MAX_STALE_MS = 10 * 60_000; + +type Account = Awaited>[number]; +type Reading = { readonly usage: unknown } | { readonly usageError: string }; + +export type UsageDependencies = { + readonly listAccounts: typeof listAccounts; + readonly storedUsage: (teamId: string) => Promise>; + readonly claimRefresh: typeof claimAccountUsageRefresh; + readonly credentials: (teamId: string) => Promise; + readonly read: (teamId: string, accountId: string, credential: EncryptedCredential | undefined) => Promise; + readonly store: typeof storeAccountUsage; + /** Runs work after the response; must absorb its rejection. */ + readonly defer: (task: Promise) => void; + readonly now: () => number; +}; + +const defaultDependencies: UsageDependencies = { + listAccounts, + storedUsage: storedAccountUsage, + claimRefresh: claimAccountUsageRefresh, + credentials: listEncryptedCredentials, + read: readCodexUsage, + store: storeAccountUsage, + defer: deferAfterResponse, + now: Date.now, +}; + +export function createAccountsWithUsage(dependencies: UsageDependencies) { + const usageRequests = new Map>(); + + async function accountsWithUsage(teamId: string, access?: CoderouterAccountAccess) { + const key = JSON.stringify([teamId, access]); + const pending = usageRequests.get(key); + if (pending) return await pending; + // Coalesce only requests that are concurrently in flight. + const request = load(teamId, access); + usageRequests.set(key, request); try { - const credential = await freshCredential({ - teamId, - accountId: account.id, - expectedRevision: credentialsByAccount.get(account.id)?.credentialRevision ?? 0, - known: credentialsByAccount.get(account.id), + return await request; + } finally { + usageRequests.delete(key); + } + } + + /** Reads, stores, and returns fresh readings for the given accounts. */ + async function refresh(teamId: string, accountIds: readonly string[]) { + const credentials = new Map( + (await dependencies.credentials(teamId)).map((credential) => [credential.accountId, credential]), + ); + const readings = new Map(); + await Promise.all(accountIds.map(async (accountId) => { + const reading = await dependencies.read(teamId, accountId, credentials.get(accountId)); + const fetchedAt = new Date(dependencies.now()); + readings.set(accountId, { + usage: "usage" in reading ? reading.usage : null, + usageError: "usageError" in reading ? reading.usageError : null, + fetchedAt, }); - if (credential.provider !== "codex") return account; - const response = await fetchProviderRead(() => fetch(CODEX_USAGE_URL, { - headers: { - authorization: `Bearer ${credential.accessToken}`, - "chatgpt-account-id": credential.accountId, - "user-agent": "coderouter/0.2", - }, - cache: "no-store", - signal: AbortSignal.timeout(5_000), - })); - if (!response.ok) { - reportCoderouterFailure( - response.status === 429 ? "provider_rate_limit" : "provider_usage", - new Error("provider usage request failed"), - { provider: account.provider, status: response.status }, - ); - return { ...account, usageError: `HTTP ${response.status}` }; - } - const usage: unknown = await response.json(); - const cooldownMs = usageCooldown(usage); - if (cooldownMs !== null) { - await markAccountCooldown(account.id, cooldownMs); + try { + await dependencies.store(teamId, accountId, reading, fetchedAt); + } catch (error) { + // The reading is still returned; the next view reads again. + reportCoderouterFailure("rds", error, { operation: "store_account_usage" }); } - return { ...account, usage }; - } catch (error) { - reportCoderouterFailure("provider_usage", error, { - provider: account.provider, - }); - return { ...account, usageError: "unavailable" }; + })); + return readings; + } + + async function refreshInBackground(teamId: string, accountIds: readonly string[]) { + const now = dependencies.now(); + const claimed = await dependencies.claimRefresh( + teamId, + accountIds, + new Date(now), + new Date(now - USAGE_FRESH_MS), + ); + if (claimed.length > 0) await refresh(teamId, claimed); + } + + async function load(teamId: string, access?: CoderouterAccountAccess) { + const startedAt = performance.now(); + addCoderouterBreadcrumb("status", "Loading account usage"); + const rdsStartedAt = performance.now(); + const [accounts, stored] = await Promise.all([ + dependencies.listAccounts(teamId, access), + dependencies.storedUsage(teamId), + ]); + const rdsMs = performance.now() - rdsStartedAt; + + const now = dependencies.now(); + const blocking: string[] = []; + const stale: string[] = []; + for (const account of accounts) { + if (!hasUsage(account)) continue; + const reading = stored.get(account.id); + const age = reading ? now - reading.fetchedAt.getTime() : Number.POSITIVE_INFINITY; + if (age > USAGE_MAX_STALE_MS) blocking.push(account.id); + else if (age > USAGE_FRESH_MS) stale.push(account.id); } - })); - addCoderouterBreadcrumb("status", "Provider usage fanout completed", { - account_count: accounts.length, - provider_ms: Math.round(performance.now() - providerStartedAt), + + const providerStartedAt = performance.now(); + const live = blocking.length > 0 ? await refresh(teamId, blocking) : new Map(); + const providerMs = performance.now() - providerStartedAt; + if (stale.length > 0) dependencies.defer(refreshInBackground(teamId, stale)); + addCoderouterBreadcrumb("status", "Account usage resolved", { + account_count: accounts.length, + provider_read_count: blocking.length, + background_refresh_count: stale.length, + provider_ms: Math.round(providerMs), + }); + + let oldest = now; + const withUsage = accounts.map((account) => { + if (!hasUsage(account)) return account; + const reading = live.get(account.id) ?? stored.get(account.id); + if (!reading) return account; + oldest = Math.min(oldest, reading.fetchedAt.getTime()); + return reading.usageError === null + ? { ...account, usage: reading.usage } + : { ...account, usageError: reading.usageError }; + }); + return { + accounts: withUsage, + usageAsOf: new Date(oldest).toISOString(), + usageGeneratedAtMs: oldest, + cacheMaxAgeSeconds: USAGE_FRESH_MS / 1_000, + timing: { + rdsMs, + providerMs, + totalMs: performance.now() - startedAt, + }, + }; + } + + return { accountsWithUsage }; +} + +export const { accountsWithUsage } = createAccountsWithUsage(defaultDependencies); + +function hasUsage(account: Account): boolean { + return account.provider === "codex" && account.state === "active"; +} + +async function readCodexUsage( + teamId: string, + accountId: string, + known: EncryptedCredential | undefined, +): Promise { + try { + const credential = await freshCredential({ + teamId, + accountId, + expectedRevision: known?.credentialRevision ?? 0, + known, + }); + if (credential.provider !== "codex") return { usageError: "unavailable" }; + const response = await fetchProviderRead(() => fetch(CODEX_USAGE_URL, { + headers: { + authorization: `Bearer ${credential.accessToken}`, + "chatgpt-account-id": credential.accountId, + "user-agent": "coderouter/0.2", + }, + cache: "no-store", + signal: AbortSignal.timeout(5_000), + })); + if (!response.ok) { + reportCoderouterFailure( + response.status === 429 ? "provider_rate_limit" : "provider_usage", + new Error("provider usage request failed"), + { provider: "codex", status: response.status }, + ); + return { usageError: `HTTP ${response.status}` }; + } + const usage: unknown = await response.json(); + const cooldownMs = usageCooldown(usage); + if (cooldownMs !== null) { + await markAccountCooldown(accountId, cooldownMs); + } + return { usage }; + } catch (error) { + reportCoderouterFailure("provider_usage", error, { provider: "codex" }); + return { usageError: "unavailable" }; + } +} + +function deferAfterResponse(task: Promise): void { + const settled = task.catch((error: unknown) => { + reportCoderouterFailure("provider_usage", error, { operation: "background_usage_refresh" }); }); - return { - accounts: withUsage, - usageAsOf: new Date().toISOString(), - usageGeneratedAtMs: Date.now(), - cacheMaxAgeSeconds: 0, - timing: { - rdsMs, - providerMs: performance.now() - providerStartedAt, - totalMs: performance.now() - startedAt, - }, - }; + try { + after(settled); + } catch { + // Unit tests and scripts have no Next request scope; the task already runs. + } } function usageCooldown(value: unknown): number | null { diff --git a/web/tests/coderouter-account-usage.test.ts b/web/tests/coderouter-account-usage.test.ts new file mode 100644 index 000000000000..3dba7a53955a --- /dev/null +++ b/web/tests/coderouter-account-usage.test.ts @@ -0,0 +1,111 @@ +import { describe, expect, mock, test } from "bun:test"; + +import { + USAGE_FRESH_MS, + USAGE_MAX_STALE_MS, + createAccountsWithUsage, + type UsageDependencies, +} from "../services/coderouter/usage"; + +const NOW = Date.parse("2026-09-23T12:00:00.000Z"); + +function account(id: string, overrides: Record = {}) { + return { id, provider: "codex", state: "active", label: `${id}@example.com`, ...overrides }; +} + +function reading(ageMs: number, usage: unknown = { plan_type: "pro" }, usageError: string | null = null) { + return { usage, usageError, fetchedAt: new Date(NOW - ageMs) }; +} + +function harness(options: { + readonly accounts: readonly Record[]; + readonly stored?: ReadonlyMap>; + readonly claimed?: (ids: readonly string[]) => readonly string[]; +}) { + const deferred: Promise[] = []; + const read = mock(async (_teamId: string, accountId: string) => ({ usage: { live: accountId } })); + const store = mock(async () => undefined); + const claimRefresh = mock(async (_teamId: string, ids: readonly string[]) => options.claimed?.(ids) ?? ids); + const dependencies: UsageDependencies = { + listAccounts: (async () => options.accounts) as never, + storedUsage: async () => options.stored ?? new Map(), + claimRefresh: claimRefresh as never, + credentials: async () => [], + read: read as never, + store: store as never, + defer: (task) => { deferred.push(task); }, + now: () => NOW, + }; + const { accountsWithUsage } = createAccountsWithUsage(dependencies); + return { accountsWithUsage, read, store, claimRefresh, deferred }; +} + +describe("coderouter status usage", () => { + test("serves fresh readings without a provider read", async () => { + const { accountsWithUsage, read, deferred } = harness({ + accounts: [account("a")], + stored: new Map([["a", reading(5_000)]]), + }); + const result = await accountsWithUsage("team-1"); + expect(read).not.toHaveBeenCalled(); + expect(deferred).toHaveLength(0); + expect(result.accounts[0]).toMatchObject({ id: "a", usage: { plan_type: "pro" } }); + expect(result.usageGeneratedAtMs).toBe(NOW - 5_000); + expect(result.cacheMaxAgeSeconds).toBe(USAGE_FRESH_MS / 1_000); + }); + + test("serves a stale reading at once and refreshes the claimed ones after the response", async () => { + const { accountsWithUsage, read, store, claimRefresh, deferred } = harness({ + accounts: [account("a"), account("b")], + stored: new Map([["a", reading(USAGE_FRESH_MS + 1)], ["b", reading(USAGE_FRESH_MS + 1)]]), + claimed: () => ["a"], + }); + const result = await accountsWithUsage("team-1"); + expect(result.accounts.map((row) => (row as { usage?: unknown }).usage)).toEqual([ + { plan_type: "pro" }, + { plan_type: "pro" }, + ]); + expect(deferred).toHaveLength(1); + await Promise.all(deferred); + expect(claimRefresh).toHaveBeenCalledWith("team-1", ["a", "b"], new Date(NOW), new Date(NOW - USAGE_FRESH_MS)); + // Another instance holds b's claim. + expect(read.mock.calls.map((call) => call[1])).toEqual(["a"]); + expect(store).toHaveBeenCalledWith("team-1", "a", { usage: { live: "a" } }, new Date(NOW)); + }); + + test("reads missing and too-old readings before answering", async () => { + const { accountsWithUsage, read, store, deferred } = harness({ + accounts: [account("missing"), account("old")], + stored: new Map([["old", reading(USAGE_MAX_STALE_MS + 1)]]), + }); + const result = await accountsWithUsage("team-1"); + expect(read.mock.calls.map((call) => call[1]).sort()).toEqual(["missing", "old"]); + expect(store).toHaveBeenCalledTimes(2); + expect(deferred).toHaveLength(0); + expect(result.accounts).toEqual([ + expect.objectContaining({ id: "missing", usage: { live: "missing" } }), + expect.objectContaining({ id: "old", usage: { live: "old" } }), + ]); + expect(result.usageGeneratedAtMs).toBe(NOW); + }); + + test("keeps a stored provider error and skips inactive and non-Codex accounts", async () => { + const { accountsWithUsage, read } = harness({ + accounts: [account("failed"), account("broken", { state: "broken" }), account("go", { provider: "opencode-go" })], + stored: new Map([["failed", reading(1_000, null, "HTTP 401")]]), + }); + const result = await accountsWithUsage("team-1"); + expect(read).not.toHaveBeenCalled(); + expect(result.accounts[0]).toMatchObject({ id: "failed", usageError: "HTTP 401" }); + expect(result.accounts[0]).not.toHaveProperty("usage"); + expect(result.accounts[1]).not.toHaveProperty("usage"); + expect(result.accounts[2]).not.toHaveProperty("usage"); + }); + + test("a failed store still returns the live reading", async () => { + const h = harness({ accounts: [account("a")] }); + h.store.mockImplementation(async () => { throw new Error("db down"); }); + const result = await h.accountsWithUsage("team-1"); + expect(result.accounts[0]).toMatchObject({ usage: { live: "a" } }); + }); +});