diff --git a/apps/cli/src/__tests__/plugin-cli-proxy.test.ts b/apps/cli/src/__tests__/plugin-cli-proxy.test.ts index 62dcec96a4..98c05f8c56 100644 --- a/apps/cli/src/__tests__/plugin-cli-proxy.test.ts +++ b/apps/cli/src/__tests__/plugin-cli-proxy.test.ts @@ -496,6 +496,46 @@ describe("runPluginCliCommand", () => { ]); }); + it("materializes an API key from stdin only in the proxied request", async () => { + const requests: string[][] = []; + vi.stubGlobal( + "fetch", + vi.fn(async (_url, init: RequestInit | undefined) => { + const parsed = JSON.parse(String(init?.body)) as { argv: string[] }; + requests.push(parsed.argv); + return new Response(JSON.stringify({ exitCode: 0 }), { status: 200 }); + }), + ); + const writes: string[] = []; + const output = { + write(value: string, callback: (error?: Error | null) => void) { + writes.push(value); + callback(); + return true; + }, + }; + const input = { + isTTY: false, + async *[Symbol.asyncIterator]() { + yield Buffer.from("sk-from-stdin\n"); + }, + }; + + await expect( + runPluginCliCommand( + "http://localhost", + "account-pool", + ["account", "add", "--provider", "claude", "--api-key-stdin"], + { stdout: output, stderr: output }, + input, + ), + ).resolves.toBe(0); + expect(requests).toEqual([ + ["account", "add", "--provider", "claude", "--api-key", "sk-from-stdin"], + ]); + expect(writes).toEqual([]); + }); + it("outlives the global fetch headers timeout while a plugin command waits on a human", async () => { const RESPONSE_DELAY_MS = 1500; const server: Server = createServer((request, response) => { diff --git a/apps/cli/src/plugin-cli-proxy.ts b/apps/cli/src/plugin-cli-proxy.ts index 97d07b4dec..d1b83377ca 100644 --- a/apps/cli/src/plugin-cli-proxy.ts +++ b/apps/cli/src/plugin-cli-proxy.ts @@ -279,6 +279,52 @@ interface PluginCliOutputStreams { stderr: PluginCliOutputStream; } +interface PluginCliInputStream extends AsyncIterable { + isTTY?: boolean; +} + +const PLUGIN_CLI_SECRET_STDIN_MAX_BYTES = 16 * 1024; + +async function materializeApiKeyStdin( + argv: readonly string[], + input: PluginCliInputStream, +): Promise { + const indexes = argv.flatMap((arg, index) => + arg === "--api-key-stdin" ? [index] : [], + ); + if (indexes.length === 0) return [...argv]; + if (indexes.length > 1 || argv.includes("--api-key")) { + throw new Error("Choose only one API-key input flag."); + } + if (input.isTTY === true) { + throw new Error("--api-key-stdin requires an API key piped on stdin."); + } + const chunks: Buffer[] = []; + let bytes = 0; + for await (const chunk of input) { + const buffer = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk); + bytes += buffer.byteLength; + if (bytes > PLUGIN_CLI_SECRET_STDIN_MAX_BYTES) { + throw new Error("API key from stdin exceeds 16 KiB."); + } + chunks.push(buffer); + } + const apiKey = Buffer.concat(chunks) + .toString("utf8") + .replace(/[\r\n]+$/u, ""); + if (apiKey.length === 0 || /[\r\n]/u.test(apiKey)) { + throw new Error("--api-key-stdin requires exactly one non-empty API key."); + } + const index = indexes[0]; + if (index === undefined) return [...argv]; + return [ + ...argv.slice(0, index), + "--api-key", + apiKey, + ...argv.slice(index + 1), + ]; +} + async function writePluginCliOutput( stream: PluginCliOutputStream, value: string, @@ -310,7 +356,18 @@ export async function runPluginCliCommand( stdout: process.stdout, stderr: process.stderr, }, + input: PluginCliInputStream = process.stdin, ): Promise { + let resolvedArgv: string[]; + try { + resolvedArgv = await materializeApiKeyStdin(argv, input); + } catch (error) { + await writePluginCliOutput( + streams.stderr, + error instanceof Error ? error.message : String(error), + ); + return 1; + } const threadId = resolveContextThreadId(); const projectId = resolveContextProjectId(); const response = await cliFetch( @@ -319,7 +376,7 @@ export async function runPluginCliCommand( method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ - argv, + argv: resolvedArgv, cwd: process.cwd(), ...(threadId ? { threadId } : {}), ...(projectId ? { projectId } : {}), diff --git a/apps/server/src/services/plugins/builtin-registry.ts b/apps/server/src/services/plugins/builtin-registry.ts index 72fb3d9c15..3349723a05 100644 --- a/apps/server/src/services/plugins/builtin-registry.ts +++ b/apps/server/src/services/plugins/builtin-registry.ts @@ -23,6 +23,11 @@ export const BUILTIN_PLUGINS_DIRECTORY_NAME = "builtin-plugins"; const REPO_PLUGINS_DIRECTORY_NAME = "plugins"; export const BUILTIN_PLUGINS = [ + { + name: "account-pool", + pluginId: "account-pool", + defaultEnabled: false, + }, { name: "ask-user-question", pluginId: "ask-user-question", 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 8b5085b1d8..2f068ebe52 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 @@ -51,6 +51,30 @@ BB_HOST_DAEMON_PORT only for an intentional non-default target. - Prefer non-interactive commands and machine-readable output for automation. - Pass `--yes` for a confirmed destructive command in a non-interactive shell. - Treat plugin commands as normal top-level commands after installation. + +The builtin Account Pool plugin is disabled by default. Enable it, add Claude +credentials, and inspect its proxy route and account quota with: + +```sh +bb plugin enable account-pool +bb pool account add --provider claude --import +printf '%s\n' "$ANTHROPIC_API_KEY" | bb pool account add --provider claude --api-key-stdin [--label ] [--priority ] +bb pool account add --provider claude --api-key [--label ] [--priority ] +bb pool account list [--json] +bb pool account remove +bb pool account enable +bb pool account disable +bb pool status [--json] [--show-key] +``` + +Newly added or enabled accounts are available without a plugin reload. The hub +bearer key appears only with `status --show-key`. 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. + - 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/apps/server/test/services/plugins/builtin-plugins.test.ts b/apps/server/test/services/plugins/builtin-plugins.test.ts index 18740dfc02..203e99bc53 100644 --- a/apps/server/test/services/plugins/builtin-plugins.test.ts +++ b/apps/server/test/services/plugins/builtin-plugins.test.ts @@ -223,6 +223,7 @@ describe("builtin plugin reconciliation", () => { it("gives every builtin plugin a deliberate settings icon", async () => { const expectedIcons = new Map([ + ["account-pool", "Layers"], ["ask-user-question", "MessageQuestion"], ["automations", "Clock"], ["concurrency-limit", "Limitation"], diff --git a/docs/configuration.md b/docs/configuration.md index e71fe7e9c4..35946645e8 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -611,6 +611,44 @@ how many connected clients received the broadcast. `spotlight` focuses the target pane and persistently dims the others; `clear-spotlight` focuses it and persistently restores undimmed splits. +## Account Pool + +The builtin Account Pool plugin is disabled on fresh installations. It stores +non-secret Claude account metadata in plugin KV, quota observations in the +plugin SQLite database, and each account token plus the generated hub bearer +key in 0600 files under `/plugins/account-pool/secrets/accounts/`. +Enable it and add at least one account: + +```sh +bb plugin enable account-pool +bb pool account add --provider claude --import +printf '%s\n' "$ANTHROPIC_API_KEY" | bb pool account add --provider claude --api-key-stdin [--label ] [--priority ] +``` + +The import path reads the Claude Code login on the bb server host. +`--api-key-stdin` reads exactly one non-empty key from piped standard input and +is the default API-key path for agents. The compatibility form `--api-key +` remains available, but exposes the secret in process arguments, shell +history, and agent transcripts. The hub starts immediately, so a newly added +or enabled account is available without a plugin reload. + +`bb pool status --show-key` is the only command that reveals the hub bearer +key. Point a client at the route printed by that command and supply the key as +`Authorization: Bearer `. 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`. +`upstreamBaseUrl` defaults to `https://api.anthropic.com` and exists only for +tests and QA with a controlled fake upstream: + +```sh +bb plugin config account-pool set switchThreshold 0.98 +bb plugin config account-pool set upstreamBaseUrl http://127.0.0.1:9000 +``` + ## bb connect `bb connect --code --server https://.getbb.app` pairs this bb diff --git a/packages/plugin-api-map/src/plugin-icons.ts b/packages/plugin-api-map/src/plugin-icons.ts index 9beacc8664..8890663ad7 100644 --- a/packages/plugin-api-map/src/plugin-icons.ts +++ b/packages/plugin-api-map/src/plugin-icons.ts @@ -33,6 +33,7 @@ interface FirstPartyPlugin { } const FIRST_PARTY_PLUGINS: Record = { + "Account Pool": { id: "account-pool", icon: Layers01Icon }, "Ask User Question": { id: "ask-user-question", icon: MessageQuestionIcon }, Automations: { id: "automations", icon: Clock01Icon }, "Custom instructions": { id: "custom-instructions", icon: Edit04Icon }, diff --git a/packages/templates/src/templates/bb-guide-plugins.md b/packages/templates/src/templates/bb-guide-plugins.md index c44c4ec90d..b529e1d03d 100644 --- a/packages/templates/src/templates/bb-guide-plugins.md +++ b/packages/templates/src/templates/bb-guide-plugins.md @@ -23,6 +23,35 @@ The builtin Custom instructions plugin adds a multiline editor under Settings → Custom instructions. Saved text is persisted on this bb host and included in agent task instructions; blank text contributes nothing. +The builtin Account Pool plugin is disabled on fresh installations. It stores +Claude account tokens in per-account 0600 secret files and proxies Anthropic +Messages API requests through the bb server. Enable it, add an account, then +point Claude Code at the route and bearer key shown by `status`: + +``` +bb plugin enable account-pool +bb pool account add --provider claude --import +printf '%s\n' "$ANTHROPIC_API_KEY" | bb pool account add --provider claude --api-key-stdin [--label ] [--priority ] +bb pool account add --provider claude --api-key [--label ] [--priority ] +bb pool account list [--json] +bb pool account remove +bb pool account enable +bb pool account disable +bb pool status [--json] [--show-key] +``` + +The hub starts immediately, even before an account is configured, so newly +added or enabled accounts are available without a plugin reload. Only +`status --show-key` reveals the hub bearer key. 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. +The `upstreamBaseUrl` setting exists for tests and QA and defaults to +`https://api.anthropic.com`; `switchThreshold` defaults to `0.98`. + The builtin Keep Awake plugin prevents macOS idle sleep while bb is running. Its settings page lets you target all hosts or selected hosts. The CLI equivalents are: diff --git a/plugins/account-pool/package.json b/plugins/account-pool/package.json new file mode 100644 index 0000000000..1a0ebd5009 --- /dev/null +++ b/plugins/account-pool/package.json @@ -0,0 +1,33 @@ +{ + "name": "bb-plugin-account-pool", + "version": "0.1.0", + "private": true, + "type": "module", + "description": "Routes Anthropic Messages API traffic across a pool of Claude accounts.", + "engines": { + "bb": ">=0.0" + }, + "bb": { + "name": "Account Pool", + "description": "Routes Anthropic Messages API traffic across a pool of Claude accounts.", + "branding": { + "icon": "Layers" + }, + "server": "./src/server.ts" + }, + "scripts": { + "typecheck": "tsc --noEmit", + "test": "vitest run --config vitest.config.ts" + }, + "dependencies": { + "zod": "^4.3.6" + }, + "devDependencies": { + "@get-bb/plugin-sdk": "workspace:*", + "@types/better-sqlite3": "^7.6.12", + "@types/node": "^22.0.0", + "better-sqlite3": "12.10.0", + "typescript": "npm:@typescript/typescript6@^6.0.2", + "vitest": "^4.1.1" + } +} diff --git a/plugins/account-pool/src/cli.ts b/plugins/account-pool/src/cli.ts new file mode 100644 index 0000000000..bbdbcb6e6d --- /dev/null +++ b/plugins/account-pool/src/cli.ts @@ -0,0 +1,238 @@ +import type { BbPluginApi, PluginCliResult } from "@get-bb/plugin-sdk"; +import { + accountAddInputSchema, + accountIdInputSchema, + type AccountSummary, + type PoolStatus, +} from "./contracts.js"; +import type { PoolOperations } from "./operations.js"; + +interface ParsedFlags { + booleans: Set; + values: Map; +} + +const HELP = [ + "Usage:", + " bb pool account add --provider claude --import [--label ] [--priority ]", + " bb pool account add --provider claude --api-key-stdin [--label ] [--priority ]", + " bb pool account add --provider claude --api-key [--label ] [--priority ] Unsafe: exposes the key in process arguments.", + " bb pool account list [--json]", + " bb pool account remove ", + " bb pool account enable ", + " bb pool account disable ", + " bb pool status [--json] [--show-key]", +].join("\n"); + +function parseFlags( + argv: readonly string[], + allowedBooleans: readonly string[], + allowedValues: readonly string[], +): ParsedFlags { + const booleans = new Set(); + const values = new Map(); + for (let index = 0; index < argv.length; index += 1) { + const arg = argv[index]; + if (arg === undefined || !arg.startsWith("--")) { + throw new Error(`Unexpected argument ${JSON.stringify(arg)}.`); + } + const name = arg.slice(2); + if (booleans.has(name) || values.has(name)) { + throw new Error(`Duplicate flag --${name}.`); + } + if (allowedBooleans.includes(name)) { + booleans.add(name); + continue; + } + if (!allowedValues.includes(name)) + throw new Error(`Unknown flag --${name}.`); + const value = argv[index + 1]; + if (value === undefined || value.startsWith("--")) { + throw new Error(`--${name} requires a value.`); + } + values.set(name, value); + index += 1; + } + return { booleans, values }; +} + +function formatReset(value: number | null): string { + return value === null ? "-" : new Date(value).toISOString(); +} + +function formatUtilization(value: number | null): string { + return value === null ? "-" : `${Math.round(value * 100)}%`; +} + +function formatAccounts(accounts: readonly AccountSummary[]): string { + if (accounts.length === 0) return "No accounts configured."; + return [ + "ID\tLabel\tKind\tEnabled\tPriority\t5h\t5h reset\t7d\t7d reset\tStatus", + ...accounts.map((account) => + [ + account.id, + account.label, + account.kind, + String(account.enabled), + String(account.priority), + formatUtilization(account.fiveHourUtilization), + formatReset(account.fiveHourResetAt), + formatUtilization(account.sevenDayUtilization), + formatReset(account.sevenDayResetAt), + account.status, + ].join("\t"), + ), + ].join("\n"); +} + +function formatStatus(status: PoolStatus, showKey: boolean): string { + return [ + `Route: ${status.route}`, + `Accepting: ${status.accepting}`, + `Enabled accounts: ${status.enabledAccountCount}`, + `In flight: ${status.inFlight}`, + ...(showKey && status.hubKey !== null ? [`Hub key: ${status.hubKey}`] : []), + "", + formatAccounts(status.accounts), + ].join("\n"); +} + +function json(value: object): string { + return `${JSON.stringify(value, null, 2)}\n`; +} + +export function registerPoolCli( + bb: Pick, + operations: PoolOperations, +): void { + bb.cli.register({ + name: "pool", + summary: "Manage Claude accounts and inspect the Account Pool hub", + commands: [ + { + name: "account-add", + summary: + "Import Claude Code OAuth credentials or add an Anthropic API key", + usage: + "bb pool account add --provider claude (--import | --api-key-stdin) [--label ] [--priority ]\nUnsafe compatibility form: bb pool account add --provider claude --api-key [--label ] [--priority ]", + }, + { + name: "account-list", + summary: "List pool accounts and observed quota", + usage: "bb pool account list [--json]", + }, + { + name: "account-remove", + summary: "Remove an account and its secret token file", + usage: "bb pool account remove ", + }, + { + name: "account-enable", + summary: "Enable an account", + usage: "bb pool account enable ", + }, + { + name: "account-disable", + summary: "Disable an account", + usage: "bb pool account disable ", + }, + { + name: "status", + summary: + "Show hub status; reveal the bearer token only with --show-key", + usage: "bb pool status [--json] [--show-key]", + }, + ], + async run(argv): Promise { + try { + if (argv.includes("--help") || argv.includes("-h")) { + return { exitCode: 0, stdout: `${HELP}\n` }; + } + if (argv[0] === "account" && argv[1] === "add") { + const flags = parseFlags( + argv.slice(2), + ["import", "api-key-stdin"], + ["provider", "api-key", "label", "priority"], + ); + const imported = flags.booleans.has("import"); + const apiKeyStdin = flags.booleans.has("api-key-stdin"); + const apiKey = flags.values.get("api-key"); + const sourceCount = + Number(imported) + + Number(apiKeyStdin) + + Number(apiKey !== undefined); + if (sourceCount !== 1) + throw new Error( + "Choose exactly one of --import, --api-key-stdin, or --api-key .", + ); + if (apiKeyStdin) { + throw new Error( + "--api-key-stdin must be invoked through the bb CLI so it can read stdin safely.", + ); + } + const priorityText = flags.values.get("priority") ?? "100"; + const input = accountAddInputSchema.parse({ + provider: flags.values.get("provider"), + source: imported ? { kind: "import" } : { kind: "api-key", apiKey }, + label: flags.values.get("label") ?? null, + priority: Number(priorityText), + }); + const account = await operations.add(input); + return { + exitCode: 0, + stdout: `Added ${account.label} (${account.id}).\n`, + }; + } + if (argv[0] === "account" && argv[1] === "list") { + const flags = parseFlags(argv.slice(2), ["json"], []); + const accounts = await operations.list(); + return { + exitCode: 0, + stdout: flags.booleans.has("json") + ? json({ accounts }) + : `${formatAccounts(accounts)}\n`, + }; + } + if ( + argv[0] === "account" && + ["remove", "enable", "disable"].includes(argv[1] ?? "") + ) { + if (argv.length !== 3) throw new Error(HELP); + const { id } = accountIdInputSchema.parse({ id: argv[2] }); + if (argv[1] === "remove") { + const removed = await operations.remove(id); + if (!removed) throw new Error(`Account ${id} does not exist.`); + return { exitCode: 0, stdout: `Removed ${id}.\n` }; + } + const account = + argv[1] === "enable" + ? await operations.enable(id) + : await operations.disable(id); + if (account === null) + throw new Error(`Account ${id} does not exist.`); + return { + exitCode: 0, + stdout: `${argv[1] === "enable" ? "Enabled" : "Disabled"} ${id}.\n`, + }; + } + if (argv[0] === "status") { + const flags = parseFlags(argv.slice(1), ["json", "show-key"], []); + const showKey = flags.booleans.has("show-key"); + const status = await operations.status(showKey); + return { + exitCode: 0, + stdout: flags.booleans.has("json") + ? json(status) + : `${formatStatus(status, showKey)}\n`, + }; + } + throw new Error(HELP); + } catch (error) { + return { + exitCode: 1, + stderr: `${error instanceof Error ? error.message : String(error)}\n`, + }; + } + }, + }); +} diff --git a/plugins/account-pool/src/contracts.ts b/plugins/account-pool/src/contracts.ts new file mode 100644 index 0000000000..b252bafdfa --- /dev/null +++ b/plugins/account-pool/src/contracts.ts @@ -0,0 +1,114 @@ +import { z } from "zod"; + +export const providerSchema = z.literal("claude"); +export const accountKindSchema = z.enum(["oauth", "api-key"]); + +export const accountSchema = z + .object({ + id: z.string().uuid(), + provider: providerSchema, + kind: accountKindSchema, + label: z.string().min(1), + email: z.string().email().nullable(), + subscriptionType: z.string().nullable(), + rateLimitTier: z.string().nullable(), + enabled: z.boolean(), + priority: z.number().int(), + createdAt: z.number().int().nonnegative(), + }) + .strict(); + +export type Account = z.infer; + +export const oauthSecretSchema = z + .object({ + kind: z.literal("oauth"), + accessToken: z.string().min(1), + refreshToken: z.string().min(1), + expiresAt: z.number().int().positive().nullable(), + }) + .strict(); + +export const apiKeySecretSchema = z + .object({ + kind: z.literal("api-key"), + apiKey: z.string().min(1), + }) + .strict(); + +export const accountSecretSchema = z.discriminatedUnion("kind", [ + oauthSecretSchema, + apiKeySecretSchema, +]); + +export type AccountSecret = z.infer; + +export const quotaSchema = z + .object({ + accountId: z.string().uuid(), + fiveHourUtilization: z.number().nullable(), + fiveHourResetAt: z.number().int().nullable(), + fiveHourStatus: z.string().nullable(), + sevenDayUtilization: z.number().nullable(), + sevenDayResetAt: z.number().int().nullable(), + sevenDayStatus: z.string().nullable(), + representativeClaim: z.string().nullable(), + bucketExhaustion: z.record(z.string(), z.number().int()), + observedAt: z.number().int().nullable(), + heldUntil: z.number().int().nullable(), + error: z.string().nullable(), + }) + .strict(); + +export type AccountQuota = z.infer; + +export const accountSummarySchema = accountSchema.extend({ + fiveHourUtilization: z.number().nullable(), + fiveHourResetAt: z.number().int().nullable(), + fiveHourStatus: z.string().nullable(), + sevenDayUtilization: z.number().nullable(), + sevenDayResetAt: z.number().int().nullable(), + sevenDayStatus: z.string().nullable(), + representativeClaim: z.string().nullable(), + bucketExhaustion: z.record(z.string(), z.number().int()), + observedAt: z.number().int().nullable(), + heldUntil: z.number().int().nullable(), + error: z.string().nullable(), + inFlight: z.number().int().nonnegative(), + status: z.enum(["disabled", "ready", "held", "exhausted", "error"]), +}); + +export type AccountSummary = z.infer; + +export const statusSchema = z + .object({ + route: z.string(), + enabledAccountCount: z.number().int().nonnegative(), + inFlight: z.number().int().nonnegative(), + accepting: z.boolean(), + hubKey: z.string().nullable(), + accounts: z.array(accountSummarySchema), + }) + .strict(); + +export type PoolStatus = z.infer; + +export const accountAddInputSchema = z + .object({ + provider: providerSchema, + source: z.discriminatedUnion("kind", [ + z.object({ kind: z.literal("import") }).strict(), + z + .object({ kind: z.literal("api-key"), apiKey: z.string().min(1) }) + .strict(), + ]), + label: z.string().min(1).nullable(), + priority: z.number().int(), + }) + .strict(); + +export type AccountAddInput = z.infer; + +export const accountIdInputSchema = z + .object({ id: z.string().uuid() }) + .strict(); diff --git a/plugins/account-pool/src/credentials.ts b/plugins/account-pool/src/credentials.ts new file mode 100644 index 0000000000..9abda506be --- /dev/null +++ b/plugins/account-pool/src/credentials.ts @@ -0,0 +1,113 @@ +import { execFile } from "node:child_process"; +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { promisify } from "node:util"; +import { z } from "zod"; + +const execFileAsync = promisify(execFile); +const KEYCHAIN_SERVICE = "Claude Code-credentials"; + +const credentialsFileSchema = z + .object({ + claudeAiOauth: z + .object({ + accessToken: z.string().min(1), + refreshToken: z.string().min(1), + expiresAt: z.number().nullish(), + subscriptionType: z.string().nullish(), + rateLimitTier: z.string().nullish(), + }) + .passthrough(), + }) + .passthrough(); + +const accountFileSchema = z.object({ + oauthAccount: z + .object({ emailAddress: z.string().email().nullish() }) + .nullish(), +}); + +export interface ImportedClaudeCredentials { + accessToken: string; + refreshToken: string; + expiresAt: number | null; + subscriptionType: string | null; + rateLimitTier: string | null; + email: string | null; +} + +function parseCredentials( + raw: string, +): Omit | null { + const trimmed = raw.trim(); + const candidates = [trimmed]; + if (/^(?:[0-9a-f]{2})+$/iu.test(trimmed)) { + candidates.push(Buffer.from(trimmed, "hex").toString("utf8")); + } + for (const candidate of candidates) { + try { + const parsed = credentialsFileSchema.safeParse(JSON.parse(candidate)); + if (!parsed.success) continue; + return { + accessToken: parsed.data.claudeAiOauth.accessToken, + refreshToken: parsed.data.claudeAiOauth.refreshToken, + expiresAt: parsed.data.claudeAiOauth.expiresAt ?? null, + subscriptionType: parsed.data.claudeAiOauth.subscriptionType ?? null, + rateLimitTier: parsed.data.claudeAiOauth.rateLimitTier ?? null, + }; + } catch {} + } + return null; +} + +async function readKeychainCredentials(): Promise { + if (process.platform !== "darwin") return null; + const username = os.userInfo().username; + const argumentSets = [ + ["find-generic-password", "-s", KEYCHAIN_SERVICE, "-a", username, "-w"], + ["find-generic-password", "-s", KEYCHAIN_SERVICE, "-w"], + ]; + for (const args of argumentSets) { + try { + const result = await execFileAsync("security", args, { timeout: 10_000 }); + if (result.stdout.trim()) return result.stdout.trim(); + } catch {} + } + return null; +} + +async function readAccountEmail(): Promise { + 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; + } catch { + return null; + } +} + +export async function importClaudeCredentials(): Promise { + const keychain = await readKeychainCredentials(); + let credentials = keychain === null ? null : parseCredentials(keychain); + if (credentials === null) { + try { + credentials = parseCredentials( + await fs.readFile( + path.join(os.homedir(), ".claude", ".credentials.json"), + "utf8", + ), + ); + } catch {} + } + if (credentials === null) { + throw new Error( + "Claude Code OAuth credentials were not found. Run `claude /login` on the bb server host, then retry.", + ); + } + return { ...credentials, email: await readAccountEmail() }; +} diff --git a/plugins/account-pool/src/hub.ts b/plugins/account-pool/src/hub.ts new file mode 100644 index 0000000000..c36b6c9fb2 --- /dev/null +++ b/plugins/account-pool/src/hub.ts @@ -0,0 +1,637 @@ +import { timingSafeEqual } from "node:crypto"; +import { z } from "zod"; +import type { + Account, + AccountQuota, + AccountSecret, + PoolStatus, +} from "./contracts.js"; +import { + accountStatus, + isQuotaExhausted, + isQuotaRejection, + quotaFromHeaders, + retryAfterMilliseconds, +} from "./quota.js"; +import type { AccountStore, QuotaStore } from "./store.js"; + +const ROUTE = "/api/v1/plugins/account-pool/http"; +const OAUTH_CLIENT_ID = "9d1c250a-e61b-44d9-88ed-5944d1962f5e"; +const DEFAULT_REFRESH_URL = "https://platform.claude.com/v1/oauth/token"; +const REFRESH_WINDOW_MS = 5 * 60 * 1_000; +const MAX_INLINE_HOLD_MS = 20_000; + +const ALLOWED_REQUEST_HEADERS = new Set([ + "accept", + "content-type", + "user-agent", + "x-app", +]); + +const ALLOWED_REQUEST_HEADER_PREFIXES = ["anthropic-", "x-stainless-"]; + +const DROPPED_RESPONSE_HEADERS = new Set([ + "content-encoding", + "content-length", + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", +]); + +const refreshResponseSchema = z + .object({ + access_token: z.string().min(1), + refresh_token: z.string().min(1).optional(), + expires_in: z.number().positive().optional(), + expires_at: z.number().positive().optional(), + }) + .passthrough(); + +export interface HubSettings { + upstreamBaseUrl: string; + switchThreshold: number; +} + +interface HubOptions { + accounts: AccountStore; + quotas: QuotaStore; + hubKey: string; + getSettings: () => HubSettings; + fetch: typeof fetch; + now: () => number; + refreshUrl: string; + drainTimeoutMs: number; +} + +interface SelectedAccount { + account: Account; + quota: AccountQuota; +} + +interface UpstreamResult { + response: Response; + controller: AbortController; + release: () => void; +} + +export class AccountPoolHub { + private accepting = false; + private readonly inFlightByAccount = new Map(); + private readonly activeControllers = new Set(); + private readonly refreshes = new Map>(); + private readonly drainWaiters = new Set<() => void>(); + + constructor(private readonly options: HubOptions) {} + + async start(signal: AbortSignal): Promise { + this.accepting = true; + await waitForAbort(signal); + await this.stop(); + } + + async stop(): Promise { + this.accepting = false; + if (this.inFlightCount() === 0) return; + let timeout: ReturnType | null = null; + await Promise.race([ + new Promise((resolve) => { + this.drainWaiters.add(resolve); + }), + new Promise((resolve) => { + timeout = setTimeout(resolve, this.options.drainTimeoutMs); + }), + ]); + if (timeout !== null) clearTimeout(timeout); + if (this.inFlightCount() === 0) return; + for (const controller of this.activeControllers) { + controller.abort( + new Error( + "Account Pool stopped before the upstream response completed.", + ), + ); + } + } + + async handle(request: Request): Promise { + if ( + !safeKeyEqual( + readBearer(request.headers.get("authorization")), + this.options.hubKey, + ) + ) { + return anthropicError( + 401, + "authentication_error", + "Invalid Account Pool bearer token.", + ); + } + if (!this.accepting) { + return anthropicError( + 503, + "api_error", + "Account Pool is not accepting requests.", + ); + } + const body = new Uint8Array(await request.arrayBuffer()); + return this.forward(request, body); + } + + async status(showKey: boolean): Promise { + const settings = this.options.getSettings(); + const now = this.options.now(); + const accounts = await this.options.accounts.list(); + return { + route: ROUTE, + enabledAccountCount: accounts.filter((account) => account.enabled).length, + inFlight: this.inFlightCount(), + accepting: this.accepting, + hubKey: showKey ? this.options.hubKey : null, + accounts: accounts.map((account) => { + const quota = this.options.quotas.get(account.id); + return { + ...account, + fiveHourUtilization: quota.fiveHourUtilization, + fiveHourResetAt: quota.fiveHourResetAt, + fiveHourStatus: quota.fiveHourStatus, + sevenDayUtilization: quota.sevenDayUtilization, + sevenDayResetAt: quota.sevenDayResetAt, + sevenDayStatus: quota.sevenDayStatus, + representativeClaim: quota.representativeClaim, + bucketExhaustion: quota.bucketExhaustion, + observedAt: quota.observedAt, + heldUntil: quota.heldUntil, + error: quota.error, + inFlight: this.inFlightByAccount.get(account.id) ?? 0, + status: accountStatus(account, quota, settings.switchThreshold, now), + }; + }), + }; + } + + private async forward(request: Request, body: Uint8Array): Promise { + const attempted = new Set(); + const accounts = await this.options.accounts.list(); + while (attempted.size < accounts.length) { + const selected = await this.select(attempted); + if (selected === null) return this.noEligibleResponse(accounts); + attempted.add(selected.account.id); + let secret: AccountSecret; + try { + secret = await this.freshSecret(selected.account); + } catch (error) { + this.markError(selected.account.id, errorMessage(error)); + continue; + } + let upstream: UpstreamResult; + try { + upstream = await this.fetchUpstream( + request, + body, + selected.account, + secret, + ); + } catch (error) { + return anthropicError( + 502, + "api_error", + "Account Pool could not reach Anthropic.", + ); + } + const observed = quotaFromHeaders( + selected.account.id, + upstream.response.headers, + this.options.quotas.get(selected.account.id), + this.options.now(), + ); + this.options.quotas.put(observed); + if ( + upstream.response.status === 429 && + isQuotaRejection(observed, this.options.now()) + ) { + await upstream.response.body?.cancel(); + upstream.release(); + continue; + } + if (upstream.response.status === 429) { + const waitMs = retryAfterMilliseconds( + upstream.response.headers.get("retry-after"), + this.options.now(), + ); + this.options.quotas.put({ + ...observed, + heldUntil: this.options.now() + waitMs, + }); + if (waitMs <= MAX_INLINE_HOLD_MS) { + await upstream.response.body?.cancel(); + upstream.release(); + await delay(waitMs); + try { + const retry = await this.fetchUpstream( + request, + body, + selected.account, + secret, + ); + const retryQuota = quotaFromHeaders( + selected.account.id, + retry.response.headers, + this.options.quotas.get(selected.account.id), + this.options.now(), + ); + if (retry.response.status === 429) { + const retryWaitMs = retryAfterMilliseconds( + retry.response.headers.get("retry-after"), + this.options.now(), + ); + this.options.quotas.put({ + ...retryQuota, + heldUntil: this.options.now() + retryWaitMs, + }); + } else { + this.options.quotas.put(retryQuota); + } + if ( + retry.response.status === 401 || + retry.response.status === 403 + ) { + const detail = await retry.response + .clone() + .text() + .catch(() => ""); + this.markError( + selected.account.id, + detail.trim() || + `Anthropic returned HTTP ${retry.response.status}.`, + ); + } + return this.clientResponse(retry); + } catch { + return anthropicError( + 502, + "api_error", + "Account Pool could not reach Anthropic.", + ); + } + } + } + if ( + upstream.response.status === 401 || + upstream.response.status === 403 + ) { + const detail = await upstream.response + .clone() + .text() + .catch(() => ""); + this.markError( + selected.account.id, + detail.trim() || + `Anthropic returned HTTP ${upstream.response.status}.`, + ); + } + return this.clientResponse(upstream); + } + return this.noEligibleResponse(accounts); + } + + private async select( + attempted: ReadonlySet, + ): Promise { + const now = this.options.now(); + const threshold = this.options.getSettings().switchThreshold; + const candidates = (await this.options.accounts.list()) + .filter((account) => account.enabled && !attempted.has(account.id)) + .map((account) => ({ + account, + quota: this.options.quotas.get(account.id), + })) + .filter(({ quota }) => quota.error === null) + .filter(({ quota }) => quota.heldUntil === null || quota.heldUntil <= now) + .filter(({ quota }) => !isQuotaExhausted(quota, threshold, now)); + candidates.sort((left, right) => { + const priority = left.account.priority - right.account.priority; + if (priority !== 0) return priority; + const inFlight = + (this.inFlightByAccount.get(left.account.id) ?? 0) - + (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) + ); + }); + return candidates[0] ?? null; + } + + private async freshSecret(account: Account): Promise { + const secret = await this.options.accounts.readSecret(account.id); + if ( + secret.kind !== "oauth" || + secret.expiresAt === null || + secret.expiresAt > this.options.now() + REFRESH_WINDOW_MS + ) { + return secret; + } + const existing = this.refreshes.get(account.id); + if (existing !== undefined) return existing; + const refresh = this.refresh(account.id, secret).finally(() => { + this.refreshes.delete(account.id); + }); + this.refreshes.set(account.id, refresh); + return refresh; + } + + private async refresh( + accountId: string, + secret: Extract, + ): Promise { + const response = await this.options.fetch(this.options.refreshUrl, { + method: "POST", + headers: { + "content-type": "application/json", + accept: "application/json", + }, + body: JSON.stringify({ + grant_type: "refresh_token", + refresh_token: secret.refreshToken, + client_id: OAUTH_CLIENT_ID, + }), + }); + if (!response.ok) { + throw new Error(`OAuth refresh failed with HTTP ${response.status}.`); + } + const parsed = refreshResponseSchema.parse(await response.json()); + const rawExpiresAt = + parsed.expires_at ?? + (parsed.expires_in === undefined + ? this.options.now() + 60 * 60 * 1_000 + : this.options.now() + parsed.expires_in * 1_000); + const refreshed: AccountSecret = { + kind: "oauth", + accessToken: parsed.access_token, + refreshToken: parsed.refresh_token ?? secret.refreshToken, + expiresAt: + rawExpiresAt < 1_000_000_000_000 + ? Math.round(rawExpiresAt * 1_000) + : Math.round(rawExpiresAt), + }; + await this.options.accounts.writeSecret(accountId, refreshed); + const quota = this.options.quotas.get(accountId); + this.options.quotas.put({ ...quota, error: null }); + return refreshed; + } + + private async fetchUpstream( + request: Request, + body: Uint8Array, + account: Account, + secret: AccountSecret, + ): Promise { + const requestUrl = new URL(request.url); + const mountedPath = requestUrl.pathname.indexOf("/http/"); + const upstreamPath = + mountedPath < 0 + ? requestUrl.pathname + : requestUrl.pathname.slice(mountedPath + 5); + const upstreamUrl = new URL( + upstreamPath.replace(/^\//u, "") + requestUrl.search, + ensureSlash(this.options.getSettings().upstreamBaseUrl), + ); + const headers = new Headers(); + for (const [name, value] of request.headers) { + if (isAllowedRequestHeader(name)) headers.append(name, value); + } + if (secret.kind === "oauth") { + headers.set("authorization", `Bearer ${secret.accessToken}`); + } else { + headers.set("x-api-key", secret.apiKey); + } + const controller = new AbortController(); + this.activeControllers.add(controller); + this.increment(account.id); + let released = false; + const release = () => { + if (released) return; + released = true; + this.activeControllers.delete(controller); + this.decrement(account.id); + }; + try { + const upstreamBody = new ArrayBuffer(body.byteLength); + new Uint8Array(upstreamBody).set(body); + const response = await this.options.fetch(upstreamUrl, { + method: request.method, + headers, + body: upstreamBody, + signal: controller.signal, + }); + return { response, controller, release }; + } catch (error) { + release(); + throw error; + } + } + + private clientResponse(upstream: UpstreamResult): Response { + const headers = new Headers(); + for (const [name, value] of upstream.response.headers) { + if (!DROPPED_RESPONSE_HEADERS.has(name.toLowerCase())) + headers.append(name, value); + } + if (upstream.response.body === null) { + upstream.release(); + return new Response(null, { + status: upstream.response.status, + statusText: upstream.response.statusText, + headers, + }); + } + const reader = upstream.response.body.getReader(); + const eventStream = + upstream.response.headers + .get("content-type") + ?.split(";", 1)[0] + ?.trim() + .toLowerCase() === "text/event-stream"; + const body = new ReadableStream({ + async pull(controller) { + try { + const chunk = await reader.read(); + if (chunk.done) { + upstream.release(); + controller.close(); + return; + } + controller.enqueue(chunk.value); + } catch (error) { + upstream.release(); + if (eventStream) { + const message = errorMessage(error); + controller.enqueue( + new TextEncoder().encode( + `event: error\ndata: ${JSON.stringify({ type: "error", error: { type: "api_error", message } })}\n\n`, + ), + ); + controller.close(); + } else { + controller.error( + error instanceof Error ? error : new Error(String(error)), + ); + } + } + }, + async cancel() { + upstream.controller.abort(); + await reader.cancel().catch(() => undefined); + upstream.release(); + }, + }); + return new Response(body, { + status: upstream.response.status, + statusText: upstream.response.statusText, + headers, + }); + } + + private noEligibleResponse(accounts: readonly Account[]): Response { + if (!accounts.some((account) => account.enabled)) { + return anthropicError( + 503, + "api_error", + "Account Pool has no enabled account", + ); + } + const now = this.options.now(); + const next = accounts + .filter((account) => account.enabled) + .flatMap((account) => { + const quota = this.options.quotas.get(account.id); + return [ + quota.heldUntil, + quota.fiveHourResetAt, + quota.sevenDayResetAt, + ].filter((value): value is number => value !== null && value > now); + }) + .sort((left, right) => left - right)[0]; + const retryAfter = Math.max( + 1, + Math.ceil(((next ?? now + 1_000) - now) / 1_000), + ); + return anthropicError( + 429, + "rate_limit_error", + "No Account Pool account is currently eligible.", + { "retry-after": String(retryAfter) }, + ); + } + + private markError(accountId: string, message: string): void { + const quota = this.options.quotas.get(accountId); + this.options.quotas.put({ ...quota, error: message.slice(0, 1_000) }); + } + + private increment(accountId: string): void { + this.inFlightByAccount.set( + accountId, + (this.inFlightByAccount.get(accountId) ?? 0) + 1, + ); + } + + private decrement(accountId: string): void { + const next = Math.max(0, (this.inFlightByAccount.get(accountId) ?? 1) - 1); + if (next === 0) this.inFlightByAccount.delete(accountId); + else this.inFlightByAccount.set(accountId, next); + if (this.inFlightCount() !== 0) return; + for (const resolve of this.drainWaiters) resolve(); + this.drainWaiters.clear(); + } + + private inFlightCount(): number { + let total = 0; + for (const count of this.inFlightByAccount.values()) total += count; + return total; + } +} + +export function createHub(options: { + accounts: AccountStore; + quotas: QuotaStore; + hubKey: string; + getSettings: () => HubSettings; + fetch?: typeof fetch; + now?: () => number; + refreshUrl?: string; + drainTimeoutMs?: number; +}): AccountPoolHub { + return new AccountPoolHub({ + accounts: options.accounts, + quotas: options.quotas, + hubKey: options.hubKey, + getSettings: options.getSettings, + fetch: options.fetch ?? fetch, + now: options.now ?? Date.now, + refreshUrl: options.refreshUrl ?? DEFAULT_REFRESH_URL, + drainTimeoutMs: options.drainTimeoutMs ?? 60_000, + }); +} + +function isAllowedRequestHeader(name: string): boolean { + const normalized = name.toLowerCase(); + return ( + ALLOWED_REQUEST_HEADERS.has(normalized) || + ALLOWED_REQUEST_HEADER_PREFIXES.some((prefix) => + normalized.startsWith(prefix), + ) + ); +} + +function ensureSlash(value: string): string { + return value.endsWith("/") ? value : `${value}/`; +} + +function readBearer(value: string | null): string | null { + if (value === null) return null; + const match = /^Bearer\s+(.+)$/iu.exec(value); + return match?.[1] ?? null; +} + +function safeKeyEqual(left: string | null, right: string): boolean { + if (left === null) return false; + const leftBytes = Buffer.from(left); + const rightBytes = Buffer.from(right); + return ( + leftBytes.length === rightBytes.length && + timingSafeEqual(leftBytes, rightBytes) + ); +} + +function anthropicError( + status: number, + type: string, + message: string, + headers?: HeadersInit, +): Response { + return Response.json( + { type: "error", error: { type, message } }, + { status, headers }, + ); +} + +function errorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} + +function waitForAbort(signal: AbortSignal): Promise { + if (signal.aborted) return Promise.resolve(); + return new Promise((resolve) => { + signal.addEventListener("abort", () => resolve(), { once: true }); + }); +} + +function delay(milliseconds: number): Promise { + return new Promise((resolve) => setTimeout(resolve, milliseconds)); +} diff --git a/plugins/account-pool/src/operations.ts b/plugins/account-pool/src/operations.ts new file mode 100644 index 0000000000..c19f90cdee --- /dev/null +++ b/plugins/account-pool/src/operations.ts @@ -0,0 +1,80 @@ +import type { Account, AccountSummary, PoolStatus } from "./contracts.js"; +import type { AccountAddInput } from "./contracts.js"; +import { + importClaudeCredentials, + type ImportedClaudeCredentials, +} from "./credentials.js"; +import type { AccountPoolHub } from "./hub.js"; +import type { AccountStore, QuotaStore } from "./store.js"; + +export class PoolOperations { + constructor( + private readonly accounts: AccountStore, + private readonly quotas: QuotaStore, + private readonly hub: AccountPoolHub, + private readonly importCredentials: () => Promise = importClaudeCredentials, + ) {} + + async add(input: AccountAddInput): Promise { + if (input.source.kind === "api-key") { + return this.accounts.add( + { + provider: input.provider, + kind: "api-key", + label: input.label ?? "Claude API key", + email: null, + subscriptionType: null, + rateLimitTier: null, + enabled: true, + priority: input.priority, + }, + { kind: "api-key", apiKey: input.source.apiKey }, + ); + } + const imported = await this.importCredentials(); + return this.accounts.add( + { + provider: input.provider, + kind: "oauth", + label: input.label ?? imported.email ?? "Claude Code account", + email: imported.email, + subscriptionType: imported.subscriptionType, + rateLimitTier: imported.rateLimitTier, + enabled: true, + priority: input.priority, + }, + { + kind: "oauth", + accessToken: imported.accessToken, + refreshToken: imported.refreshToken, + expiresAt: imported.expiresAt, + }, + ); + } + + async list(): Promise { + return (await this.hub.status(false)).accounts; + } + + async remove(id: string): Promise { + const removed = await this.accounts.remove(id); + if (removed) this.quotas.remove(id); + return removed; + } + + async enable(id: string): Promise { + const account = await this.accounts.setEnabled(id, true); + if (account === null) return null; + const quota = this.quotas.get(id); + this.quotas.put({ ...quota, error: null, heldUntil: null }); + return account; + } + + async disable(id: string): Promise { + return this.accounts.setEnabled(id, false); + } + + status(showKey: boolean): Promise { + return this.hub.status(showKey); + } +} diff --git a/plugins/account-pool/src/quota.ts b/plugins/account-pool/src/quota.ts new file mode 100644 index 0000000000..3e3abf25fd --- /dev/null +++ b/plugins/account-pool/src/quota.ts @@ -0,0 +1,144 @@ +import type { Account, AccountQuota, AccountSummary } from "./contracts.js"; + +const PREFIX = "anthropic-ratelimit-unified-"; + +function parseNumber(value: string | null): number | null { + if (value === null || value.trim() === "") return null; + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : null; +} + +function parseReset(value: string | null): number | null { + if (value === null || value.trim() === "") return null; + const numeric = Number(value); + if (Number.isFinite(numeric)) { + return numeric < 1_000_000_000_000 + ? Math.round(numeric * 1_000) + : Math.round(numeric); + } + const parsed = Date.parse(value); + return Number.isNaN(parsed) ? null : parsed; +} + +export function quotaFromHeaders( + accountId: string, + headers: Headers, + previous: AccountQuota, + now: number, +): AccountQuota { + const fiveHourUtilization = parseNumber( + headers.get(`${PREFIX}5h-utilization`), + ); + const fiveHourResetAt = parseReset(headers.get(`${PREFIX}5h-reset`)); + const fiveHourStatus = headers.get(`${PREFIX}5h-status`); + const sevenDayUtilization = parseNumber( + headers.get(`${PREFIX}7d-utilization`), + ); + 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 observed = + fiveHourUtilization !== null || + fiveHourResetAt !== null || + fiveHourStatus !== null || + sevenDayUtilization !== null || + sevenDayResetAt !== null || + sevenDayStatus !== null || + representativeClaim !== null; + return { + accountId, + fiveHourUtilization: fiveHourUtilization ?? previous.fiveHourUtilization, + fiveHourResetAt: fiveHourResetAt ?? previous.fiveHourResetAt, + fiveHourStatus: fiveHourStatus ?? previous.fiveHourStatus, + sevenDayUtilization: sevenDayUtilization ?? previous.sevenDayUtilization, + sevenDayResetAt: sevenDayResetAt ?? previous.sevenDayResetAt, + sevenDayStatus: sevenDayStatus ?? previous.sevenDayStatus, + representativeClaim: representativeClaim ?? previous.representativeClaim, + bucketExhaustion, + observedAt: observed ? now : previous.observedAt, + heldUntil: previous.heldUntil, + error: previous.error, + }; +} + +function activeWindow( + utilization: number | null, + status: string | null, + resetAt: number | null, + threshold: number, + now: number, +): boolean { + if (resetAt !== null && resetAt <= now) return false; + return ( + status?.toLowerCase() === "rejected" || + (utilization !== null && utilization >= threshold) + ); +} + +export function isQuotaExhausted( + quota: AccountQuota, + threshold: number, + now: number, +): boolean { + return ( + activeWindow( + quota.fiveHourUtilization, + quota.fiveHourStatus, + quota.fiveHourResetAt, + threshold, + now, + ) || + activeWindow( + quota.sevenDayUtilization, + quota.sevenDayStatus, + quota.sevenDayResetAt, + threshold, + now, + ) + ); +} + +export function isQuotaRejection(quota: AccountQuota, now: number): boolean { + return ( + activeWindow(null, quota.fiveHourStatus, quota.fiveHourResetAt, 1, now) || + activeWindow(null, quota.sevenDayStatus, quota.sevenDayResetAt, 1, now) + ); +} + +export function accountStatus( + account: Account, + quota: AccountQuota, + threshold: number, + now: number, +): AccountSummary["status"] { + 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"; + return "ready"; +} + +export function retryAfterMilliseconds( + value: string | null, + now: number, +): number { + if (value === null) return 1_000; + const seconds = Number(value); + if (Number.isFinite(seconds)) return Math.max(0, Math.ceil(seconds * 1_000)); + const date = Date.parse(value); + return Number.isNaN(date) ? 1_000 : Math.max(0, date - now); +} diff --git a/plugins/account-pool/src/rpc.ts b/plugins/account-pool/src/rpc.ts new file mode 100644 index 0000000000..600955708e --- /dev/null +++ b/plugins/account-pool/src/rpc.ts @@ -0,0 +1,55 @@ +import { defineRpcContract } from "@get-bb/plugin-sdk"; +import { z } from "zod"; +import { + accountAddInputSchema, + accountIdInputSchema, + accountSchema, + accountSummarySchema, + statusSchema, +} from "./contracts.js"; +import type { PoolOperations } from "./operations.js"; + +export const accountPoolRpcContract = defineRpcContract({ + "account.add": { + input: accountAddInputSchema, + output: accountSchema, + }, + "account.list": { + input: z.null(), + output: z.array(accountSummarySchema), + }, + "account.remove": { + input: accountIdInputSchema, + output: z.object({ removed: z.boolean() }).strict(), + }, + "account.enable": { + input: accountIdInputSchema, + output: z.object({ account: accountSchema.nullable() }).strict(), + }, + "account.disable": { + input: accountIdInputSchema, + output: z.object({ account: accountSchema.nullable() }).strict(), + }, + status: { + input: z.null(), + output: statusSchema, + }, +}); + +export function createRpcHandlers(operations: PoolOperations) { + return { + "account.add": (input: Parameters[0]) => + operations.add(input), + "account.list": () => operations.list(), + "account.remove": async ({ id }: { id: string }) => ({ + removed: await operations.remove(id), + }), + "account.enable": async ({ id }: { id: string }) => ({ + account: await operations.enable(id), + }), + "account.disable": async ({ id }: { id: string }) => ({ + account: await operations.disable(id), + }), + status: () => operations.status(false), + }; +} diff --git a/plugins/account-pool/src/server.test.ts b/plugins/account-pool/src/server.test.ts new file mode 100644 index 0000000000..9ba8d8fffd --- /dev/null +++ b/plugins/account-pool/src/server.test.ts @@ -0,0 +1,778 @@ +import fs from "node:fs/promises"; +import http, { type IncomingMessage, type ServerResponse } from "node:http"; +import path from "node:path"; +import { mkdtemp } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { createFakePluginHost } from "@get-bb/plugin-sdk/testing"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { + accountSchema, + accountSecretSchema, + accountSummarySchema, + statusSchema, + type AccountSummary, +} from "./contracts.js"; +import { z } from "zod"; +import type { ImportedClaudeCredentials } from "./credentials.js"; +import { + createAccountPoolPlugin, + helloResponse, + type AccountPoolPluginOptions, +} from "./server.js"; + +type UpstreamHandler = ( + request: IncomingMessage, + response: ServerResponse, +) => void | Promise; + +interface Upstream { + url: string; + close: () => Promise; +} + +interface Fixture { + dataDir: string; + host: ReturnType; + service: ReturnType< + ReturnType["harness"]["behavior"]["runService"] + >; + key: string; + account: AccountSummary; +} + +const cleanups: Array<() => Promise> = []; + +afterEach(async () => { + while (cleanups.length > 0) await cleanups.pop()?.(); +}); + +async function startUpstream(handler: UpstreamHandler): Promise { + const server = http.createServer((request, response) => { + Promise.resolve(handler(request, response)).catch((error) => { + response.statusCode = 500; + response.end(error instanceof Error ? error.message : String(error)); + }); + }); + await new Promise((resolve) => { + server.listen(0, "127.0.0.1", resolve); + }); + const address = server.address(); + if (address === null || typeof address === "string") { + throw new Error("Fake upstream did not bind a TCP port."); + } + const url = `http://127.0.0.1:${address.port}`; + return { + url, + close: async () => { + server.closeAllConnections(); + await new Promise((resolve, reject) => { + server.close((error) => { + if (error) reject(error); + else resolve(); + }); + }); + }, + }; +} + +async function readRequestBody(request: IncomingMessage): Promise { + const chunks: Buffer[] = []; + for await (const chunk of request) { + chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)); + } + return Buffer.concat(chunks); +} + +function importedCredentials( + overrides: Partial = {}, +): ImportedClaudeCredentials { + return { + accessToken: "oauth-access", + refreshToken: "oauth-refresh", + expiresAt: Date.now() + 60 * 60 * 1_000, + subscriptionType: "max", + rateLimitTier: "max_5x", + email: "pool@example.com", + ...overrides, + }; +} + +async function createFixture(args: { + upstreamUrl: string; + options?: AccountPoolPluginOptions; + source?: "api-key" | "import"; + apiKey?: string; + priority?: number; +}): Promise { + const dataDir = await mkdtemp(path.join(tmpdir(), "bb-account-pool-")); + const host = createFakePluginHost({ + pluginId: "account-pool", + dataDir, + settings: { + upstreamBaseUrl: args.upstreamUrl, + switchThreshold: 0.98, + }, + }); + const plugin = createAccountPoolPlugin(args.options); + await plugin(host.bb); + const accountMetadata = accountSchema.parse( + await host.harness.behavior.callRpc("account.add", { + provider: "claude", + source: + args.source === "import" + ? { kind: "import" } + : { kind: "api-key", apiKey: args.apiKey ?? "sk-account" }, + label: null, + priority: args.priority ?? 100, + }), + ); + const service = host.harness.behavior.runService("hub"); + await vi.waitFor(async () => { + const result = await host.harness.behavior.runCli([ + "status", + "--json", + "--show-key", + ]); + expect(result.exitCode).toBe(0); + expect(statusSchema.parse(JSON.parse(result.stdout)).accepting).toBe(true); + }); + const statusResult = await host.harness.behavior.runCli([ + "status", + "--json", + "--show-key", + ]); + const status = statusSchema.parse(JSON.parse(statusResult.stdout)); + if (status.hubKey === null) throw new Error("Hub key was not returned."); + const account = status.accounts.find( + (candidate) => candidate.id === accountMetadata.id, + ); + if (account === undefined) throw new Error("Added account was not listed."); + cleanups.push(async () => { + service.controller.abort(); + await service.done; + await host.harness.lifecycle.dispose(); + await fs.rm(dataDir, { recursive: true, force: true }); + }); + return { dataDir, host, service, key: status.hubKey, account }; +} + +function authHeaders(key: string): Record { + return { + authorization: `Bearer ${key}`, + "content-type": "application/json", + "anthropic-version": "2023-06-01", + }; +} + +async function addApiAccount( + fixture: Fixture, + apiKey: string, + priority = 100, +): Promise { + const added = accountSchema.parse( + await fixture.host.harness.behavior.callRpc("account.add", { + provider: "claude", + source: { kind: "api-key", apiKey }, + label: apiKey, + priority, + }), + ); + const list = z + .array(accountSummarySchema) + .parse(await fixture.host.harness.behavior.callRpc("account.list", null)); + const found = list.find((account) => account.id === added.id); + if (found === undefined) throw new Error("Added account was not listed."); + return found; +} + +describe("Account Pool plugin", () => { + it("forwards the next request after adding the first account through the CLI", async () => { + let forwarded = 0; + const upstream = await startUpstream(async (request, response) => { + forwarded += 1; + await readRequestBody(request); + response.writeHead(200, { "content-type": "application/json" }); + response.end('{"forwarded":true}'); + }); + cleanups.push(upstream.close); + const dataDir = await mkdtemp( + path.join(tmpdir(), "bb-account-pool-empty-"), + ); + const host = createFakePluginHost({ + pluginId: "account-pool", + dataDir, + settings: { upstreamBaseUrl: upstream.url, switchThreshold: 0.98 }, + }); + await createAccountPoolPlugin()(host.bb); + const service = host.harness.behavior.runService("hub"); + cleanups.push(async () => { + service.controller.abort(); + await service.done; + await host.harness.lifecycle.dispose(); + await fs.rm(dataDir, { recursive: true, force: true }); + }); + const statusResult = await host.harness.behavior.runCli([ + "status", + "--json", + "--show-key", + ]); + const status = statusSchema.parse(JSON.parse(statusResult.stdout)); + expect(status.accepting).toBe(true); + if (status.hubKey === null) throw new Error("Expected a hub key."); + expect(host.harness.inspection.needsConfigurationMessages).toEqual([ + "Add and enable a Claude account with `bb pool account add`.", + ]); + const hello = helloResponse(); + expect(hello.status).toBe(200); + const unavailable = await host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { headers: authHeaders(status.hubKey), body: "{}" }, + ); + expect(unavailable.status).toBe(503); + expect(await unavailable.json()).toEqual({ + type: "error", + error: { + type: "api_error", + message: "Account Pool has no enabled account", + }, + }); + expect(forwarded).toBe(0); + const added = await host.harness.behavior.runCli([ + "account", + "add", + "--provider", + "claude", + "--api-key", + "sk-cli-secret", + "--label", + "CLI account", + "--priority", + "7", + ]); + expect(added.exitCode).toBe(0); + expect(added.stdout).not.toContain("sk-cli-secret"); + expect(added.stdout).not.toContain("reload"); + const forwardedResponse = await host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { headers: authHeaders(status.hubKey), body: "{}" }, + ); + expect(forwardedResponse.status).toBe(200); + expect(await forwardedResponse.text()).toBe('{"forwarded":true}'); + expect(forwarded).toBe(1); + }); + + it("exposes every account CLI operation", async () => { + const upstream = await startUpstream(async (request, response) => { + await readRequestBody(request); + response.writeHead(200, { "content-type": "application/json" }); + response.end("{}"); + }); + cleanups.push(upstream.close); + const fixture = await createFixture({ upstreamUrl: upstream.url }); + const help = await fixture.host.harness.behavior.runCli([ + "account", + "add", + "--help", + ]); + expect(help.exitCode).toBe(0); + expect(help.stdout).toContain("--api-key-stdin"); + expect(help.stdout).toContain("Unsafe: exposes the key"); + const list = await fixture.host.harness.behavior.runCli([ + "account", + "list", + "--json", + ]); + const listed = z + .object({ accounts: z.array(accountSummarySchema) }) + .strict() + .parse(JSON.parse(list.stdout)); + const account = listed.accounts[0]; + if (account === undefined) throw new Error("CLI account was not listed."); + expect(account).toMatchObject({ label: "Claude API key", priority: 100 }); + expect( + ( + await fixture.host.harness.behavior.runCli([ + "account", + "disable", + account.id, + ]) + ).exitCode, + ).toBe(0); + expect( + z + .array(accountSummarySchema) + .parse( + await fixture.host.harness.behavior.callRpc("account.list", null), + )[0]?.status, + ).toBe("disabled"); + expect( + ( + await fixture.host.harness.behavior.runCli([ + "account", + "enable", + account.id, + ]) + ).exitCode, + ).toBe(0); + const publicStatus = statusSchema.parse( + JSON.parse( + (await fixture.host.harness.behavior.runCli(["status", "--json"])) + .stdout, + ), + ); + expect(publicStatus.accepting).toBe(true); + expect(publicStatus.hubKey).toBeNull(); + const secretStatus = statusSchema.parse( + JSON.parse( + ( + await fixture.host.harness.behavior.runCli([ + "status", + "--json", + "--show-key", + ]) + ).stdout, + ), + ); + if (secretStatus.hubKey === null) throw new Error("Expected a hub key."); + const counted = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages/count_tokens", + { headers: authHeaders(secretStatus.hubKey), body: "{}" }, + ); + expect(counted.status).toBe(200); + expect(await counted.text()).toBe("{}"); + expect( + ( + await fixture.host.harness.behavior.runCli([ + "account", + "remove", + account.id, + ]) + ).exitCode, + ).toBe(0); + expect( + await fixture.host.harness.behavior.callRpc("account.list", null), + ).toEqual([]); + }); + + it("requires the hub key and forwards a streaming SSE response byte for byte", async () => { + const seen: { + url: string; + authorization: string | undefined; + clientApiKey: string | undefined; + beta: string | undefined; + version: string | undefined; + userAgent: string | undefined; + app: string | undefined; + stainlessRetry: string | undefined; + cookie: string | undefined; + gateMachineId: string | undefined; + forwarded: string | undefined; + cfRay: string | undefined; + body: Buffer; + }[] = []; + const first = Buffer.from('event: message_start\ndata: {"one":1}\n\n'); + const second = Buffer.from('event: message_stop\ndata: {"two":2}\n\n'); + const upstream = await startUpstream(async (request, response) => { + seen.push({ + url: request.url ?? "", + authorization: request.headers.authorization, + clientApiKey: request.headers["x-api-key"]?.toString(), + beta: request.headers["anthropic-beta"]?.toString(), + version: request.headers["anthropic-version"]?.toString(), + userAgent: request.headers["user-agent"]?.toString(), + app: request.headers["x-app"]?.toString(), + stainlessRetry: request.headers["x-stainless-retry-count"]?.toString(), + cookie: request.headers.cookie, + gateMachineId: request.headers["x-bb-gate-machine-id"]?.toString(), + forwarded: request.headers.forwarded, + cfRay: request.headers["cf-ray"]?.toString(), + body: await readRequestBody(request), + }); + response.writeHead(200, { + "content-type": "text/event-stream", + "anthropic-ratelimit-unified-5h-utilization": "0.25", + "anthropic-ratelimit-unified-5h-reset": "4102444800", + "anthropic-ratelimit-unified-5h-status": "allowed", + "anthropic-ratelimit-unified-7d-utilization": "0.5", + "anthropic-ratelimit-unified-7d-reset": "4102448400", + "anthropic-ratelimit-unified-7d-status": "allowed", + "anthropic-ratelimit-unified-representative-claim": "claim-a", + "anthropic-ratelimit-unified-status": "rejected", + "anthropic-ratelimit-unified-overage-status": "rejected", + "anthropic-ratelimit-unified-7d_oi-status": "rejected", + "anthropic-ratelimit-unified-7d_oi-reset": "4102452000", + }); + response.write(first); + setTimeout(() => response.end(second), 60); + }); + cleanups.push(upstream.close); + const fixture = await createFixture({ + upstreamUrl: upstream.url, + source: "import", + options: { importCredentials: async () => importedCredentials() }, + }); + const unauthorized = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { headers: { "content-type": "application/json" }, body: "{}" }, + ); + expect(unauthorized.status).toBe(401); + expect(seen).toHaveLength(0); + const body = Buffer.from('{"model":"claude-test","stream":true}'); + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages?beta=true", + { + headers: { + ...authHeaders(fixture.key), + "x-api-key": "client-key-must-not-forward", + "anthropic-beta": "feature-a,feature-b", + "user-agent": "claude-code-test", + "x-app": "cli", + "x-stainless-retry-count": "2", + cookie: "bb_session=browser-secret", + "x-bb-gate-machine-id": "machine-stable-id", + forwarded: "for=192.0.2.1", + "cf-ray": "edge-request-id", + "accept-encoding": "gzip", + }, + body, + }, + ); + expect(response.status).toBe(200); + const reader = response.body?.getReader(); + if (reader === undefined) throw new Error("Expected an SSE response body."); + const firstRead = await Promise.race([ + reader.read(), + new Promise((_, reject) => + setTimeout( + () => reject(new Error("First SSE chunk was buffered.")), + 30, + ), + ), + ]); + expect(firstRead.done).toBe(false); + expect(Buffer.from(firstRead.value ?? []).equals(first)).toBe(true); + const remaining: Buffer[] = []; + while (true) { + const chunk = await reader.read(); + if (chunk.done) break; + remaining.push(Buffer.from(chunk.value)); + } + expect( + Buffer.concat([Buffer.from(firstRead.value ?? []), ...remaining]), + ).toEqual(Buffer.concat([first, second])); + expect(seen).toEqual([ + { + url: "/v1/messages?beta=true", + authorization: "Bearer oauth-access", + clientApiKey: undefined, + beta: "feature-a,feature-b", + version: "2023-06-01", + userAgent: "claude-code-test", + app: "cli", + stainlessRetry: "2", + cookie: undefined, + gateMachineId: undefined, + forwarded: undefined, + cfRay: undefined, + body, + }, + ]); + const accounts = z + .array(accountSummarySchema) + .parse(await fixture.host.harness.behavior.callRpc("account.list", null)); + expect(accounts[0]).toMatchObject({ + fiveHourUtilization: 0.25, + sevenDayUtilization: 0.5, + fiveHourStatus: "allowed", + sevenDayStatus: "allowed", + representativeClaim: "claim-a", + bucketExhaustion: { "7d_oi": 4_102_452_000_000 }, + }); + expect(fixture.host.harness.inspection.registrations.httpRoutes).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + method: "HEAD", + path: "/api/hello", + auth: "none", + }), + ]), + ); + }); + + it("errors a failed non-SSE stream without appending an SSE frame", async () => { + const partial = Buffer.from('{"partial":'); + const upstream = await startUpstream(async (request, response) => { + await readRequestBody(request); + response.writeHead(200, { "content-type": "application/json" }); + response.write(partial); + setTimeout(() => response.destroy(new Error("upstream failed")), 30); + }); + cleanups.push(upstream.close); + const fixture = await createFixture({ upstreamUrl: upstream.url }); + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { headers: authHeaders(fixture.key), body: "{}" }, + ); + const reader = response.body?.getReader(); + if (reader === undefined) throw new Error("Expected a streaming body."); + const first = await reader.read(); + expect(first.done).toBe(false); + expect(Buffer.from(first.value ?? [])).toEqual(partial); + await expect(reader.read()).rejects.toThrow(); + }); + + it("skips threshold-exhausted accounts and rotates quota rejections", async () => { + const keys: string[] = []; + let requestNumber = 0; + const upstream = await startUpstream((request, response) => { + requestNumber += 1; + keys.push(request.headers["x-api-key"]?.toString() ?? ""); + if (requestNumber === 1) { + response.writeHead(200, { + "content-type": "application/json", + "anthropic-ratelimit-unified-5h-utilization": "0.99", + "anthropic-ratelimit-unified-5h-reset": "4102444800", + "anthropic-ratelimit-unified-5h-status": "allowed", + }); + response.end('{"first":true}'); + return; + } + if (requestNumber === 2) { + response.writeHead(429, { + "content-type": "application/json", + "anthropic-ratelimit-unified-5h-status": "rejected", + "anthropic-ratelimit-unified-5h-reset": "4102444800", + }); + response.end('{"rejected":true}'); + return; + } + response.writeHead(200, { "content-type": "application/json" }); + response.end('{"rotated":true}'); + }); + cleanups.push(upstream.close); + const fixture = await createFixture({ + upstreamUrl: upstream.url, + apiKey: "sk-one", + }); + await addApiAccount(fixture, "sk-two"); + await addApiAccount(fixture, "sk-three"); + const first = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { + headers: authHeaders(fixture.key), + body: "{}", + }, + ); + expect(first.status).toBe(200); + await first.text(); + const rotated = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { + headers: authHeaders(fixture.key), + body: "{}", + }, + ); + expect(rotated.status).toBe(200); + expect(await rotated.text()).toBe('{"rotated":true}'); + expect(keys).toEqual(["sk-one", "sk-two", "sk-three"]); + }); + + it("paces a per-minute 429 on the same account without rotating", async () => { + const keys: string[] = []; + const times: number[] = []; + const upstream = await startUpstream((request, response) => { + keys.push(request.headers["x-api-key"]?.toString() ?? ""); + times.push(Date.now()); + if (keys.length === 1) { + response.writeHead(429, { + "content-type": "application/json", + "retry-after": "0.04", + }); + response.end('{"minute":true}'); + return; + } + response.writeHead(200, { "content-type": "application/json" }); + response.end('{"retried":true}'); + }); + cleanups.push(upstream.close); + const fixture = await createFixture({ + upstreamUrl: upstream.url, + apiKey: "sk-one", + }); + await addApiAccount(fixture, "sk-two"); + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { + headers: authHeaders(fixture.key), + body: "{}", + }, + ); + expect(response.status).toBe(200); + expect(await response.text()).toBe('{"retried":true}'); + expect(keys).toEqual(["sk-one", "sk-one"]); + expect((times[1] ?? 0) - (times[0] ?? 0)).toBeGreaterThanOrEqual(30); + }); + + it("serializes refresh, writes new tokens with 0600 mode, and uses them", async () => { + let refreshCalls = 0; + const authorizations: Array = []; + const upstream = await startUpstream(async (request, response) => { + if (request.url === "/oauth/token") { + refreshCalls += 1; + await new Promise((resolve) => setTimeout(resolve, 25)); + response.writeHead(200, { "content-type": "application/json" }); + response.end( + JSON.stringify({ + access_token: "oauth-new", + refresh_token: "refresh-new", + expires_in: 3600, + }), + ); + return; + } + authorizations.push(request.headers.authorization); + await readRequestBody(request); + response.writeHead(200, { "content-type": "application/json" }); + response.end("{}"); + }); + cleanups.push(upstream.close); + const fixture = await createFixture({ + upstreamUrl: upstream.url, + source: "import", + options: { + importCredentials: async () => + importedCredentials({ expiresAt: Date.now() + 1_000 }), + refreshUrl: `${upstream.url}/oauth/token`, + }, + }); + const requests = [1, 2].map(() => + fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { + headers: authHeaders(fixture.key), + body: "{}", + }), + ); + const responses = await Promise.all(requests); + await Promise.all(responses.map((response) => response.text())); + expect(refreshCalls).toBe(1); + expect(authorizations).toEqual(["Bearer oauth-new", "Bearer oauth-new"]); + const secretPath = path.join( + fixture.dataDir, + "plugins", + "account-pool", + "secrets", + "accounts", + `account-${fixture.account.id}.json`, + ); + const secret = accountSecretSchema.parse( + JSON.parse(await fs.readFile(secretPath, "utf8")), + ); + expect(secret).toMatchObject({ + kind: "oauth", + accessToken: "oauth-new", + refreshToken: "refresh-new", + }); + expect((await fs.stat(secretPath)).mode & 0o777).toBe(0o600); + }); + + it("marks refresh and upstream authorization failures as account errors", async () => { + const upstream = await startUpstream((request, response) => { + if (request.url === "/oauth/token") { + response.writeHead(401, { "content-type": "application/json" }); + response.end("{}"); + return; + } + response.writeHead(401, { "content-type": "application/json" }); + response.end('{"error":{"message":"bad account"}}'); + }); + cleanups.push(upstream.close); + const refreshFixture = await createFixture({ + upstreamUrl: upstream.url, + source: "import", + options: { + importCredentials: async () => + importedCredentials({ expiresAt: Date.now() + 1_000 }), + refreshUrl: `${upstream.url}/oauth/token`, + }, + }); + const refreshResponse = + await refreshFixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { headers: authHeaders(refreshFixture.key), body: "{}" }, + ); + expect(refreshResponse.status).toBe(429); + const refreshAccounts = z + .array(accountSummarySchema) + .parse( + await refreshFixture.host.harness.behavior.callRpc( + "account.list", + null, + ), + ); + expect(refreshAccounts[0]?.status).toBe("error"); + expect(refreshAccounts[0]?.error).toContain("OAuth refresh failed"); + + const authFixture = await createFixture({ upstreamUrl: upstream.url }); + const authResponse = await authFixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { headers: authHeaders(authFixture.key), body: "{}" }, + ); + expect(authResponse.status).toBe(401); + await authResponse.text(); + const authAccounts = z + .array(accountSummarySchema) + .parse( + await authFixture.host.harness.behavior.callRpc("account.list", null), + ); + expect(authAccounts[0]?.status).toBe("error"); + expect(authAccounts[0]?.error).toContain("bad account"); + }); + + it("drains completed streams and aborts a stuck stream after the stop deadline", async () => { + const upstream = await startUpstream((_request, response) => { + response.writeHead(200, { "content-type": "text/event-stream" }); + response.write("data: started\n\n"); + }); + cleanups.push(upstream.close); + const fixture = await createFixture({ + upstreamUrl: upstream.url, + options: { drainTimeoutMs: 40 }, + }); + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { + headers: authHeaders(fixture.key), + body: "{}", + }, + ); + const reader = response.body?.getReader(); + if (reader === undefined) throw new Error("Expected a streaming body."); + const first = await reader.read(); + expect(new TextDecoder().decode(first.value)).toBe("data: started\n\n"); + const startedAt = Date.now(); + fixture.service.controller.abort(); + await fixture.service.done; + expect(Date.now() - startedAt).toBeGreaterThanOrEqual(30); + const stopped = await reader.read(); + expect(new TextDecoder().decode(stopped.value)).toContain( + "Account Pool stopped", + ); + const rejected = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { headers: authHeaders(fixture.key), body: "{}" }, + ); + expect(rejected.status).toBe(503); + }); +}); diff --git a/plugins/account-pool/src/server.ts b/plugins/account-pool/src/server.ts new file mode 100644 index 0000000000..f7daf598fe --- /dev/null +++ b/plugins/account-pool/src/server.ts @@ -0,0 +1,117 @@ +import path from "node:path"; +import type { BbPluginApi } from "@get-bb/plugin-sdk"; +import { z } from "zod"; +import { registerPoolCli } from "./cli.js"; +import type { ImportedClaudeCredentials } from "./credentials.js"; +import { createHub } from "./hub.js"; +import { PoolOperations } from "./operations.js"; +import { accountPoolRpcContract, createRpcHandlers } from "./rpc.js"; +import { AccountStore, QUOTA_MIGRATIONS, QuotaStore } from "./store.js"; + +export interface AccountPoolPluginOptions { + fetch?: typeof fetch; + now?: () => number; + refreshUrl?: string; + drainTimeoutMs?: number; + importCredentials?: () => Promise; +} + +export function helloResponse(): Response { + return new Response(null, { status: 200 }); +} + +const upstreamSchema = z + .string() + .url() + .refine((value) => { + const protocol = new URL(value).protocol; + return protocol === "http:" || protocol === "https:"; + }, "Must be an HTTP or HTTPS URL."); + +export function createAccountPoolPlugin( + options: AccountPoolPluginOptions = {}, +) { + return async function accountPoolPlugin(bb: BbPluginApi): Promise { + const settings = bb.settings.define({ + upstreamBaseUrl: { + type: "string", + label: "Anthropic upstream base URL", + description: + "Override only for tests and QA. Production traffic uses https://api.anthropic.com.", + default: "https://api.anthropic.com", + experimental_schema: upstreamSchema, + }, + switchThreshold: { + type: "number", + label: "Quota switch threshold", + description: + "Stop selecting an account when its 5-hour or 7-day utilization reaches this fraction.", + default: 0.98, + experimental_schema: z.number().min(0).max(1), + }, + }); + let currentSettings = await settings.get(); + settings.onChange((next) => { + currentSettings = next; + }); + const secretDir = path.join( + bb.server.experimental_dataDir, + "plugins", + bb.pluginId, + "secrets", + "accounts", + ); + const accounts = new AccountStore(bb.storage.kv, secretDir); + await accounts.initialize(); + const db = bb.storage.database(); + bb.storage.migrate(db, QUOTA_MIGRATIONS); + const quotas = new QuotaStore(db); + const hubKey = await accounts.hubKey(); + const hub = createHub({ + accounts, + quotas, + hubKey, + getSettings: () => currentSettings, + fetch: options.fetch, + now: options.now, + refreshUrl: options.refreshUrl, + drainTimeoutMs: options.drainTimeoutMs, + }); + const operations = new PoolOperations( + accounts, + quotas, + hub, + options.importCredentials, + ); + if ((await accounts.list()).every((account) => !account.enabled)) { + bb.status.needsConfiguration( + "Add and enable a Claude account with `bb pool account add`.", + ); + } + bb.rpc.register(accountPoolRpcContract, createRpcHandlers(operations)); + registerPoolCli(bb, operations); + bb.http.route( + "POST", + "/v1/messages", + (context) => hub.handle(context.req.raw), + { auth: "none" }, + ); + bb.http.route( + "POST", + "/v1/messages/count_tokens", + (context) => hub.handle(context.req.raw), + { auth: "none" }, + ); + bb.http.route( + "HEAD", + "/api/hello", + () => helloResponse(), + { auth: "none" }, + ); + bb.background.service("hub", { + start: (signal) => hub.start(signal), + }); + }; +} + +export default createAccountPoolPlugin(); diff --git a/plugins/account-pool/src/store.test.ts b/plugins/account-pool/src/store.test.ts new file mode 100644 index 0000000000..9e929a4768 --- /dev/null +++ b/plugins/account-pool/src/store.test.ts @@ -0,0 +1,66 @@ +import fs from "node:fs/promises"; +import path from "node:path"; +import { tmpdir } from "node:os"; +import { mkdtemp } from "node:fs/promises"; +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"; + +const cleanups: Array<() => Promise> = []; + +afterEach(async () => { + while (cleanups.length > 0) await cleanups.pop()?.(); +}); + +function delayedAccountReads(kv: PluginKvStorage): PluginKvStorage { + return { + async get(key: string): Promise { + const value = await kv.get(key); + if (key === "accounts:v1") { + await new Promise((resolve) => setTimeout(resolve, 50)); + } + return value; + }, + set: (key, value) => kv.set(key, value), + delete: (key) => kv.delete(key), + list: (prefix) => kv.list(prefix), + }; +} + +describe("AccountStore", () => { + it("preserves both accounts added concurrently", async () => { + const dataDir = await mkdtemp(path.join(tmpdir(), "bb-account-store-")); + const host = createFakePluginHost({ pluginId: "account-pool", dataDir }); + const secretsDir = path.join(dataDir, "secrets"); + const store = new AccountStore( + delayedAccountReads(host.bb.storage.kv), + secretsDir, + ); + await store.initialize(); + cleanups.push(async () => { + await host.harness.lifecycle.dispose(); + await fs.rm(dataDir, { recursive: true, force: true }); + }); + const account = (label: string): Omit => ({ + provider: "claude", + kind: "api-key", + label, + email: null, + subscriptionType: null, + rateLimitTier: null, + enabled: true, + priority: 100, + }); + + const [first, second] = await Promise.all([ + store.add(account("first"), { kind: "api-key", apiKey: "sk-first" }), + store.add(account("second"), { kind: "api-key", apiKey: "sk-second" }), + ]); + + expect((await store.list()).map((entry) => entry.id).sort()).toEqual( + [first.id, second.id].sort(), + ); + }); +}); diff --git a/plugins/account-pool/src/store.ts b/plugins/account-pool/src/store.ts new file mode 100644 index 0000000000..9fc2b4f8f3 --- /dev/null +++ b/plugins/account-pool/src/store.ts @@ -0,0 +1,286 @@ +import { randomBytes, randomUUID } from "node:crypto"; +import fs from "node:fs/promises"; +import path from "node:path"; +import type Database from "better-sqlite3"; +import { z } from "zod"; +import type { PluginKvStorage } from "@get-bb/plugin-sdk"; +import { + accountSchema, + accountSecretSchema, + quotaSchema, + type Account, + type AccountQuota, + type AccountSecret, +} from "./contracts.js"; + +const ACCOUNTS_KEY = "accounts:v1"; +const accountsSchema = z.array(accountSchema); + +export class AccountStore { + private mutationLock: Promise | null = null; + + constructor( + private readonly kv: PluginKvStorage, + private readonly secretsDir: string, + ) {} + + async initialize(): Promise { + await fs.mkdir(this.secretsDir, { recursive: true, mode: 0o700 }); + await fs.chmod(this.secretsDir, 0o700); + } + + async list(): Promise { + const value = await this.kv.get(ACCOUNTS_KEY); + if (value === undefined) return []; + return accountsSchema.parse(value); + } + + async get(id: string): Promise { + return (await this.list()).find((account) => account.id === id) ?? null; + } + + async add( + input: Omit, + secret: AccountSecret, + ): Promise { + return this.serialized(async () => { + const account = accountSchema.parse({ + ...input, + id: randomUUID(), + createdAt: Date.now(), + }); + await this.writeSecret(account.id, secret); + try { + const accounts = await this.list(); + accounts.push(account); + await this.kv.set(ACCOUNTS_KEY, accounts); + } catch (error) { + await fs.rm(this.accountSecretPath(account.id), { force: true }); + throw error; + } + return account; + }); + } + + async remove(id: string): Promise { + return this.serialized(async () => { + const accounts = await this.list(); + const next = accounts.filter((account) => account.id !== id); + if (next.length === accounts.length) return false; + await this.kv.set(ACCOUNTS_KEY, next); + await fs.rm(this.accountSecretPath(id), { force: true }); + return true; + }); + } + + async setEnabled(id: string, enabled: boolean): 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, enabled }); + 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"), + ); + return accountSecretSchema.parse(parsed); + } + + async writeSecret(id: string, secret: AccountSecret): Promise { + await this.initialize(); + const value = `${JSON.stringify(accountSecretSchema.parse(secret))}\n`; + const destination = this.accountSecretPath(id); + const temporary = `${destination}.${randomUUID()}.tmp`; + await fs.writeFile(temporary, value, { encoding: "utf8", mode: 0o600 }); + await fs.rename(temporary, destination); + await fs.chmod(destination, 0o600); + } + + async hubKey(): Promise { + await this.initialize(); + const file = path.join(this.secretsDir, "hub-key"); + try { + const existing = (await fs.readFile(file, "utf8")).trim(); + if (existing) return existing; + } catch (error) { + if (!isMissingFile(error)) throw error; + } + const key = randomBytes(32).toString("base64url"); + const temporary = `${file}.${randomUUID()}.tmp`; + await fs.writeFile(temporary, `${key}\n`, { + encoding: "utf8", + mode: 0o600, + }); + try { + await fs.link(temporary, file); + } catch (error) { + if (!isExistingFile(error)) throw error; + } finally { + await fs.rm(temporary, { force: true }); + } + await fs.chmod(file, 0o600); + return (await fs.readFile(file, "utf8")).trim(); + } + + private accountSecretPath(id: string): string { + z.string().uuid().parse(id); + return path.join(this.secretsDir, `account-${id}.json`); + } + + private async serialized(action: () => Promise): Promise { + const previous = this.mutationLock ?? Promise.resolve(); + let release = () => {}; + const current = new Promise((resolve) => { + release = resolve; + }); + const tail = previous.then(() => current); + this.mutationLock = tail; + await previous; + try { + return await action(); + } finally { + release(); + if (this.mutationLock === tail) this.mutationLock = null; + } + } +} + +function isMissingFile(error: unknown): boolean { + return error instanceof Error && "code" in error && error.code === "ENOENT"; +} + +function isExistingFile(error: unknown): boolean { + return error instanceof Error && "code" in error && error.code === "EEXIST"; +} + +const quotaRowSchema = z + .object({ + account_id: z.string().uuid(), + five_hour_utilization: z.number().nullable(), + five_hour_reset_at: z.number().int().nullable(), + five_hour_status: z.string().nullable(), + seven_day_utilization: z.number().nullable(), + seven_day_reset_at: z.number().int().nullable(), + seven_day_status: z.string().nullable(), + representative_claim: z.string().nullable(), + bucket_exhaustion_json: z.string(), + observed_at: z.number().int().nullable(), + held_until: z.number().int().nullable(), + error: z.string().nullable(), + }) + .strict(); + +const EMPTY_QUOTA = { + fiveHourUtilization: null, + fiveHourResetAt: null, + fiveHourStatus: null, + sevenDayUtilization: null, + sevenDayResetAt: null, + sevenDayStatus: null, + representativeClaim: null, + bucketExhaustion: {}, + observedAt: null, + heldUntil: null, + error: null, +}; + +export class QuotaStore { + constructor(private readonly db: Database.Database) {} + + get(accountId: string): AccountQuota { + const row = quotaRowSchema + .optional() + .parse( + this.db + .prepare("SELECT * FROM account_quota WHERE account_id = ?") + .get(accountId), + ); + if (row === undefined) { + return quotaSchema.parse({ accountId, ...EMPTY_QUOTA }); + } + return quotaSchema.parse({ + accountId: row.account_id, + fiveHourUtilization: row.five_hour_utilization, + fiveHourResetAt: row.five_hour_reset_at, + fiveHourStatus: row.five_hour_status, + sevenDayUtilization: row.seven_day_utilization, + sevenDayResetAt: row.seven_day_reset_at, + sevenDayStatus: row.seven_day_status, + representativeClaim: row.representative_claim, + bucketExhaustion: JSON.parse(row.bucket_exhaustion_json), + observedAt: row.observed_at, + heldUntil: row.held_until, + error: row.error, + }); + } + + put(quota: AccountQuota): void { + const value = quotaSchema.parse(quota); + this.db + .prepare( + `INSERT INTO account_quota ( + 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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(account_id) DO UPDATE SET + five_hour_utilization = excluded.five_hour_utilization, + five_hour_reset_at = excluded.five_hour_reset_at, + five_hour_status = excluded.five_hour_status, + seven_day_utilization = excluded.seven_day_utilization, + 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, + observed_at = excluded.observed_at, + held_until = excluded.held_until, + error = excluded.error`, + ) + .run( + value.accountId, + value.fiveHourUtilization, + value.fiveHourResetAt, + value.fiveHourStatus, + value.sevenDayUtilization, + value.sevenDayResetAt, + value.sevenDayStatus, + value.representativeClaim, + JSON.stringify(value.bucketExhaustion), + value.observedAt, + value.heldUntil, + value.error, + ); + } + + remove(accountId: string): void { + this.db + .prepare("DELETE FROM account_quota WHERE account_id = ?") + .run(accountId); + } +} + +export const QUOTA_MIGRATIONS = [ + `CREATE TABLE account_quota ( + account_id TEXT PRIMARY KEY, + five_hour_utilization REAL, + five_hour_reset_at INTEGER, + five_hour_status TEXT, + seven_day_utilization REAL, + seven_day_reset_at INTEGER, + seven_day_status TEXT, + representative_claim TEXT, + bucket_exhaustion_json TEXT NOT NULL DEFAULT '{}', + observed_at INTEGER, + held_until INTEGER, + error TEXT + )`, +]; diff --git a/plugins/account-pool/tsconfig.json b/plugins/account-pool/tsconfig.json new file mode 100644 index 0000000000..c7a065987d --- /dev/null +++ b/plugins/account-pool/tsconfig.json @@ -0,0 +1,22 @@ +{ + "compilerOptions": { + "strict": true, + "target": "ES2022", + "module": "ESNext", + "moduleResolution": "bundler", + "lib": ["ES2022", "DOM", "DOM.Iterable"], + "esModuleInterop": true, + "noEmit": true, + "skipLibCheck": true, + "types": ["node", "vitest/globals"], + "paths": { + "@get-bb/plugin-sdk": [ + "../../packages/plugin-sdk/bundled-types/bb-plugin-sdk.d.ts" + ], + "@get-bb/plugin-sdk/testing": [ + "../../packages/plugin-sdk/bundled-types/bb-plugin-sdk-testing.d.ts" + ] + } + }, + "include": ["src/**/*.ts"] +} diff --git a/plugins/account-pool/vitest.config.ts b/plugins/account-pool/vitest.config.ts new file mode 100644 index 0000000000..6ab8633a37 --- /dev/null +++ b/plugins/account-pool/vitest.config.ts @@ -0,0 +1,9 @@ +import { defineWorkspaceTestConfig } from "../../vitest.shared.js"; + +export default defineWorkspaceTestConfig({ + test: { + name: "bb-plugin-account-pool", + include: ["**/*.test.ts"], + exclude: ["dist/**", "node_modules/**"], + }, +}); diff --git a/plugins/bb-official.json b/plugins/bb-official.json index adee7eca4f..5332a5eb2f 100644 --- a/plugins/bb-official.json +++ b/plugins/bb-official.json @@ -1,4 +1,8 @@ { + "account-pool": { + "category": "agents-and-providers", + "screenshots": [] + }, "ask-user-question": { "category": "thread-content", "screenshots": [] diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 33aabdd167..5ca3391758 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2746,6 +2746,31 @@ importers: specifier: ^3.0.0 version: 3.2.6(@types/debug@4.1.13)(@types/node@24.12.4)(jiti@2.7.0)(jsdom@29.0.1(@noble/hashes@2.4.0))(lightningcss@1.32.0)(msw@2.12.14(@types/node@24.12.4)(@typescript/typescript6@6.0.2))(terser@5.50.0)(tsx@4.23.1)(yaml@2.9.0) + plugins/account-pool: + dependencies: + zod: + specifier: 4.3.6 + version: 4.3.6 + devDependencies: + '@get-bb/plugin-sdk': + specifier: workspace:* + version: link:../../packages/plugin-sdk + '@types/better-sqlite3': + specifier: ^7.6.12 + version: 7.6.13 + '@types/node': + specifier: ^22.0.0 + version: 22.19.10 + better-sqlite3: + specifier: 12.10.0 + version: 12.10.0 + typescript: + specifier: npm:@typescript/typescript6@^6.0.2 + version: '@typescript/typescript6@6.0.2' + vitest: + specifier: ^4.1.1 + version: 4.1.1(@opentelemetry/api@1.9.1)(@types/node@22.19.10)(jsdom@29.0.1(@noble/hashes@2.4.0))(msw@2.12.14(@types/node@22.19.10)(@typescript/typescript6@6.0.2))(vite@8.0.12(@types/node@22.19.10)(esbuild@0.28.1)(jiti@2.7.0)(terser@5.50.0)(tsx@4.23.1)(yaml@2.9.0)) + plugins/ask-user-question: dependencies: '@bb/plugin-interaction-contracts': diff --git a/turbo.json b/turbo.json index d8c4fec01d..3dcf8a84d7 100644 --- a/turbo.json +++ b/turbo.json @@ -772,6 +772,9 @@ "bb-plugin-ask-user-question#typecheck": { "dependsOn": ["@get-bb/plugin-sdk#build:types", "topo"] }, + "bb-plugin-account-pool#typecheck": { + "dependsOn": ["@get-bb/plugin-sdk#build:types", "topo"] + }, "bb-plugin-automations#typecheck": { "dependsOn": ["@get-bb/plugin-sdk#build:types", "topo"] },