diff --git a/apps/server/src/services/skills/builtin-skills/bb-cli/SKILL.md b/apps/server/src/services/skills/builtin-skills/bb-cli/SKILL.md index 6e53da9b31..b7e1bdabe7 100644 --- a/apps/server/src/services/skills/builtin-skills/bb-cli/SKILL.md +++ b/apps/server/src/services/skills/builtin-skills/bb-cli/SKILL.md @@ -85,9 +85,13 @@ prior token valid for ten minutes. Agents should pipe API keys to `--api-key-stdin`; `--api-key ` is an unsafe compatibility form that exposes the key in process arguments, shell history, and agent transcripts. Prefer `--import` for -an existing Claude Code login. JSON account status -reports rejected upstream bucket resets under `bucketExhaustion`; this -diagnostic field does not affect account selection. +an existing Claude Code login. OAuth quota refreshes on add or enable and every +five minutes while an account is idle. Account tables add columns for observed +model-family buckets; JSON status exposes their utilization, reset, status, +observation time, and source under `familyWeekly`. Selection skips an account +whose requested family is spent while retaining it for other families. A +present `metadata.user_id` account UUID is aligned with the selected OAuth +account. - Inspect real status, logs, API results, or diffs instead of assumptions. - Keep file paths on the machine that owns the selected workspace. diff --git a/docs/configuration.md b/docs/configuration.md index eb382fb50f..62986cbbb0 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -658,11 +658,18 @@ for ten minutes so in-flight requests can drain. Bypass or restore routing for one thread with `bb pool bypass ` or `bb pool bypass --off`. Account listing, enable, disable, and removal are available through `bb pool account list|enable|disable|remove`. -JSON account status includes rejected upstream bucket resets under -`bucketExhaustion`. The field is diagnostic and does not affect selection. - -Two settings control routing. `switchThreshold` is the 5-hour or 7-day quota -fraction at which an account stops receiving traffic and defaults to `0.98`. +OAuth accounts refresh quota from Anthropic's usage endpoint when added or +enabled and every five minutes while idle. `account list` adds columns for the +family buckets Anthropic reports; JSON status exposes their utilization, +reset, status, observation time, and `header` or `usage` source under +`familyWeekly`. Requests route around an account spent for their model family +without disabling that account for other families. Imported and newly signed-in +accounts retain their Anthropic account UUID, and the hub aligns a present +`metadata.user_id` account component with the selected account. + +Two settings control routing. `switchThreshold` is the shared or requested +model-family quota fraction at which an account stops receiving matching +traffic and defaults to `0.98`. `upstreamBaseUrl` defaults to `https://api.anthropic.com` and exists only for tests and QA with a controlled fake upstream: diff --git a/packages/templates/src/templates/bb-guide-plugins.md b/packages/templates/src/templates/bb-guide-plugins.md index b7ce0ac740..09a1f48c0c 100644 --- a/packages/templates/src/templates/bb-guide-plugins.md +++ b/packages/templates/src/templates/bb-guide-plugins.md @@ -62,9 +62,13 @@ safely. Rotation keeps the prior token valid for ten minutes. Agents should use `--api-key-stdin`, which reads exactly one non-empty key from piped standard input. The compatibility form `--api-key ` exposes the key in process arguments, shell history, and agent transcripts. Prefer `--import` when Claude -Code is already signed in. JSON account status reports rejected upstream -bucket resets under `bucketExhaustion`; this is diagnostic status and does not -affect selection. +Code is already signed in. OAuth quota refreshes on add or enable and every +five minutes while the account is idle. Account tables add columns for the +family buckets Anthropic reports, and JSON status exposes the same observations +under `familyWeekly`. Selection skips an account only for a spent requested +family while retaining it for other families. When Claude Code supplies an +account UUID in `metadata.user_id`, the hub aligns it with the selected OAuth +account. The `upstreamBaseUrl` setting exists for tests and QA and defaults to `https://api.anthropic.com`; `switchThreshold` defaults to `0.98`. diff --git a/plugins/account-pool/app.test.tsx b/plugins/account-pool/app.test.tsx index ea610ba353..71b9297339 100644 --- a/plugins/account-pool/app.test.tsx +++ b/plugins/account-pool/app.test.tsx @@ -15,6 +15,7 @@ function account(): AccountSummary { kind: "oauth", label: "Personal Claude", email: "person@example.com", + accountUuid: null, subscriptionType: "max", rateLimitTier: "default_claude_max_5x", enabled: true, @@ -27,7 +28,13 @@ function account(): AccountSummary { sevenDayResetAt: null, sevenDayStatus: null, representativeClaim: null, - bucketExhaustion: {}, + familyWeekly: { + fable: null, + sonnet: null, + opus: null, + haiku: null, + other: null, + }, observedAt: 1, heldUntil: null, error: null, diff --git a/plugins/account-pool/src/cli.ts b/plugins/account-pool/src/cli.ts index 3a13487160..539f1e785f 100644 --- a/plugins/account-pool/src/cli.ts +++ b/plugins/account-pool/src/cli.ts @@ -6,6 +6,8 @@ import { loginCompleteInputSchema, tokenRotateInputSchema, type AccountSummary, + type FamilyQuota, + type ModelFamily, type PoolStatus, } from "./contracts.js"; import type { PoolOperations } from "./operations.js"; @@ -72,10 +74,47 @@ function formatUtilization(value: number | null): string { return value === null ? "-" : `${Math.round(value * 100)}%`; } +const MODEL_FAMILIES: ModelFamily[] = [ + "fable", + "sonnet", + "opus", + "haiku", + "other", +]; + +function familyLabel(family: ModelFamily): string { + return family[0]?.toUpperCase() + family.slice(1); +} + +function formatFamilyQuota(quota: FamilyQuota | null): string { + if (quota === null) return "-"; + return [ + formatUtilization(quota.utilization), + quota.status ?? "-", + formatReset(quota.resetAt), + quota.source, + ].join(" "); +} + function formatAccounts(accounts: readonly AccountSummary[]): string { if (accounts.length === 0) return "No accounts configured."; + const families = MODEL_FAMILIES.filter((family) => + accounts.some((account) => account.familyWeekly[family] !== null), + ); return [ - "ID\tLabel\tKind\tEnabled\tPriority\t5h\t5h reset\t7d\t7d reset\tStatus", + [ + "ID", + "Label", + "Kind", + "Enabled", + "Priority", + "5h", + "5h reset", + "7d", + "7d reset", + ...families.map(familyLabel), + "Status", + ].join("\t"), ...accounts.map((account) => [ account.id, @@ -87,6 +126,9 @@ function formatAccounts(accounts: readonly AccountSummary[]): string { formatReset(account.fiveHourResetAt), formatUtilization(account.sevenDayUtilization), formatReset(account.sevenDayResetAt), + ...families.map((family) => + formatFamilyQuota(account.familyWeekly[family]), + ), account.status, ].join("\t"), ), diff --git a/plugins/account-pool/src/contracts.ts b/plugins/account-pool/src/contracts.ts index 5f9e952a8a..accdeac026 100644 --- a/plugins/account-pool/src/contracts.ts +++ b/plugins/account-pool/src/contracts.ts @@ -2,6 +2,39 @@ import { z } from "zod"; export const providerSchema = z.literal("claude"); export const accountKindSchema = z.enum(["oauth", "api-key"]); +export const modelFamilySchema = z.enum([ + "fable", + "sonnet", + "opus", + "haiku", + "other", +]); + +export type ModelFamily = z.infer; + +export const familyQuotaSchema = z + .object({ + utilization: z.number().nullable(), + resetAt: z.number().int().nullable(), + status: z.string().nullable(), + observedAt: z.number().int(), + source: z.enum(["header", "usage"]), + }) + .strict(); + +export type FamilyQuota = z.infer; + +export const familyWeeklySchema = z + .object({ + fable: familyQuotaSchema.nullable(), + sonnet: familyQuotaSchema.nullable(), + opus: familyQuotaSchema.nullable(), + haiku: familyQuotaSchema.nullable(), + other: familyQuotaSchema.nullable(), + }) + .strict(); + +export type FamilyWeekly = z.infer; export const accountSchema = z .object({ @@ -10,6 +43,7 @@ export const accountSchema = z kind: accountKindSchema, label: z.string().min(1), email: z.string().email().nullable(), + accountUuid: z.string().uuid().nullable().default(null), subscriptionType: z.string().nullable(), rateLimitTier: z.string().nullable(), enabled: z.boolean(), @@ -53,7 +87,7 @@ export const quotaSchema = z sevenDayResetAt: z.number().int().nullable(), sevenDayStatus: z.string().nullable(), representativeClaim: z.string().nullable(), - bucketExhaustion: z.record(z.string(), z.number().int()), + familyWeekly: familyWeeklySchema, observedAt: z.number().int().nullable(), heldUntil: z.number().int().nullable(), error: z.string().nullable(), @@ -70,7 +104,7 @@ export const accountSummarySchema = accountSchema.extend({ sevenDayResetAt: z.number().int().nullable(), sevenDayStatus: z.string().nullable(), representativeClaim: z.string().nullable(), - bucketExhaustion: z.record(z.string(), z.number().int()), + familyWeekly: familyWeeklySchema, observedAt: z.number().int().nullable(), heldUntil: z.number().int().nullable(), error: z.string().nullable(), diff --git a/plugins/account-pool/src/credentials.ts b/plugins/account-pool/src/credentials.ts index 9abda506be..be9475a756 100644 --- a/plugins/account-pool/src/credentials.ts +++ b/plugins/account-pool/src/credentials.ts @@ -24,7 +24,10 @@ const credentialsFileSchema = z const accountFileSchema = z.object({ oauthAccount: z - .object({ emailAddress: z.string().email().nullish() }) + .object({ + emailAddress: z.string().email().nullish(), + accountUuid: z.string().uuid().nullish(), + }) .nullish(), }); @@ -35,11 +38,12 @@ export interface ImportedClaudeCredentials { subscriptionType: string | null; rateLimitTier: string | null; email: string | null; + accountUuid: string | null; } function parseCredentials( raw: string, -): Omit | null { +): Omit | null { const trimmed = raw.trim(); const candidates = [trimmed]; if (/^(?:[0-9a-f]{2})+$/iu.test(trimmed)) { @@ -77,17 +81,23 @@ async function readKeychainCredentials(): Promise { return null; } -async function readAccountEmail(): Promise { +async function readAccountIdentity(): Promise<{ + email: string | null; + accountUuid: string | null; +}> { try { const value = JSON.parse( await fs.readFile(path.join(os.homedir(), ".claude.json"), "utf8"), ); const parsed = accountFileSchema.safeParse(value); return parsed.success - ? (parsed.data.oauthAccount?.emailAddress ?? null) - : null; + ? { + email: parsed.data.oauthAccount?.emailAddress ?? null, + accountUuid: parsed.data.oauthAccount?.accountUuid ?? null, + } + : { email: null, accountUuid: null }; } catch { - return null; + return { email: null, accountUuid: null }; } } @@ -109,5 +119,5 @@ export async function importClaudeCredentials(): Promise number; refreshUrl: string; + usageUrl: string; + profileUrl: string; + usageRefreshIntervalMs: number; drainTimeoutMs: number; } @@ -83,16 +104,114 @@ export class AccountPoolHub { private readonly inFlightByAccount = new Map(); private readonly activeControllers = new Set(); private readonly refreshes = new Map>(); + private readonly usageRefreshes = new Map>(); + private readonly lastUsageRefreshAt = new Map(); private readonly drainWaiters = new Set<() => void>(); constructor(private readonly options: HubOptions) {} async start(signal: AbortSignal): Promise { this.accepting = true; - await waitForAbort(signal); + while (!signal.aborted) { + await this.refreshUsage(); + await waitForDelay(this.options.usageRefreshIntervalMs, signal); + } await this.stop(); } + async refreshUsage(accountId?: string, force = false): Promise { + const accounts = (await this.options.accounts.list()).filter( + (account) => + account.enabled && + account.kind === "oauth" && + (accountId === undefined || account.id === accountId), + ); + await Promise.all( + accounts.map((account) => this.refreshAccountUsage(account, force)), + ); + } + + private async refreshAccountUsage( + account: Account, + force: boolean, + ): Promise { + if ((this.inFlightByAccount.get(account.id) ?? 0) > 0) return; + const now = this.options.now(); + const lastRefreshAt = this.lastUsageRefreshAt.get(account.id); + if ( + !force && + lastRefreshAt !== undefined && + now - lastRefreshAt < this.options.usageRefreshIntervalMs + ) + return; + const running = this.usageRefreshes.get(account.id); + if (running !== undefined) return running; + this.lastUsageRefreshAt.set(account.id, now); + const refresh = this.fetchAccountUsage(account).finally(() => { + this.usageRefreshes.delete(account.id); + }); + this.usageRefreshes.set(account.id, refresh); + return refresh; + } + + private async fetchAccountUsage(account: Account): Promise { + try { + const secret = await this.freshSecret(account); + if (secret.kind !== "oauth") return; + const response = await this.options.fetch(this.options.usageUrl, { + headers: { + authorization: `Bearer ${secret.accessToken}`, + "anthropic-beta": OAUTH_BETA, + accept: "application/json", + }, + signal: AbortSignal.timeout(USAGE_REQUEST_TIMEOUT_MS), + }); + if (response.ok) { + const payload = await response.json().catch(() => null); + if (typeof payload === "object" && payload !== null) { + const quota = quotaFromUsage( + account.id, + payload, + this.options.quotas.get(account.id), + this.options.now(), + ); + if (quota !== null) this.options.quotas.put(quota); + } + } else { + await response.body?.cancel(); + } + if (account.accountUuid === null) { + await this.backfillAccountUuid(account.id, secret.accessToken); + } + } catch {} + } + + private async backfillAccountUuid( + accountId: string, + accessToken: string, + ): Promise { + const response = await this.options.fetch(this.options.profileUrl, { + headers: { + authorization: `Bearer ${accessToken}`, + "anthropic-beta": OAUTH_BETA, + accept: "application/json", + }, + signal: AbortSignal.timeout(USAGE_REQUEST_TIMEOUT_MS), + }); + if (!response.ok) { + await response.body?.cancel(); + return; + } + const payload = await response.json().catch(() => null); + const parsed = profileResponseSchema.safeParse(payload); + const accountUuid = parsed.success + ? (parsed.data.account?.uuid ?? null) + : null; + if (accountUuid !== null) { + await this.options.accounts.setAccountUuid(accountId, accountUuid); + } + } + async stop(): Promise { this.accepting = false; if (this.inFlightCount() === 0) return; @@ -160,7 +279,7 @@ export class AccountPoolHub { sevenDayResetAt: quota.sevenDayResetAt, sevenDayStatus: quota.sevenDayStatus, representativeClaim: quota.representativeClaim, - bucketExhaustion: quota.bucketExhaustion, + familyWeekly: quota.familyWeekly, observedAt: quota.observedAt, heldUntil: quota.heldUntil, error: quota.error, @@ -174,9 +293,11 @@ export class AccountPoolHub { private async forward(request: Request, body: Uint8Array): Promise { const attempted = new Set(); const accounts = await this.options.accounts.list(); + const parsedBody = parseRequestBody(body); while (attempted.size < accounts.length) { - const selected = await this.select(attempted); - if (selected === null) return this.noEligibleResponse(accounts); + const selected = await this.select(attempted, parsedBody.family); + if (selected === null) + return this.noEligibleResponse(accounts, parsedBody.family); attempted.add(selected.account.id); let secret: AccountSecret; try { @@ -189,7 +310,7 @@ export class AccountPoolHub { try { upstream = await this.fetchUpstream( request, - body, + parsedBody.forAccount(selected.account.accountUuid), selected.account, secret, ); @@ -204,12 +325,13 @@ export class AccountPoolHub { selected.account.id, upstream.response.headers, this.options.quotas.get(selected.account.id), + parsedBody.family, this.options.now(), ); this.options.quotas.put(observed); if ( upstream.response.status === 429 && - isQuotaRejection(observed, this.options.now()) + isQuotaRejection(upstream.response.headers) ) { await upstream.response.body?.cancel(); upstream.release(); @@ -231,7 +353,7 @@ export class AccountPoolHub { try { const retry = await this.fetchUpstream( request, - body, + parsedBody.forAccount(selected.account.accountUuid), selected.account, secret, ); @@ -239,6 +361,7 @@ export class AccountPoolHub { selected.account.id, retry.response.headers, this.options.quotas.get(selected.account.id), + parsedBody.family, this.options.now(), ); if (retry.response.status === 429) { @@ -293,11 +416,12 @@ export class AccountPoolHub { } return this.clientResponse(upstream); } - return this.noEligibleResponse(accounts); + return this.noEligibleResponse(accounts, parsedBody.family); } private async select( attempted: ReadonlySet, + family: ModelFamily, ): Promise { const now = this.options.now(); const threshold = this.options.getSettings().switchThreshold; @@ -309,7 +433,7 @@ export class AccountPoolHub { })) .filter(({ quota }) => quota.error === null) .filter(({ quota }) => quota.heldUntil === null || quota.heldUntil <= now) - .filter(({ quota }) => !isQuotaExhausted(quota, threshold, now)); + .filter(({ quota }) => !isQuotaExhausted(quota, family, threshold, now)); candidates.sort((left, right) => { const priority = left.account.priority - right.account.priority; if (priority !== 0) return priority; @@ -318,8 +442,9 @@ export class AccountPoolHub { (this.inFlightByAccount.get(right.account.id) ?? 0); if (inFlight !== 0) return inFlight; return ( - (left.quota.sevenDayResetAt ?? Number.MAX_SAFE_INTEGER) - - (right.quota.sevenDayResetAt ?? Number.MAX_SAFE_INTEGER) + (governingWeeklyResetAt(left.quota, family) ?? + Number.MAX_SAFE_INTEGER) - + (governingWeeklyResetAt(right.quota, family) ?? Number.MAX_SAFE_INTEGER) ); }); return candidates[0] ?? null; @@ -495,7 +620,10 @@ export class AccountPoolHub { }); } - private noEligibleResponse(accounts: readonly Account[]): Response { + private noEligibleResponse( + accounts: readonly Account[], + family: ModelFamily, + ): Response { if (!accounts.some((account) => account.enabled)) { return anthropicError( 503, @@ -512,6 +640,7 @@ export class AccountPoolHub { quota.heldUntil, quota.fiveHourResetAt, quota.sevenDayResetAt, + governingWeeklyResetAt(quota, family), ].filter((value): value is number => value !== null && value > now); }) .sort((left, right) => left - right)[0]; @@ -563,6 +692,9 @@ export function createHub(options: { fetch?: typeof fetch; now?: () => number; refreshUrl?: string; + usageUrl?: string; + profileUrl?: string; + usageRefreshIntervalMs?: number; drainTimeoutMs?: number; }): AccountPoolHub { return new AccountPoolHub({ @@ -573,6 +705,10 @@ export function createHub(options: { fetch: options.fetch ?? fetch, now: options.now ?? Date.now, refreshUrl: options.refreshUrl ?? DEFAULT_REFRESH_URL, + usageUrl: options.usageUrl ?? DEFAULT_USAGE_URL, + profileUrl: options.profileUrl ?? DEFAULT_PROFILE_URL, + usageRefreshIntervalMs: + options.usageRefreshIntervalMs ?? DEFAULT_USAGE_REFRESH_INTERVAL_MS, drainTimeoutMs: options.drainTimeoutMs ?? 60_000, }); } @@ -613,10 +749,22 @@ function errorMessage(error: unknown): string { return error instanceof Error ? error.message : String(error); } -function waitForAbort(signal: AbortSignal): Promise { +function waitForDelay( + milliseconds: number, + signal: AbortSignal, +): Promise { if (signal.aborted) return Promise.resolve(); return new Promise((resolve) => { - signal.addEventListener("abort", () => resolve(), { once: true }); + const timeout = setTimeout(() => { + signal.removeEventListener("abort", abort); + resolve(); + }, milliseconds); + timeout.unref(); + const abort = () => { + clearTimeout(timeout); + resolve(); + }; + signal.addEventListener("abort", abort, { once: true }); }); } diff --git a/plugins/account-pool/src/oauth-login.test.ts b/plugins/account-pool/src/oauth-login.test.ts index aa22772635..5b875422f5 100644 --- a/plugins/account-pool/src/oauth-login.test.ts +++ b/plugins/account-pool/src/oauth-login.test.ts @@ -59,6 +59,7 @@ function savedAccount(authenticated: ClaudeOAuthAccount): Account { kind: "oauth", label: authenticated.label, email: authenticated.email, + accountUuid: authenticated.accountUuid, subscriptionType: authenticated.subscriptionType, rateLimitTier: authenticated.rateLimitTier, enabled: true, @@ -96,6 +97,7 @@ describe("Claude OAuth login", () => { response.end( JSON.stringify({ account: { + uuid: "11111111-2222-4333-8444-555555555555", email: "person@example.com", display_name: "Personal Claude", has_claude_max: true, @@ -149,6 +151,7 @@ describe("Claude OAuth login", () => { ).resolves.toMatchObject({ label: "Personal Claude", email: "person@example.com", + accountUuid: "11111111-2222-4333-8444-555555555555", subscriptionType: "max", rateLimitTier: "default_claude_max_5x", }); diff --git a/plugins/account-pool/src/oauth-login.ts b/plugins/account-pool/src/oauth-login.ts index 2ca9d375f3..f780671c3f 100644 --- a/plugins/account-pool/src/oauth-login.ts +++ b/plugins/account-pool/src/oauth-login.ts @@ -24,6 +24,7 @@ const profileResponseSchema = z .object({ account: z .object({ + uuid: z.string().uuid().nullish(), email: z.string().email().nullish(), display_name: z.string().trim().min(1).nullish(), has_claude_max: z.boolean().nullish(), @@ -53,6 +54,7 @@ interface LoginSession { export interface ClaudeOAuthAccount { label: string; email: string | null; + accountUuid: string | null; subscriptionType: string | null; rateLimitTier: string | null; accessToken: string; @@ -228,6 +230,7 @@ export class ClaudeOAuthLogin { profile.organization?.name ?? "Claude account", email, + accountUuid: profile.account.uuid ?? null, subscriptionType, rateLimitTier: profile.account.rate_limit_tier ?? diff --git a/plugins/account-pool/src/operations.ts b/plugins/account-pool/src/operations.ts index 05c0d05715..13425af715 100644 --- a/plugins/account-pool/src/operations.ts +++ b/plugins/account-pool/src/operations.ts @@ -46,6 +46,9 @@ export class PoolOperations { private readonly now: () => number = Date.now, private readonly importCredentials: () => Promise = importClaudeCredentials, private readonly onAccountsChanged: () => void = () => {}, + private readonly onAccountEnabled: ( + accountId: string, + ) => Promise = async () => {}, ) {} async add(input: AccountAddInput): Promise { @@ -56,6 +59,7 @@ export class PoolOperations { kind: "api-key", label: input.label ?? "Claude API key", email: null, + accountUuid: null, subscriptionType: null, rateLimitTier: null, enabled: true, @@ -64,6 +68,7 @@ export class PoolOperations { { kind: "api-key", apiKey: input.source.apiKey }, ); this.onAccountsChanged(); + await this.onAccountEnabled(account.id); return account; } const imported = await this.importCredentials(); @@ -73,6 +78,7 @@ export class PoolOperations { kind: "oauth", label: input.label ?? imported.email ?? "Claude Code account", email: imported.email, + accountUuid: imported.accountUuid, subscriptionType: imported.subscriptionType, rateLimitTier: imported.rateLimitTier, enabled: true, @@ -86,6 +92,7 @@ export class PoolOperations { }, ); this.onAccountsChanged(); + await this.onAccountEnabled(account.id); return account; } @@ -96,6 +103,7 @@ export class PoolOperations { kind: "oauth", label: authenticated.label, email: authenticated.email, + accountUuid: authenticated.accountUuid, subscriptionType: authenticated.subscriptionType, rateLimitTier: authenticated.rateLimitTier, enabled: true, @@ -109,6 +117,7 @@ export class PoolOperations { }, ); this.onAccountsChanged(); + await this.onAccountEnabled(account.id); return account; } @@ -131,6 +140,7 @@ export class PoolOperations { const quota = this.quotas.get(id); this.quotas.put({ ...quota, error: null, heldUntil: null }); this.onAccountsChanged(); + await this.onAccountEnabled(account.id); return account; } diff --git a/plugins/account-pool/src/quota.ts b/plugins/account-pool/src/quota.ts index 3e3abf25fd..6678672f25 100644 --- a/plugins/account-pool/src/quota.ts +++ b/plugins/account-pool/src/quota.ts @@ -1,6 +1,20 @@ -import type { Account, AccountQuota, AccountSummary } from "./contracts.js"; +import type { + Account, + AccountQuota, + AccountSummary, + FamilyQuota, + ModelFamily, +} from "./contracts.js"; const PREFIX = "anthropic-ratelimit-unified-"; +const SCOPED_HEADER = + /^anthropic-ratelimit-unified-7d_(.+)-(utilization|reset|status)$/u; + +interface ScopedHeaderValues { + utilization: number | null; + resetAt: number | null; + status: string | null; +} function parseNumber(value: string | null): number | null { if (value === null || value.trim() === "") return null; @@ -8,8 +22,15 @@ function parseNumber(value: string | null): number | null { return Number.isFinite(parsed) ? parsed : null; } -function parseReset(value: string | null): number | null { - if (value === null || value.trim() === "") return null; +export function parseReset(value: string | number | null): number | null { + if (value === null) return null; + if (typeof value === "number") { + if (!Number.isFinite(value)) return null; + return value < 1_000_000_000_000 + ? Math.round(value * 1_000) + : Math.round(value); + } + if (value.trim() === "") return null; const numeric = Number(value); if (Number.isFinite(numeric)) { return numeric < 1_000_000_000_000 @@ -20,10 +41,41 @@ function parseReset(value: string | null): number | null { return Number.isNaN(parsed) ? null : parsed; } +export function modelFamily(model: string | null): ModelFamily { + if (model === null) return "other"; + const normalized = model.toLowerCase(); + if (normalized.includes("fable")) return "fable"; + if (normalized.includes("sonnet")) return "sonnet"; + if (normalized.includes("opus")) return "opus"; + if (normalized.includes("haiku")) return "haiku"; + return "other"; +} + +function scopedHeaderValues(headers: Headers): ScopedHeaderValues | null { + const buckets = new Map(); + for (const [name, value] of headers) { + const match = SCOPED_HEADER.exec(name.toLowerCase()); + const bucket = match?.[1]; + const field = match?.[2]; + if (bucket === undefined || field === undefined) continue; + const current = buckets.get(bucket) ?? { + utilization: null, + resetAt: null, + status: null, + }; + if (field === "utilization") current.utilization = parseNumber(value); + if (field === "reset") current.resetAt = parseReset(value); + if (field === "status") current.status = value; + buckets.set(bucket, current); + } + return [...buckets.values()][0] ?? null; +} + export function quotaFromHeaders( accountId: string, headers: Headers, previous: AccountQuota, + family: ModelFamily, now: number, ): AccountQuota { const fiveHourUtilization = parseNumber( @@ -37,20 +89,21 @@ export function quotaFromHeaders( const sevenDayResetAt = parseReset(headers.get(`${PREFIX}7d-reset`)); const sevenDayStatus = headers.get(`${PREFIX}7d-status`); const representativeClaim = headers.get(`${PREFIX}representative-claim`); - const bucketExhaustion = { ...previous.bucketExhaustion }; - for (const [name, value] of headers) { - if (!name.startsWith(PREFIX) || !name.endsWith("-status")) continue; - if (name === `${PREFIX}5h-status` || name === `${PREFIX}7d-status`) - continue; - const bucket = name.slice(PREFIX.length, -"-status".length); - if (bucket === "" || bucket === "overage") continue; - if (value.toLowerCase() !== "rejected") { - delete bucketExhaustion[bucket]; - continue; - } - const reset = parseReset(headers.get(`${PREFIX}${bucket}-reset`)); - bucketExhaustion[bucket] = reset ?? now; - } + const scoped = scopedHeaderValues(headers); + const priorFamily = previous.familyWeekly[family]; + const familyWeekly: AccountQuota["familyWeekly"] = + scoped === null + ? previous.familyWeekly + : { + ...previous.familyWeekly, + [family]: { + utilization: scoped.utilization ?? priorFamily?.utilization ?? null, + resetAt: scoped.resetAt ?? priorFamily?.resetAt ?? null, + status: scoped.status ?? priorFamily?.status ?? null, + observedAt: now, + source: "header", + }, + }; const observed = fiveHourUtilization !== null || fiveHourResetAt !== null || @@ -58,7 +111,8 @@ export function quotaFromHeaders( sevenDayUtilization !== null || sevenDayResetAt !== null || sevenDayStatus !== null || - representativeClaim !== null; + representativeClaim !== null || + scoped !== null; return { accountId, fiveHourUtilization: fiveHourUtilization ?? previous.fiveHourUtilization, @@ -68,7 +122,7 @@ export function quotaFromHeaders( sevenDayResetAt: sevenDayResetAt ?? previous.sevenDayResetAt, sevenDayStatus: sevenDayStatus ?? previous.sevenDayStatus, representativeClaim: representativeClaim ?? previous.representativeClaim, - bucketExhaustion, + familyWeekly, observedAt: observed ? now : previous.observedAt, heldUntil: previous.heldUntil, error: previous.error, @@ -89,7 +143,18 @@ function activeWindow( ); } -export function isQuotaExhausted( +function activeFamilyWindow( + quota: FamilyQuota | null, + threshold: number, + now: number, +): boolean { + return ( + quota !== null && + activeWindow(quota.utilization, quota.status, quota.resetAt, threshold, now) + ); +} + +export function isSharedQuotaExhausted( quota: AccountQuota, threshold: number, now: number, @@ -112,13 +177,39 @@ export function isQuotaExhausted( ); } -export function isQuotaRejection(quota: AccountQuota, now: number): boolean { +export function isQuotaExhausted( + quota: AccountQuota, + family: ModelFamily, + threshold: number, + now: number, +): boolean { return ( - activeWindow(null, quota.fiveHourStatus, quota.fiveHourResetAt, 1, now) || - activeWindow(null, quota.sevenDayStatus, quota.sevenDayResetAt, 1, now) + isSharedQuotaExhausted(quota, threshold, now) || + activeFamilyWindow(quota.familyWeekly[family], threshold, now) ); } +export function isQuotaRejection(headers: Headers): boolean { + for (const [name, value] of headers) { + if (value.toLowerCase() !== "rejected") continue; + const normalized = name.toLowerCase(); + if ( + normalized === `${PREFIX}5h-status` || + normalized === `${PREFIX}7d-status` || + SCOPED_HEADER.test(normalized) + ) + return true; + } + return false; +} + +export function governingWeeklyResetAt( + quota: AccountQuota, + family: ModelFamily, +): number | null { + return quota.familyWeekly[family]?.resetAt ?? quota.sevenDayResetAt; +} + export function accountStatus( account: Account, quota: AccountQuota, @@ -128,7 +219,7 @@ export function accountStatus( if (!account.enabled) return "disabled"; if (quota.error !== null) return "error"; if (quota.heldUntil !== null && quota.heldUntil > now) return "held"; - if (isQuotaExhausted(quota, threshold, now)) return "exhausted"; + if (isSharedQuotaExhausted(quota, threshold, now)) return "exhausted"; return "ready"; } diff --git a/plugins/account-pool/src/request-body.ts b/plugins/account-pool/src/request-body.ts new file mode 100644 index 0000000000..fed450c832 --- /dev/null +++ b/plugins/account-pool/src/request-body.ts @@ -0,0 +1,72 @@ +import { z } from "zod"; +import type { ModelFamily } from "./contracts.js"; +import { modelFamily } from "./quota.js"; + +const requestSchema = z + .object({ + model: z.string().nullish(), + metadata: z + .object({ user_id: z.string().nullish() }) + .passthrough() + .nullish(), + }) + .passthrough(); + +const encodedUserSchema = z + .object({ account_uuid: z.string().uuid().nullish() }) + .passthrough(); + +const ACCOUNT_COMPONENT = + /(^|_)account_([0-9a-f]{8}(?:-[0-9a-f]{4}){3}-[0-9a-f]{12})(?=_|$)/iu; + +export interface ParsedRequestBody { + family: ModelFamily; + forAccount: (accountUuid: string | null) => Uint8Array; +} + +function rewriteUserId(userId: string, accountUuid: string): string | null { + try { + const encoded = encodedUserSchema.safeParse(JSON.parse(userId)); + if (encoded.success && encoded.data.account_uuid !== undefined) { + if (encoded.data.account_uuid === accountUuid) return null; + return JSON.stringify({ ...encoded.data, account_uuid: accountUuid }); + } + } catch {} + if (!ACCOUNT_COMPONENT.test(userId)) return null; + const rewritten = userId.replace( + ACCOUNT_COMPONENT, + (_match, prefix: string) => `${prefix}account_${accountUuid}`, + ); + return rewritten === userId ? null : rewritten; +} + +export function parseRequestBody(body: Uint8Array): ParsedRequestBody { + const original = body; + try { + const parsed = requestSchema.safeParse( + JSON.parse(new TextDecoder().decode(body)), + ); + if (!parsed.success) { + return { family: "other", forAccount: () => original }; + } + const request = parsed.data; + return { + family: modelFamily(request.model ?? null), + forAccount(accountUuid) { + if (accountUuid === null) return original; + const userId = request.metadata?.user_id; + if (userId === undefined || userId === null) return original; + const rewritten = rewriteUserId(userId, accountUuid); + if (rewritten === null) return original; + return new TextEncoder().encode( + JSON.stringify({ + ...request, + metadata: { ...request.metadata, user_id: rewritten }, + }), + ); + }, + }; + } catch { + return { family: "other", forAccount: () => original }; + } +} diff --git a/plugins/account-pool/src/server.test.ts b/plugins/account-pool/src/server.test.ts index 0812589e22..6de81500dd 100644 --- a/plugins/account-pool/src/server.test.ts +++ b/plugins/account-pool/src/server.test.ts @@ -135,6 +135,7 @@ function importedCredentials( subscriptionType: "max", rateLimitTier: "max_5x", email: "pool@example.com", + accountUuid: "11111111-1111-4111-8111-111111111111", ...overrides, }; } @@ -156,7 +157,10 @@ async function createFixture(args: { }, sdk: sdkStubs(), }); - const plugin = createAccountPoolPlugin(args.options); + const plugin = createAccountPoolPlugin({ + usageUrl: "data:application/json,{}", + ...args.options, + }); await plugin(host.bb); const accountMetadata = accountSchema.parse( await host.harness.behavior.callRpc("account.add", { @@ -486,6 +490,7 @@ describe("Account Pool plugin", () => { response.end( JSON.stringify({ account: { + uuid: "22222222-2222-4222-8222-222222222222", email: "login@example.com", display_name: "Logged-in Claude", has_claude_pro: true, @@ -509,6 +514,7 @@ describe("Account Pool plugin", () => { oauthAuthorizeUrl: `${oauth.url}/authorize`, oauthTokenUrl: `${oauth.url}/token`, oauthProfileUrl: `${oauth.url}/profile`, + usageUrl: "data:application/json,{}", })(host.bb); cleanups.push(async () => { await host.harness.lifecycle.dispose(); @@ -970,7 +976,7 @@ describe("Account Pool plugin", () => { ); expect(unauthorized.status).toBe(401); expect(seen).toHaveLength(0); - const body = Buffer.from('{"model":"claude-test","stream":true}'); + const body = Buffer.from('{"model":"claude-fable-5","stream":true}'); const response = await fixture.host.harness.behavior.fetchHttp( "POST", "/v1/messages?beta=true", @@ -1040,7 +1046,15 @@ describe("Account Pool plugin", () => { fiveHourStatus: "allowed", sevenDayStatus: "allowed", representativeClaim: "claim-a", - bucketExhaustion: { "7d_oi": 4_102_452_000_000 }, + familyWeekly: { + fable: { + utilization: null, + resetAt: 4_102_452_000_000, + status: "rejected", + observedAt: expect.any(Number), + source: "header", + }, + }, }); expect(fixture.host.harness.inspection.registrations.httpRoutes).toEqual( expect.arrayContaining([ @@ -1134,6 +1148,314 @@ describe("Account Pool plugin", () => { expect(keys).toEqual(["sk-one", "sk-two", "sk-three"]); }); + it("routes around a Fable-spent account while retaining it for Opus", async () => { + const keys: string[] = []; + const upstream = await startUpstream(async (request, response) => { + keys.push(request.headers["x-api-key"]?.toString() ?? ""); + await readRequestBody(request); + if (keys.length === 1) { + response.writeHead(200, { + "content-type": "application/json", + "anthropic-ratelimit-unified-5h-status": "allowed", + "anthropic-ratelimit-unified-7d-status": "allowed", + "anthropic-ratelimit-unified-7d_oi-utilization": "0.99", + "anthropic-ratelimit-unified-7d_oi-reset": "4102452000", + "anthropic-ratelimit-unified-7d_oi-status": "allowed", + }); + } else { + response.writeHead(200, { "content-type": "application/json" }); + } + response.end("{}"); + }); + cleanups.push(upstream.close); + const fixture = await createFixture({ + upstreamUrl: upstream.url, + apiKey: "sk-one", + }); + await addApiAccount(fixture, "sk-two"); + + for (const model of [ + "claude-fable-5", + "claude-fable-5", + "claude-opus-4-1", + ]) { + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { + headers: authHeaders(fixture.key), + body: JSON.stringify({ model }), + }, + ); + expect(response.status).toBe(200); + await response.text(); + } + + expect(keys).toEqual(["sk-one", "sk-two", "sk-one"]); + }); + + it("rotates a family-only 429 without exhausting other families", async () => { + const keys: string[] = []; + const upstream = await startUpstream(async (request, response) => { + keys.push(request.headers["x-api-key"]?.toString() ?? ""); + await readRequestBody(request); + if (keys.length === 1) { + response.writeHead(429, { + "content-type": "application/json", + "anthropic-ratelimit-unified-5h-status": "allowed", + "anthropic-ratelimit-unified-7d-status": "allowed", + "anthropic-ratelimit-unified-7d_oi-reset": "4102452000", + "anthropic-ratelimit-unified-7d_oi-status": "rejected", + }); + response.end('{"rejected":true}'); + return; + } + response.writeHead(200, { "content-type": "application/json" }); + response.end("{}"); + }); + cleanups.push(upstream.close); + const fixture = await createFixture({ + upstreamUrl: upstream.url, + apiKey: "sk-one", + }); + await addApiAccount(fixture, "sk-two"); + + const fable = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { + headers: authHeaders(fixture.key), + body: JSON.stringify({ model: "claude-fable-5" }), + }, + ); + expect(fable.status).toBe(200); + await fable.text(); + const opus = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { + headers: authHeaders(fixture.key), + body: JSON.stringify({ model: "claude-opus-4-1" }), + }, + ); + expect(opus.status).toBe(200); + await opus.text(); + + expect(keys).toEqual(["sk-one", "sk-two", "sk-one"]); + const accounts = z + .array(accountSummarySchema) + .parse(await fixture.host.harness.behavior.callRpc("account.list", null)); + expect(accounts[0]).toMatchObject({ + status: "ready", + familyWeekly: { + fable: { status: "rejected", source: "header" }, + }, + }); + }); + + it("refreshes usage on import and routes from its family observations", async () => { + const authorizations: Array = []; + const usageCalls = new Map(); + const upstream = await startUpstream(async (request, response) => { + if (request.url === "/usage") { + const authorization = request.headers.authorization; + usageCalls.set( + authorization ?? "", + (usageCalls.get(authorization ?? "") ?? 0) + 1, + ); + response.writeHead(200, { "content-type": "application/json" }); + response.end( + JSON.stringify({ + five_hour: { utilization: 10, resets_at: "4102444800" }, + seven_day: { utilization: 20, resets_at: "4102448400" }, + limits: [ + { + kind: "weekly_scoped", + group: "weekly", + percent: authorization === "Bearer oauth-a" ? 100 : 0, + resets_at: "4102452000", + scope: { model: { display_name: "Fable" } }, + }, + ], + }), + ); + return; + } + if (request.url === "/profile") { + response.writeHead(200, { "content-type": "application/json" }); + response.end( + JSON.stringify({ + account: { uuid: "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa" }, + }), + ); + return; + } + authorizations.push(request.headers.authorization); + await readRequestBody(request); + response.writeHead(200, { "content-type": "application/json" }); + response.end("{}"); + }); + cleanups.push(upstream.close); + const imports = [ + importedCredentials({ + accessToken: "oauth-a", + accountUuid: null, + }), + importedCredentials({ + accessToken: "oauth-b", + accountUuid: "bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb", + }), + ]; + let importIndex = 0; + const fixture = await createFixture({ + upstreamUrl: upstream.url, + source: "import", + options: { + usageUrl: `${upstream.url}/usage`, + oauthProfileUrl: `${upstream.url}/profile`, + importCredentials: async () => { + const imported = imports[importIndex]; + importIndex += 1; + if (imported === undefined) throw new Error("No import fixture."); + return imported; + }, + }, + }); + const second = accountSchema.parse( + await fixture.host.harness.behavior.callRpc("account.add", { + provider: "claude", + source: { kind: "import" }, + label: "second", + priority: 100, + }), + ); + expect(usageCalls).toEqual( + new Map([ + ["Bearer oauth-a", 1], + ["Bearer oauth-b", 1], + ]), + ); + await fixture.host.harness.behavior.callRpc("account.disable", { + id: second.id, + }); + await fixture.host.harness.behavior.callRpc("account.enable", { + id: second.id, + }); + expect(usageCalls.get("Bearer oauth-a")).toBe(1); + expect(usageCalls.get("Bearer oauth-b")).toBe(2); + const listed = await fixture.host.harness.behavior.runCli([ + "account", + "list", + ]); + expect(listed.stdout).toContain("Fable"); + expect(listed.stdout).toContain("100% rejected"); + expect(listed.stdout).toContain("0% allowed"); + const accounts = z + .array(accountSummarySchema) + .parse(await fixture.host.harness.behavior.callRpc("account.list", null)); + expect(accounts[0]?.accountUuid).toBe( + "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa", + ); + + for (const model of ["claude-fable-5", "claude-opus-4-1"]) { + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { + headers: authHeaders(fixture.key), + body: JSON.stringify({ model }), + }, + ); + expect(response.status).toBe(200); + await response.text(); + } + + expect(authorizations).toEqual(["Bearer oauth-b", "Bearer oauth-a"]); + }); + + it("rewrites both known metadata account UUID formats", async () => { + const bodies: Buffer[] = []; + const upstream = await startUpstream(async (request, response) => { + bodies.push(await readRequestBody(request)); + response.writeHead(200, { "content-type": "application/json" }); + response.end("{}"); + }); + cleanups.push(upstream.close); + const accountUuid = "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa"; + const oldUuid = "bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb"; + const fixture = await createFixture({ + upstreamUrl: upstream.url, + source: "import", + options: { + importCredentials: async () => importedCredentials({ accountUuid }), + }, + }); + const inputs = [ + JSON.stringify({ + model: "claude-fable-5", + metadata: { + user_id: JSON.stringify({ + device_id: "device", + account_uuid: oldUuid, + }), + }, + }), + JSON.stringify({ + model: "claude-fable-5", + metadata: { + user_id: `user_hash_account_${oldUuid}_session_cccccccc-cccc-4ccc-8ccc-cccccccccccc`, + }, + }), + ]; + for (const body of inputs) { + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { headers: authHeaders(fixture.key), body }, + ); + await response.text(); + } + expect(bodies).toHaveLength(2); + expect(bodies.every((body) => body.toString().includes(accountUuid))).toBe( + true, + ); + expect(bodies.every((body) => !body.toString().includes(oldUuid))).toBe( + true, + ); + }); + + it("preserves request bytes when an account UUID rewrite cannot apply", async () => { + const bodies: Buffer[] = []; + const upstream = await startUpstream(async (request, response) => { + bodies.push(await readRequestBody(request)); + response.writeHead(200, { "content-type": "application/json" }); + response.end("{}"); + }); + cleanups.push(upstream.close); + const fixture = await createFixture({ upstreamUrl: upstream.url }); + const inputs = [ + JSON.stringify({ + model: "claude-fable-5", + metadata: { + user_id: JSON.stringify({ + account_uuid: "bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb", + }), + }, + }), + '{ "model": "claude-fable-5", "messages": [] }', + "not-json-at-all", + ]; + for (const body of inputs) { + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { headers: authHeaders(fixture.key), body }, + ); + await response.text(); + } + expect(bodies.map((body) => body.toString())).toEqual(inputs); + }); + it("paces a per-minute 429 on the same account without rotating", async () => { const keys: string[] = []; const times: number[] = []; diff --git a/plugins/account-pool/src/server.ts b/plugins/account-pool/src/server.ts index f1a0ece7ff..88555a3f24 100644 --- a/plugins/account-pool/src/server.ts +++ b/plugins/account-pool/src/server.ts @@ -20,6 +20,8 @@ export interface AccountPoolPluginOptions { fetch?: typeof fetch; now?: () => number; refreshUrl?: string; + usageUrl?: string; + usageRefreshIntervalMs?: number; drainTimeoutMs?: number; disposeTimeoutMs?: number; importCredentials?: () => Promise; @@ -60,7 +62,7 @@ export function createAccountPoolPlugin( type: "number", label: "Quota switch threshold", description: - "Stop selecting an account when its 5-hour or 7-day utilization reaches this fraction.", + "Stop selecting an account when its shared or requested-family quota reaches this fraction.", default: 0.98, experimental_schema: z.number().min(0).max(1), }, @@ -95,6 +97,9 @@ export function createAccountPoolPlugin( fetch: options.fetch, now, refreshUrl: options.refreshUrl, + usageUrl: options.usageUrl, + profileUrl: options.oauthProfileUrl, + usageRefreshIntervalMs: options.usageRefreshIntervalMs, drainTimeoutMs: options.drainTimeoutMs, }); const operations = new PoolOperations( @@ -109,6 +114,7 @@ export function createAccountPoolPlugin( now, options.importCredentials, () => bb.realtime.publish(ACCOUNT_POOL_ACCOUNTS_CHANGED, {}), + (accountId) => hub.refreshUsage(accountId, true), ); const login = new ClaudeOAuthLogin({ fetch: options.fetch, diff --git a/plugins/account-pool/src/store.test.ts b/plugins/account-pool/src/store.test.ts index 9e929a4768..982beb679f 100644 --- a/plugins/account-pool/src/store.test.ts +++ b/plugins/account-pool/src/store.test.ts @@ -2,11 +2,12 @@ import fs from "node:fs/promises"; import path from "node:path"; import { tmpdir } from "node:os"; import { mkdtemp } from "node:fs/promises"; +import Database from "better-sqlite3"; import type { PluginKvStorage } from "@get-bb/plugin-sdk"; import { createFakePluginHost } from "@get-bb/plugin-sdk/testing"; import { afterEach, describe, expect, it } from "vitest"; import type { Account } from "./contracts.js"; -import { AccountStore } from "./store.js"; +import { AccountStore, QUOTA_MIGRATIONS, QuotaStore } from "./store.js"; const cleanups: Array<() => Promise> = []; @@ -30,6 +31,37 @@ function delayedAccountReads(kv: PluginKvStorage): PluginKvStorage { } describe("AccountStore", () => { + it("loads account metadata written before account UUIDs were stored", async () => { + const dataDir = await mkdtemp(path.join(tmpdir(), "bb-account-store-")); + const host = createFakePluginHost({ pluginId: "account-pool", dataDir }); + const store = new AccountStore( + host.bb.storage.kv, + path.join(dataDir, "secrets"), + ); + await host.bb.storage.kv.set("accounts:v1", [ + { + id: "11111111-1111-4111-8111-111111111111", + provider: "claude", + kind: "oauth", + label: "existing", + email: null, + subscriptionType: null, + rateLimitTier: null, + enabled: true, + priority: 100, + createdAt: 1, + }, + ]); + cleanups.push(async () => { + await host.harness.lifecycle.dispose(); + await fs.rm(dataDir, { recursive: true, force: true }); + }); + + expect(await store.list()).toEqual([ + expect.objectContaining({ accountUuid: null }), + ]); + }); + it("preserves both accounts added concurrently", async () => { const dataDir = await mkdtemp(path.join(tmpdir(), "bb-account-store-")); const host = createFakePluginHost({ pluginId: "account-pool", dataDir }); @@ -48,6 +80,7 @@ describe("AccountStore", () => { kind: "api-key", label, email: null, + accountUuid: null, subscriptionType: null, rateLimitTier: null, enabled: true, @@ -64,3 +97,55 @@ describe("AccountStore", () => { ); }); }); + +describe("QuotaStore", () => { + it("migrates an existing quota table to family observations", () => { + const database = new Database(":memory:"); + const initial = QUOTA_MIGRATIONS[0]; + const familyMigration = QUOTA_MIGRATIONS[1]; + if (initial === undefined || familyMigration === undefined) { + throw new Error("Expected quota migrations."); + } + database.exec(initial); + database + .prepare( + `INSERT INTO account_quota ( + account_id, seven_day_utilization, bucket_exhaustion_json + ) VALUES (?, ?, ?)`, + ) + .run( + "11111111-1111-4111-8111-111111111111", + 0.5, + '{"7d_oi":4102452000000}', + ); + database.exec(familyMigration); + const quotas = new QuotaStore(database); + + const migrated = quotas.get("11111111-1111-4111-8111-111111111111"); + expect(migrated.sevenDayUtilization).toBe(0.5); + expect(migrated.familyWeekly).toEqual({ + fable: null, + sonnet: null, + opus: null, + haiku: null, + other: null, + }); + quotas.put({ + ...migrated, + familyWeekly: { + ...migrated.familyWeekly, + fable: { + utilization: 1, + resetAt: 4_102_452_000_000, + status: "rejected", + observedAt: 10, + source: "usage", + }, + }, + }); + expect( + quotas.get("11111111-1111-4111-8111-111111111111").familyWeekly.fable, + ).toMatchObject({ utilization: 1, source: "usage" }); + database.close(); + }); +}); diff --git a/plugins/account-pool/src/store.ts b/plugins/account-pool/src/store.ts index 1f2c5ccf23..c9af62564e 100644 --- a/plugins/account-pool/src/store.ts +++ b/plugins/account-pool/src/store.ts @@ -120,6 +120,23 @@ export class AccountStore { }); } + async setAccountUuid( + id: string, + accountUuid: string, + ): Promise { + return this.serialized(async () => { + const accounts = await this.list(); + const index = accounts.findIndex((account) => account.id === id); + if (index < 0) return null; + const current = accounts[index]; + if (current === undefined) return null; + const updated = accountSchema.parse({ ...current, accountUuid }); + accounts[index] = updated; + await this.kv.set(ACCOUNTS_KEY, accounts); + return updated; + }); + } + async readSecret(id: string): Promise { const parsed = JSON.parse( await fs.readFile(this.accountSecretPath(id), "utf8"), @@ -447,6 +464,7 @@ const quotaRowSchema = z seven_day_status: z.string().nullable(), representative_claim: z.string().nullable(), bucket_exhaustion_json: z.string(), + family_weekly_json: z.string(), observed_at: z.number().int().nullable(), held_until: z.number().int().nullable(), error: z.string().nullable(), @@ -461,7 +479,13 @@ const EMPTY_QUOTA = { sevenDayResetAt: null, sevenDayStatus: null, representativeClaim: null, - bucketExhaustion: {}, + familyWeekly: { + fable: null, + sonnet: null, + opus: null, + haiku: null, + other: null, + }, observedAt: null, heldUntil: null, error: null, @@ -490,7 +514,7 @@ export class QuotaStore { sevenDayResetAt: row.seven_day_reset_at, sevenDayStatus: row.seven_day_status, representativeClaim: row.representative_claim, - bucketExhaustion: JSON.parse(row.bucket_exhaustion_json), + familyWeekly: JSON.parse(row.family_weekly_json), observedAt: row.observed_at, heldUntil: row.held_until, error: row.error, @@ -505,8 +529,8 @@ export class QuotaStore { account_id, five_hour_utilization, five_hour_reset_at, five_hour_status, seven_day_utilization, seven_day_reset_at, seven_day_status, representative_claim, bucket_exhaustion_json, - observed_at, held_until, error - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + family_weekly_json, observed_at, held_until, error + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, '{}', ?, ?, ?, ?) ON CONFLICT(account_id) DO UPDATE SET five_hour_utilization = excluded.five_hour_utilization, five_hour_reset_at = excluded.five_hour_reset_at, @@ -515,7 +539,7 @@ export class QuotaStore { seven_day_reset_at = excluded.seven_day_reset_at, seven_day_status = excluded.seven_day_status, representative_claim = excluded.representative_claim, - bucket_exhaustion_json = excluded.bucket_exhaustion_json, + family_weekly_json = excluded.family_weekly_json, observed_at = excluded.observed_at, held_until = excluded.held_until, error = excluded.error`, @@ -529,7 +553,7 @@ export class QuotaStore { value.sevenDayResetAt, value.sevenDayStatus, value.representativeClaim, - JSON.stringify(value.bucketExhaustion), + JSON.stringify(value.familyWeekly), value.observedAt, value.heldUntil, value.error, @@ -558,4 +582,5 @@ export const QUOTA_MIGRATIONS = [ held_until INTEGER, error TEXT )`, + `ALTER TABLE account_quota ADD COLUMN family_weekly_json TEXT NOT NULL DEFAULT '{"fable":null,"sonnet":null,"opus":null,"haiku":null,"other":null}'`, ]; diff --git a/plugins/account-pool/src/usage.ts b/plugins/account-pool/src/usage.ts new file mode 100644 index 0000000000..a89ed21181 --- /dev/null +++ b/plugins/account-pool/src/usage.ts @@ -0,0 +1,183 @@ +import { z } from "zod"; +import type { + AccountQuota, + FamilyQuota, + FamilyWeekly, + ModelFamily, +} from "./contracts.js"; +import { modelFamily, parseReset } from "./quota.js"; + +const numberValueSchema = z.union([z.number(), z.string()]); + +const usageBucketSchema = z + .object({ + utilization: numberValueSchema.nullish(), + used_percentage: numberValueSchema.nullish(), + usedPercentage: numberValueSchema.nullish(), + resets_at: z.union([z.number(), z.string()]).nullish(), + reset_at: z.union([z.number(), z.string()]).nullish(), + resetAt: z.union([z.number(), z.string()]).nullish(), + status: z.string().nullish(), + }) + .passthrough(); + +const usageLimitSchema = z + .object({ + kind: z.string().nullish(), + group: z.string().nullish(), + percent: numberValueSchema.nullish(), + resets_at: z.union([z.number(), z.string()]).nullish(), + status: z.string().nullish(), + scope: z + .object({ + model: z + .object({ display_name: z.string().nullish() }) + .passthrough() + .nullish(), + }) + .passthrough() + .nullish(), + }) + .passthrough(); + +const usagePayloadSchema = z + .object({ + five_hour: usageBucketSchema.nullish(), + seven_day: usageBucketSchema.nullish(), + seven_day_fable: usageBucketSchema.nullish(), + seven_day_sonnet: usageBucketSchema.nullish(), + seven_day_opus: usageBucketSchema.nullish(), + seven_day_haiku: usageBucketSchema.nullish(), + limits: z.array(usageLimitSchema).optional(), + }) + .passthrough(); + +type UsageBucket = z.infer; + +function percentage(value: string | number | null | undefined): number | null { + if (value === null || value === undefined || value === "") return null; + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed / 100 : null; +} + +function normalizeBucket(bucket: UsageBucket, now: number): FamilyQuota | null { + const utilization = percentage( + bucket.used_percentage ?? + bucket.utilization ?? + bucket.usedPercentage ?? + null, + ); + const resetAt = parseReset( + bucket.resets_at ?? bucket.reset_at ?? bucket.resetAt ?? null, + ); + if (utilization === null && resetAt === null && bucket.status == null) + return null; + return { + utilization, + resetAt, + status: + bucket.status ?? + (utilization === null ? null : utilization >= 1 ? "rejected" : "allowed"), + observedAt: now, + source: "usage", + }; +} + +function stronger( + current: FamilyQuota | null, + candidate: FamilyQuota | null, +): FamilyQuota | null { + if (candidate === null) return current; + if (current === null) return candidate; + return (candidate.utilization ?? -1) > (current.utilization ?? -1) + ? candidate + : current; +} + +function limitBucket( + percent: string | number | null | undefined, + reset: string | number | null | undefined, + status: string | null | undefined, + now: number, +): FamilyQuota | null { + return normalizeBucket( + { + utilization: percent, + resets_at: reset, + status, + }, + now, + ); +} + +export function quotaFromUsage( + accountId: string, + payload: object, + previous: AccountQuota, + now: number, +): AccountQuota | null { + const parsed = usagePayloadSchema.safeParse(payload); + if (!parsed.success) return null; + const data = parsed.data; + if ( + data.five_hour == null && + data.seven_day == null && + data.seven_day_fable == null && + data.seven_day_sonnet == null && + data.seven_day_opus == null && + data.seven_day_haiku == null && + data.limits === undefined + ) + return null; + const familyWeekly: FamilyWeekly = + data.limits === undefined + ? { ...previous.familyWeekly } + : { + fable: null, + sonnet: null, + opus: null, + haiku: null, + other: null, + }; + for (const limit of data.limits ?? []) { + if ( + limit.kind !== "weekly_scoped" && + !(limit.group === "weekly" && limit.scope?.model?.display_name != null) + ) + continue; + const displayName = limit.scope?.model?.display_name; + if (displayName == null) continue; + const family = modelFamily(displayName); + familyWeekly[family] = stronger( + familyWeekly[family], + limitBucket(limit.percent, limit.resets_at, limit.status, now), + ); + } + const slots: Array<[ModelFamily, UsageBucket | null | undefined]> = [ + ["fable", data.seven_day_fable], + ["sonnet", data.seven_day_sonnet], + ["opus", data.seven_day_opus], + ["haiku", data.seven_day_haiku], + ]; + for (const [family, bucket] of slots) { + if (bucket == null) continue; + const normalized = normalizeBucket(bucket, now); + if (normalized !== null) familyWeekly[family] = normalized; + } + const fiveHour = + data.five_hour == null ? null : normalizeBucket(data.five_hour, now); + const sevenDay = + data.seven_day == null ? null : normalizeBucket(data.seven_day, now); + return { + ...previous, + accountId, + fiveHourUtilization: fiveHour?.utilization ?? previous.fiveHourUtilization, + fiveHourResetAt: fiveHour?.resetAt ?? previous.fiveHourResetAt, + fiveHourStatus: fiveHour?.status ?? previous.fiveHourStatus, + sevenDayUtilization: sevenDay?.utilization ?? previous.sevenDayUtilization, + sevenDayResetAt: sevenDay?.resetAt ?? previous.sevenDayResetAt, + sevenDayStatus: sevenDay?.status ?? previous.sevenDayStatus, + familyWeekly, + observedAt: now, + }; +}