From b9fe105b12b3c791ddbe9d9c74cf8f1920247321 Mon Sep 17 00:00:00 2001 From: Galkurta Date: Sat, 5 Sep 2026 21:42:01 +0700 Subject: [PATCH] feat: implement stitch account pooling system with load balancing, health tracking, and transport management --- .codegraph/.gitignore | 5 + README.md | 27 + package.json | 2 +- src/commands/accounts/command.ts | 17 + src/commands/accounts/handler.ts | 63 +++ src/commands/dashboard/command.ts | 23 + src/commands/dashboard/handler.ts | 163 ++++++ src/commands/dashboard/page.ts | 46 ++ .../doctor/steps/ApiKeyDetectedStep.ts | 5 +- src/commands/pool/command.ts | 16 + src/commands/pool/handler.ts | 21 + src/commands/projects/command.ts | 29 ++ src/commands/proxy/LoggingCallToolHandler.ts | 119 +++-- src/commands/proxy/handler.ts | 97 +++- src/lib/log/append.test.ts | 7 +- .../stitch-pool/account-health-tracker.ts | 144 +++++ src/services/stitch-pool/account-pool.ts | 493 ++++++++++++++++++ src/services/stitch-pool/account-selector.ts | 53 ++ .../stitch-pool/credential-provider.ts | 126 +++++ src/services/stitch-pool/errors.ts | 101 ++++ src/services/stitch-pool/index.ts | 8 + .../stitch-pool/project-affinity-store.ts | 192 +++++++ src/services/stitch-pool/transport.ts | 132 +++++ src/services/stitch-pool/types.ts | 59 +++ .../services/stitch-pool/account-pool.test.ts | 273 ++++++++++ 25 files changed, 2159 insertions(+), 62 deletions(-) create mode 100644 .codegraph/.gitignore create mode 100644 src/commands/accounts/command.ts create mode 100644 src/commands/accounts/handler.ts create mode 100644 src/commands/dashboard/command.ts create mode 100644 src/commands/dashboard/handler.ts create mode 100644 src/commands/dashboard/page.ts create mode 100644 src/commands/pool/command.ts create mode 100644 src/commands/pool/handler.ts create mode 100644 src/commands/projects/command.ts create mode 100644 src/services/stitch-pool/account-health-tracker.ts create mode 100644 src/services/stitch-pool/account-pool.ts create mode 100644 src/services/stitch-pool/account-selector.ts create mode 100644 src/services/stitch-pool/credential-provider.ts create mode 100644 src/services/stitch-pool/errors.ts create mode 100644 src/services/stitch-pool/index.ts create mode 100644 src/services/stitch-pool/project-affinity-store.ts create mode 100644 src/services/stitch-pool/transport.ts create mode 100644 src/services/stitch-pool/types.ts create mode 100644 tests/services/stitch-pool/account-pool.test.ts diff --git a/.codegraph/.gitignore b/.codegraph/.gitignore new file mode 100644 index 0000000..d20c0fe --- /dev/null +++ b/.codegraph/.gitignore @@ -0,0 +1,5 @@ +# CodeGraph data files — local to each machine, not for committing. +# Ignore everything in .codegraph/ except this file itself, so transient +# files (the database, daemon.pid, sockets, logs) never show up in git. +* +!.gitignore diff --git a/README.md b/README.md index 49e6794..51cdcd9 100644 --- a/README.md +++ b/README.md @@ -127,12 +127,16 @@ stitch-mcp | `serve -p ` | Preview project screens locally | | `screens -p ` | Browse screens in terminal | | `view` | Interactive resource browser | +| `dashboard` | Run the local account and quota dashboard | | **Build** | | | `site -p ` | Generate Astro project from screens | | `snapshot` | Save screen state to file | | **Integration** | | | `tool [name]` | Invoke MCP tools from CLI | | `proxy` | Run MCP proxy for agents | +| `accounts list|test|enable|disable` | Manage API-key accounts | +| `pool status|reset` | Inspect or reset pool state | +| `projects` | Aggregate projects across accounts | Run any command with `--help` for full options. @@ -150,6 +154,25 @@ npx @_davideast/stitch-mcp init export STITCH_API_KEY="your-api-key" ``` +**Multiple API-key accounts:** Configure account IDs as references to environment variables. Secrets stay outside the pool configuration: + +```bash +export STITCH_ACCOUNTS='[{"id":"personal","env":"STITCH_API_KEY_PERSONAL"},{"id":"work","env":"STITCH_API_KEY_WORK"}]' +export STITCH_API_KEY_PERSONAL="your-personal-api-key" +export STITCH_API_KEY_WORK="your-work-api-key" +export STITCH_ACCOUNT_STRATEGY="least_used" +``` + +`STITCH_ACCOUNT_STRATEGY` accepts `least_used` (default) or `round_robin`. Project affinity and account health are stored in `~/.stitch-mcp/account-pool.json`. + +Start the local dashboard: + +```bash +stitch-mcp dashboard +``` + +Open the dashboard's **Quota bridge** section, copy the script, and run it from the signed-in Stitch Settings page. The bridge sends only quota numbers to localhost; browser cookies remain in Stitch. + **Manual (existing gcloud):** If you already have gcloud configured: ```bash @@ -178,6 +201,10 @@ Then use the proxy with `STITCH_USE_SYSTEM_GCLOUD=1`: | Variable | Description | |----------|-------------| +| `STITCH_ACCOUNTS` | JSON or `id:ENV_VAR` references for multiple API-key accounts | +| `STITCH_ACCOUNT_STRATEGY` | Account selection strategy: `least_used` or `round_robin` | +| `STITCH_POOL_STATE_FILE` | Optional pool state path override | +| `STITCH_ACCOUNT_COOLDOWN_MS` | Optional 429 cooldown duration in milliseconds | | `STITCH_API_KEY` | API key for direct authentication (skips OAuth) | | `STITCH_ACCESS_TOKEN` | Pre-existing access token | | `STITCH_USE_SYSTEM_GCLOUD` | Use system gcloud config instead of isolated config | diff --git a/package.json b/package.json index e136c0e..5216d56 100644 --- a/package.json +++ b/package.json @@ -24,9 +24,9 @@ ], "scripts": { "build": "bun run scripts/build.ts && tsc -p tsconfig.build.json && bun scripts/generate-bin.ts", + "typecheck": "tsc -p tsconfig.build.json --noEmit", "dev": "bun run src/cli.ts", "test": "bun test --preload ./tests/setup.ts", - "prepublishOnly": "bun run build", "verify-pack": "bun scripts/verify-pack.ts", "release": "np" }, diff --git a/src/commands/accounts/command.ts b/src/commands/accounts/command.ts new file mode 100644 index 0000000..a968874 --- /dev/null +++ b/src/commands/accounts/command.ts @@ -0,0 +1,17 @@ +import type { CommandDefinition } from '../../framework/CommandDefinition.js'; +import { AccountsCommandHandler } from './handler.js'; + +export const command: CommandDefinition> = { + name: 'accounts', + description: 'Manage Stitch API-key accounts', + arguments: ' [id]', + action: async (action, _options, command) => { + const id = command.args[1]; + try { + await new AccountsCommandHandler().execute(action, id); + } catch (error) { + console.error(error instanceof Error ? error.message : String(error)); + process.exitCode = 1; + } + }, +}; diff --git a/src/commands/accounts/handler.ts b/src/commands/accounts/handler.ts new file mode 100644 index 0000000..c7ea77f --- /dev/null +++ b/src/commands/accounts/handler.ts @@ -0,0 +1,63 @@ +import { type AccountPool, createAccountPoolFromEnvironment } from '../../services/stitch-pool/account-pool.js'; + +export class AccountsCommandHandler { + constructor(private readonly createPool = createAccountPoolFromEnvironment) {} + + async execute(action: string | undefined, accountId?: string): Promise { + const pool = await this.createPool(); + if (!pool) throw new Error('No Stitch API-key accounts configured. Set STITCH_ACCOUNTS or STITCH_API_KEY.'); + + switch (action) { + case 'list': + this.print(pool.status().accounts); + return; + case 'test': + await this.testAccounts(pool); + return; + case 'enable': { + const id = this.requireAccountId(accountId, pool, 'enable'); + await pool.enableAccount(id); + this.print(pool.status().accounts.find((account) => account.id === id)); + return; + } + case 'disable': { + const id = this.requireAccountId(accountId, pool, 'disable'); + await pool.disableAccount(id); + this.print(pool.status().accounts.find((account) => account.id === id)); + return; + } + default: + throw new Error('Usage: accounts list|test|enable |disable '); + } + } + + private async testAccounts(pool: AccountPool): Promise { + const results: Array> = []; + for (const account of pool.status().accounts) { + try { + await pool.testAccount(account.id); + results.push({ id: account.id, ok: true }); + } catch (error) { + results.push({ + id: account.id, + ok: false, + error: error instanceof Error ? error.message : String(error), + }); + } + } + await pool.flush(); + this.print(results); + } + + private requireAccountId(accountId: string | undefined, pool: AccountPool, action: string): string { + if (!accountId) throw new Error(`Usage: accounts ${action} `); + if (!pool.status().accounts.some((account) => account.id === accountId)) { + throw new Error(`Unknown Stitch account: ${accountId}`); + } + return accountId; + } + + private print(value: unknown): void { + console.log(JSON.stringify(value, null, 2)); + } +} diff --git a/src/commands/dashboard/command.ts b/src/commands/dashboard/command.ts new file mode 100644 index 0000000..6c9cc8f --- /dev/null +++ b/src/commands/dashboard/command.ts @@ -0,0 +1,23 @@ +import type { CommandDefinition } from '../../framework/CommandDefinition.js'; +import { DashboardHandler } from './handler.js'; + +interface DashboardOptions { + port: number; + host: string; +} + +export const command: CommandDefinition = { + name: 'dashboard', + description: 'Run the local Stitch account and quota dashboard', + options: [ + { flags: '--port ', description: 'Dashboard port', fn: (value) => Number.parseInt(value, 10), defaultValue: 4173 }, + { flags: '--host ', description: 'Bind host', defaultValue: '127.0.0.1' }, + ], + action: async (_args, options) => { + const result = await new DashboardHandler().execute(options); + if (!result.success) { + console.error(result.error); + process.exitCode = 1; + } + }, +}; diff --git a/src/commands/dashboard/handler.ts b/src/commands/dashboard/handler.ts new file mode 100644 index 0000000..3a3b95e --- /dev/null +++ b/src/commands/dashboard/handler.ts @@ -0,0 +1,163 @@ +import { randomUUID } from 'node:crypto'; +import { AccountPool, createAccountPoolFromEnvironment } from '../../services/stitch-pool/account-pool.js'; +import { renderDashboardPage } from './page.js'; + +export interface DashboardInput { + port: number; + host: string; +} + +export interface DashboardResult { + success: boolean; + url?: string; + error?: string; +} + +export class DashboardHandler { + constructor(private readonly createPool = createAccountPoolFromEnvironment) {} + + async execute(input: DashboardInput): Promise { + const pool = await this.createPool(); + if (!pool) return { success: false, error: 'No Stitch API-key accounts configured' }; + if (!Number.isInteger(input.port) || input.port < 1 || input.port > 65_535) { + return { success: false, error: 'Dashboard port must be between 1 and 65535' }; + } + + try { + const bridgeToken = randomUUID(); + const server = Bun.serve({ + hostname: input.host, + port: input.port, + fetch: async (request) => this.handleRequest(request, pool, bridgeToken), + }); + const url = `http://${input.host}:${server.port}`; + console.log(`Stitch dashboard running at ${url}`); + console.log('Quota bridge: open Stitch, then paste the generated bridge script from the dashboard.'); + return { success: true, url }; + } catch (error) { + return { success: false, error: error instanceof Error ? error.message : String(error) }; + } + } + + private async handleRequest(request: Request, pool: AccountPool, bridgeToken: string): Promise { + const url = new URL(request.url); + if (url.pathname === '/' && request.method === 'GET') { + return new Response(renderDashboardPage(pool.status().accounts, bridgeToken), { + headers: { 'Content-Type': 'text/html; charset=utf-8', 'Cache-Control': 'no-store' }, + }); + } + if (url.pathname === '/api/status' && request.method === 'GET') { + return this.json(pool.status()); + } + if (url.pathname === '/api/bridge' && request.method === 'GET') { + const accountId = url.searchParams.get('accountId'); + if (!accountId || !pool.status().accounts.some((account) => account.id === accountId)) { + return this.json({ error: 'Unknown account' }, 400); + } + return new Response(makeBridgeScript(new URL('/api/quota', request.url).toString(), bridgeToken, accountId), { + headers: { 'Content-Type': 'text/plain; charset=utf-8', 'Cache-Control': 'no-store' }, + }); + } + if (url.pathname === '/api/quota' && request.method === 'OPTIONS') { + return this.cors(new Response(null, { status: 204 })); + } + if (url.pathname === '/api/quota' && request.method === 'POST') { + return this.receiveQuota(request, pool, bridgeToken); + } + return new Response('Not found', { status: 404 }); + } + + private async receiveQuota(request: Request, pool: AccountPool, bridgeToken: string): Promise { + if (request.headers.get('Authorization') !== `Bearer ${bridgeToken}`) { + return this.cors(this.json({ error: 'Unauthorized' }, 401)); + } + const origin = request.headers.get('Origin'); + if (origin && origin !== 'https://stitch.withgoogle.com') { + return this.cors(this.json({ error: 'Origin not allowed' }, 403)); + } + const raw = await request.text(); + if (raw.length > 8_192) return this.cors(this.json({ error: 'Payload too large' }, 413)); + + try { + const body = JSON.parse(raw) as Record; + const accountId = typeof body.accountId === 'string' ? body.accountId : undefined; + if (!accountId || !pool.status().accounts.some((account) => account.id === accountId)) { + return this.cors(this.json({ error: 'Unknown account' }, 400)); + } + await pool.updateQuota(accountId, { + used: readQuotaNumber(body.used), + allocated: readQuotaNumber(body.allocated), + imageProUsed: readQuotaNumber(body.imageProUsed), + imageProAllocated: readQuotaNumber(body.imageProAllocated), + observedAt: Date.now(), + }); + return this.cors(this.json({ ok: true, quota: pool.status().quota })); + } catch { + return this.cors(this.json({ error: 'Invalid quota payload' }, 400)); + } + } + + private json(value: unknown, status = 200): Response { + return new Response(JSON.stringify(value), { + status, + headers: { 'Content-Type': 'application/json; charset=utf-8', 'Cache-Control': 'no-store' }, + }); + } + + private cors(response: Response): Response { + response.headers.set('Access-Control-Allow-Origin', 'https://stitch.withgoogle.com'); + response.headers.set('Access-Control-Allow-Headers', 'Authorization, Content-Type'); + response.headers.set('Access-Control-Allow-Methods', 'POST, OPTIONS'); + return response; + } +} + +function readQuotaNumber(value: unknown): number | undefined { + return typeof value === 'number' && Number.isFinite(value) && value >= 0 ? value : undefined; +} + +function makeBridgeScript(targetUrl: string, bridgeToken: string, accountId: string): string { + const encodedTarget = JSON.stringify(targetUrl); + const encodedToken = JSON.stringify(bridgeToken); + const encodedAccount = JSON.stringify(accountId); + return `(() => { + const target = ${encodedTarget}; + const bridgeToken = ${encodedToken}; + const accountId = ${encodedAccount}; + const send = () => { + const requestId = String(Date.now()); + const at = window.WIZ_global_data && window.WIZ_global_data.SNlM0e; + if (!at) { + console.error('Stitch session token unavailable; reload Settings and try again.'); + return; + } + const rpc = '[[["N5xENe","[]",null,"' + requestId + '"]]]'; + fetch('/_/Nemo/data/batchexecute?rpcids=N5xENe&source-path=%2Fsettings', { + method: 'POST', + credentials: 'include', + headers: { 'Content-Type': 'application/x-www-form-urlencoded;charset=UTF-8' }, + body: 'f.req=' + encodeURIComponent(rpc) + '&at=' + encodeURIComponent(at) + '&', + }).then((response) => response.text()).then((raw) => { + const start = raw.indexOf('[["wrb.fr"'); + const payload = start >= 0 ? raw.slice(start).split('\\n')[0] : ''; + if (!payload) throw new Error('Unexpected Stitch quota response.'); + const outer = JSON.parse(payload); + const fields = JSON.parse(outer[0][2]); + return fetch(target, { + method: 'POST', + headers: { 'Authorization': 'Bearer ' + bridgeToken, 'Content-Type': 'application/json' }, + body: JSON.stringify({ + accountId, + used: fields[8] == null ? 0 : fields[8], + allocated: fields[9], + imageProUsed: fields[10] == null ? 0 : fields[10], + imageProAllocated: fields[11], + }), + }); + }).then(() => console.log('Stitch quota sent to local dashboard.')).catch((error) => console.error('Stitch quota bridge failed:', error)); + }; + send(); + clearInterval(window.__stitchQuotaBridge); + window.__stitchQuotaBridge = setInterval(send, 60000); +})();`; +} diff --git a/src/commands/dashboard/page.ts b/src/commands/dashboard/page.ts new file mode 100644 index 0000000..37e6c78 --- /dev/null +++ b/src/commands/dashboard/page.ts @@ -0,0 +1,46 @@ +import type { AccountStatus } from '../../services/stitch-pool/types.js'; + +export function renderDashboardPage(accounts: AccountStatus[], bridgeToken: string): string { + const safeAccounts = accounts.map((account) => ({ + id: account.id, + envVar: account.envVar, + legacy: account.legacy, + enabled: account.enabled, + })); + const accountJson = JSON.stringify(safeAccounts).replace(/ + + + + +Stitch Pool Control + + + +
+
Stitch account pool

Control surface

Connecting…
+
+

See the quota signal before the pool hits the wall.

Local operational view for sticky project routing, account health, and the quota snapshot supplied by your signed-in Stitch browser session.

+

Accounts

Loading pool…
AccountStateRequestsLast usedFailures
Loading…
+

Project affinity

sticky routing
No affinity yet.

Quota bridge

browser session only

Run the generated script from the Stitch Settings page. It reads GetUserSettings with your existing browser session and sends only numeric quota fields to this local dashboard.

Cookie values never leave the browser.
After copying: open stitch.withgoogle.com/settings, open DevTools Console, paste, and run.
+
+
+ + +`; +} diff --git a/src/commands/doctor/steps/ApiKeyDetectedStep.ts b/src/commands/doctor/steps/ApiKeyDetectedStep.ts index 0cdc1cb..2232bd0 100644 --- a/src/commands/doctor/steps/ApiKeyDetectedStep.ts +++ b/src/commands/doctor/steps/ApiKeyDetectedStep.ts @@ -23,10 +23,7 @@ export class ApiKeyDetectedStep implements CommandStep { return { success: false, error: new Error(message) }; } - const masked = apiKey.length > 7 - ? `${apiKey.slice(0, 4)}...${apiKey.slice(-3)}` - : '***'; - const message = `Detected (${masked})`; + const message = 'Detected'; context.checks.push({ name: 'API Key', passed: true, diff --git a/src/commands/pool/command.ts b/src/commands/pool/command.ts new file mode 100644 index 0000000..f07a003 --- /dev/null +++ b/src/commands/pool/command.ts @@ -0,0 +1,16 @@ +import type { CommandDefinition } from '../../framework/CommandDefinition.js'; +import { PoolCommandHandler } from './handler.js'; + +export const command: CommandDefinition> = { + name: 'pool', + description: 'Inspect or reset the Stitch account pool', + arguments: '', + action: async (action) => { + try { + await new PoolCommandHandler().execute(action); + } catch (error) { + console.error(error instanceof Error ? error.message : String(error)); + process.exitCode = 1; + } + }, +}; diff --git a/src/commands/pool/handler.ts b/src/commands/pool/handler.ts new file mode 100644 index 0000000..97f66c0 --- /dev/null +++ b/src/commands/pool/handler.ts @@ -0,0 +1,21 @@ +import { createAccountPoolFromEnvironment } from '../../services/stitch-pool/account-pool.js'; + +export class PoolCommandHandler { + constructor(private readonly createPool = createAccountPoolFromEnvironment) {} + + async execute(action: string | undefined): Promise { + const pool = await this.createPool(); + if (!pool) throw new Error('No Stitch API-key accounts configured. Set STITCH_ACCOUNTS or STITCH_API_KEY.'); + + if (action === 'status') { + console.log(JSON.stringify(pool.status(), null, 2)); + return; + } + if (action === 'reset') { + await pool.reset(); + console.log(JSON.stringify(pool.status(), null, 2)); + return; + } + throw new Error('Usage: pool status|reset'); + } +} diff --git a/src/commands/projects/command.ts b/src/commands/projects/command.ts new file mode 100644 index 0000000..013f8c8 --- /dev/null +++ b/src/commands/projects/command.ts @@ -0,0 +1,29 @@ +import { StitchToolClient } from '@google/stitch-sdk'; +import type { CommandDefinition } from '../../framework/CommandDefinition.js'; +import { createAccountPoolFromEnvironment } from '../../services/stitch-pool/account-pool.js'; + +export const command: CommandDefinition> = { + name: 'projects', + description: 'List Stitch projects across configured accounts', + action: async () => { + try { + const pool = await createAccountPoolFromEnvironment(); + if (pool) { + const result = await pool.call('list_projects'); + await pool.flush(); + console.log(JSON.stringify(result, null, 2)); + return; + } + + const client = new StitchToolClient(); + try { + console.log(JSON.stringify(await client.callTool('list_projects', {}), null, 2)); + } finally { + await client.close(); + } + } catch (error) { + console.error(error instanceof Error ? error.message : String(error)); + process.exitCode = 1; + } + }, +}; diff --git a/src/commands/proxy/LoggingCallToolHandler.ts b/src/commands/proxy/LoggingCallToolHandler.ts index e867be7..837accd 100644 --- a/src/commands/proxy/LoggingCallToolHandler.ts +++ b/src/commands/proxy/LoggingCallToolHandler.ts @@ -1,5 +1,7 @@ import { CallToolRequestSchema } from '@modelcontextprotocol/sdk/types.js'; import type { CaptureSpec } from '../../lib/log/capture/spec.js'; +import type { AccountPool } from '../../services/stitch-pool/account-pool.js'; +import { redactSecrets } from '../../services/stitch-pool/transport.js'; const DEFAULT_STITCH_MCP_URL = 'https://stitch.googleapis.com/mcp'; @@ -8,58 +10,96 @@ interface ForwardOptions { url: string; } +interface CallToolForwarder { + call(name: string, args: Record): Promise; + secrets?: string[]; +} + +interface CallToolRequest { + params: { + name: string; + arguments?: Record; + }; +} + +interface RequestServer { + setRequestHandler(schema: unknown, handler: (request: CallToolRequest) => Promise): void; + _requestHandlers?: Map Promise>; +} + /** * Replace the SDK proxy's tools/call handler with a capture-wrapped variant. * Must be called AFTER {@link StitchProxy.start} (which registers the original). */ -export function installLoggingCallToolHandler(proxy: any, capture: CaptureSpec, opts?: Partial): void { - const apiKey = opts?.apiKey ?? process.env.STITCH_API_KEY; - const url = opts?.url ?? process.env.STITCH_MCP_URL ?? DEFAULT_STITCH_MCP_URL; - if (!apiKey) throw new Error('logging proxy requires STITCH_API_KEY'); - - // The SDK exposes McpServer at proxy.server, whose underlying Server is at .server. - const server = proxy?.server?.server; - if (!server || typeof server.setRequestHandler !== 'function') { - throw new Error('cannot install logging handler: proxy.server.server.setRequestHandler missing'); - } +export function installCallToolHandler(proxy: unknown, forwarder: CallToolForwarder, capture?: CaptureSpec): void { + const server = getRequestServer(proxy); + if (!server) throw new Error('cannot install call handler: proxy.server.server.setRequestHandler missing'); - server.setRequestHandler(CallToolRequestSchema, async (request: any) => { + const originalHandler = server._requestHandlers?.get('tools/call'); + server.setRequestHandler(CallToolRequestSchema, async (request) => { const { name, arguments: args } = request.params; + const requestArgs = args ?? {}; + const secrets = forwarder.secrets ?? []; + const safeArgs = redactSecrets(requestArgs, secrets) as Record; const startedAt = new Date().toISOString(); const t0 = Date.now(); - let result: any; + let result: unknown; try { - result = await forwardToStitch({ apiKey, url }, name, args ?? {}); + // Keep SDK-owned virtual tools local; their implementation uses the SDK client. + if (name === 'download_assets' && originalHandler) { + result = await originalHandler(request); + } else { + result = await forwarder.call(name, requestArgs); + } } catch (err) { - const message = err instanceof Error ? err.message : String(err); + const message = redactSecrets(err instanceof Error ? err.message : String(err), secrets) as string; result = { content: [{ type: 'text', text: `Error calling ${name}: ${message}` }], isError: true }; } - const finishedAt = new Date().toISOString(); - const durationMs = Date.now() - t0; - // best-effort capture — never propagate failures to the MCP client - try { - await capture.capture({ - tool: name, - args: (args ?? {}) as Record, - result, - duration_ms: durationMs, - started_at: startedAt, - finished_at: finishedAt, - }); - } catch (err) { - const msg = err instanceof Error ? err.message : String(err); - // eslint-disable-next-line no-console - console.error(`[stitch-mcp log] capture failed: ${msg}`); + if (capture) { + try { + await capture.capture({ + tool: name, + args: safeArgs, + result: redactSecrets(result, secrets), + duration_ms: Date.now() - t0, + started_at: startedAt, + finished_at: new Date().toISOString(), + }); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + // eslint-disable-next-line no-console + console.error(`[stitch-mcp log] capture failed: ${message}`); + } } - return result; }); } +export function installLoggingCallToolHandler( + proxy: unknown, + capture: CaptureSpec, + opts?: Partial & { pool?: AccountPool }, +): void { + const apiKey = opts?.apiKey ?? process.env.STITCH_API_KEY; + const url = opts?.url ?? process.env.STITCH_MCP_URL ?? DEFAULT_STITCH_MCP_URL; + const pool = opts?.pool; + if (pool) { + installCallToolHandler(proxy, { + call: (name, args) => pool.call(name, args), + }, capture); + return; + } + if (!apiKey) throw new Error('logging proxy requires STITCH_API_KEY'); + installCallToolHandler(proxy, { + call: (name, args) => forwardToStitch({ apiKey, url }, name, args), + secrets: [apiKey], + }, capture); +} + /** Minimal JSON-RPC forwarder, mirroring the SDK's private forwardToStitch. */ -async function forwardToStitch(opts: ForwardOptions, name: string, args: Record): Promise { +async function forwardToStitch(opts: ForwardOptions, name: string, args: Record): Promise { const body = { jsonrpc: '2.0', method: 'tools/call', @@ -76,7 +116,16 @@ async function forwardToStitch(opts: ForwardOptions, name: string, args: Record< body: JSON.stringify(body), }); if (!res.ok) throw new Error(`Stitch API error (${res.status}): ${await res.text()}`); - const json = (await res.json()) as { result?: unknown; error?: { message: string } }; - if (json.error) throw new Error(`Stitch RPC error: ${json.error.message}`); - return json.result; + const json = await res.json() as unknown; + if (!json || typeof json !== 'object') return undefined; + const envelope = json as { result?: unknown; error?: { message?: string } }; + if (envelope.error) throw new Error(`Stitch RPC error: ${envelope.error.message ?? 'unknown error'}`); + return envelope.result; +} + +function getRequestServer(proxy: unknown): RequestServer | undefined { + if (!proxy || typeof proxy !== 'object') return undefined; + const candidate = proxy as { server?: { server?: RequestServer } }; + const server = candidate.server?.server; + return server && typeof server.setRequestHandler === 'function' ? server : undefined; } diff --git a/src/commands/proxy/handler.ts b/src/commands/proxy/handler.ts index 7db848b..8b98199 100644 --- a/src/commands/proxy/handler.ts +++ b/src/commands/proxy/handler.ts @@ -2,7 +2,8 @@ import { StitchProxy } from '@google/stitch-sdk'; import type { StitchProxy as StitchProxyType } from '@google/stitch-sdk'; import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'; import { createCaptureHandler, isLogEnabled } from '../../lib/log/factory.js'; -import { installLoggingCallToolHandler } from './LoggingCallToolHandler.js'; +import { type AccountPool, createAccountPoolFromEnvironment } from '../../services/stitch-pool/account-pool.js'; +import { installCallToolHandler, installLoggingCallToolHandler } from './LoggingCallToolHandler.js'; interface ProxyCommandInput { port?: number; @@ -15,42 +16,100 @@ interface ProxyCommandResult { error?: { code: string; message: string; recoverable: boolean }; } +interface ProxyServerLike { + server?: { + server?: { + setRequestHandler?: unknown; + }; + }; +} + export class ProxyCommandHandler { private createProxy: (opts: { apiKey?: string }) => StitchProxyType; private createTransport: () => StdioServerTransport; + private createPool: () => Promise; constructor(deps?: { createProxy?: (opts: { apiKey?: string }) => StitchProxyType; createTransport?: () => StdioServerTransport; + createPool?: () => Promise; }) { this.createProxy = deps?.createProxy ?? ((opts) => new StitchProxy(opts)); this.createTransport = deps?.createTransport ?? (() => new StdioServerTransport()); + this.createPool = deps?.createPool ?? (() => createAccountPoolFromEnvironment()); } - async execute(input: ProxyCommandInput): Promise { + async execute(_input: ProxyCommandInput): Promise { try { - const proxy = this.createProxy({ - apiKey: process.env.STITCH_API_KEY, - }); - const transport = this.createTransport(); - await proxy.start(transport); - - // Override the SDK's tools/call handler with a capture-wrapped variant. - // Must run AFTER proxy.start() (which registers the original handler). - if (isLogEnabled()) { + const pool = await this.createPool(); + const attemptedAccounts = new Set(); + let startupCredential = pool?.primaryCredential; + let serverRetries = 0; + let activeProxy: StitchProxyType | undefined; + let activeTransport: StdioServerTransport | undefined; + + while (!activeProxy || !activeTransport) { + const apiKey = pool ? startupCredential?.apiKey : process.env.STITCH_API_KEY; + if (pool && !startupCredential) throw new Error('No healthy Stitch accounts available'); + + const proxy = this.createProxy({ apiKey }); + const transport = this.createTransport(); + if (startupCredential) { + if (!pool) throw new Error('Missing account pool for credential'); + attemptedAccounts.add(startupCredential.id); + pool.recordStartupAttempt(startupCredential.id); + } + try { - installLoggingCallToolHandler(proxy, createCaptureHandler()); - } catch (err) { - const msg = err instanceof Error ? err.message : String(err); - // eslint-disable-next-line no-console - console.error(`[stitch-mcp log] failed to install capture handler: ${msg}`); + await proxy.start(transport); + if (startupCredential) pool!.recordStartupSuccess(startupCredential.id); + activeProxy = proxy; + activeTransport = transport; + } catch (error) { + if (!startupCredential || !pool) throw error; + const classification = pool.recordStartupFailure(startupCredential.id, error); + await proxy.close().catch(() => undefined); + if (classification.kind === 'server' && serverRetries < 2) { + await wait(250 * 2 ** serverRetries); + serverRetries += 1; + continue; + } + if (!isStartupFailover(classification.kind)) throw error; + startupCredential = pool.startupCredential(attemptedAccounts); + serverRetries = 0; } } - await transport.onclose; + if (!activeProxy || !activeTransport) throw new Error('Proxy failed to start'); + if (hasCallHandler(activeProxy)) { + const capture = isLogEnabled() ? createCaptureHandler() : undefined; + if (pool) { + installCallToolHandler(activeProxy, pool, capture); + } else if (capture) { + installLoggingCallToolHandler(activeProxy, capture); + } + } + + await activeTransport.onclose; + await pool?.flush(); return { success: true, data: { status: 'running' } }; - } catch (e: any) { - return { success: false, error: { code: 'PROXY_START_ERROR', message: e.message, recoverable: false } }; + } catch (e: unknown) { + const message = e instanceof Error ? e.message : String(e); + return { success: false, error: { code: 'PROXY_START_ERROR', message, recoverable: false } }; } } } + +function isStartupFailover(kind: string): boolean { + return kind === 'unauthorized' || kind === 'rate_limited' || kind === 'server'; +} + +function wait(milliseconds: number): Promise { + return new Promise((resolve) => setTimeout(resolve, milliseconds)); +} + +function hasCallHandler(proxy: unknown): boolean { + if (!proxy || typeof proxy !== 'object') return false; + const server = (proxy as ProxyServerLike).server?.server; + return typeof server?.setRequestHandler === 'function'; +} diff --git a/src/lib/log/append.test.ts b/src/lib/log/append.test.ts index b793056..70b29c5 100644 --- a/src/lib/log/append.test.ts +++ b/src/lib/log/append.test.ts @@ -1,5 +1,4 @@ -import { describe, test, expect, beforeEach, afterEach } from 'bun:test'; -import { mkdtemp, readFile, rm } from 'node:fs/promises'; +import { mkdtemp, readFile, rm, writeFile } from 'node:fs/promises'; import { tmpdir } from 'node:os'; import { join } from 'node:path'; import { appendEvent } from './append.js'; @@ -61,7 +60,9 @@ describe('appendEvent', () => { }); test('returns EVENT_WRITE_FAILED when destination cannot be written', async () => { - const r = await appendEvent('/proc/cannot-write-here/x.jsonl', validEvent); + const blockedParent = join(root, 'blocked'); + await writeFile(blockedParent, 'not-a-directory'); + const r = await appendEvent(join(blockedParent, 'events.jsonl'), validEvent); expect(r.success).toBe(false); if (r.success) return; expect(r.error.code).toBe('EVENT_WRITE_FAILED'); diff --git a/src/services/stitch-pool/account-health-tracker.ts b/src/services/stitch-pool/account-health-tracker.ts new file mode 100644 index 0000000..fb346c0 --- /dev/null +++ b/src/services/stitch-pool/account-health-tracker.ts @@ -0,0 +1,144 @@ +import type { ErrorKind } from './errors.js'; +import type { AccountStatus, PersistedAccountState } from './types.js'; + +interface MutableAccountState extends PersistedAccountState { + enabled: boolean; + usageCount: number; + consecutiveFailures: number; +} + +export interface AccountHealthTrackerOptions { + now?: () => number; + cooldownMs?: number; + onChange?: (accountId: string, state: PersistedAccountState) => void; +} + +export class AccountHealthTracker { + private readonly states = new Map(); + private readonly now: () => number; + private readonly cooldownMs: number; + private readonly onChange?: (accountId: string, state: PersistedAccountState) => void; + + constructor(accountIds: string[], persisted: Record = {}, options: AccountHealthTrackerOptions = {}) { + this.now = options.now ?? Date.now; + this.cooldownMs = options.cooldownMs ?? readCooldownFromEnvironment(); + this.onChange = options.onChange; + for (const id of accountIds) { + const saved = persisted[id] ?? {}; + this.states.set(id, { + enabled: saved.enabled !== false, + disabledReason: saved.disabledReason, + cooldownUntil: saved.cooldownUntil, + usageCount: saved.usageCount ?? 0, + lastUsedAt: saved.lastUsedAt, + consecutiveFailures: saved.consecutiveFailures ?? 0, + }); + } + } + + isAvailable(accountId: string, now = this.now()): boolean { + const state = this.states.get(accountId); + if (!state?.enabled) return false; + return state.cooldownUntil === undefined || state.cooldownUntil <= now; + } + + recordRequest(accountId: string): void { + const state = this.require(accountId); + state.usageCount += 1; + state.lastUsedAt = this.now(); + this.persist(accountId); + } + + recordSuccess(accountId: string): void { + const state = this.require(accountId); + state.consecutiveFailures = 0; + if (state.cooldownUntil !== undefined && state.cooldownUntil <= this.now()) state.cooldownUntil = undefined; + this.persist(accountId); + } + + recordFailure(accountId: string, kind: ErrorKind, retryAfterMs?: number): void { + const state = this.require(accountId); + state.consecutiveFailures += 1; + if (kind === 'unauthorized') { + state.enabled = false; + state.disabledReason = 'unauthorized'; + state.cooldownUntil = undefined; + } else if (kind === 'rate_limited') { + state.cooldownUntil = this.now() + Math.max(this.cooldownMs, retryAfterMs ?? 0); + } + this.persist(accountId); + } + + enable(accountId: string): void { + const state = this.require(accountId); + state.enabled = true; + state.disabledReason = undefined; + state.cooldownUntil = undefined; + state.consecutiveFailures = 0; + this.persist(accountId); + } + + disable(accountId: string, reason = 'manual'): void { + const state = this.require(accountId); + state.enabled = false; + state.disabledReason = reason; + state.cooldownUntil = undefined; + this.persist(accountId); + } + + reset(): void { + for (const [id, state] of this.states) { + state.enabled = true; + state.disabledReason = undefined; + state.cooldownUntil = undefined; + state.usageCount = 0; + state.lastUsedAt = undefined; + state.consecutiveFailures = 0; + this.persist(id); + } + } + + status(accountId: string, envVar: string, legacy: boolean, configured: boolean): AccountStatus { + const state = this.require(accountId); + const status: AccountStatus = { + id: accountId, + envVar, + legacy, + configured, + enabled: state.enabled, + usageCount: state.usageCount, + consecutiveFailures: state.consecutiveFailures, + }; + if (state.cooldownUntil !== undefined) status.cooldownUntil = state.cooldownUntil; + if (state.lastUsedAt !== undefined) status.lastUsedAt = state.lastUsedAt; + if (state.disabledReason !== undefined) status.disabledReason = state.disabledReason; + return status; + } + + snapshot(): Record { + return Object.fromEntries([...this.states].map(([id, state]) => [id, { ...state }])); + } + + private require(accountId: string): MutableAccountState { + const state = this.states.get(accountId); + if (!state) throw new Error(`Unknown Stitch account: ${accountId}`); + return state; + } + + private persist(accountId: string): void { + const state = this.require(accountId); + this.onChange?.(accountId, { + enabled: state.enabled, + disabledReason: state.disabledReason, + cooldownUntil: state.cooldownUntil, + usageCount: state.usageCount, + lastUsedAt: state.lastUsedAt, + consecutiveFailures: state.consecutiveFailures, + }); + } +} + +function readCooldownFromEnvironment(): number { + const configured = Number(process.env.STITCH_ACCOUNT_COOLDOWN_MS); + return Number.isFinite(configured) && configured >= 0 ? configured : 30_000; +} diff --git a/src/services/stitch-pool/account-pool.ts b/src/services/stitch-pool/account-pool.ts new file mode 100644 index 0000000..fedade2 --- /dev/null +++ b/src/services/stitch-pool/account-pool.ts @@ -0,0 +1,493 @@ +import { AccountHealthTracker } from './account-health-tracker.js'; +import { AccountSelector, type SelectableAccount } from './account-selector.js'; +import { classifyError, type ErrorClassification } from './errors.js'; +import { CredentialProvider } from './credential-provider.js'; +import { ProjectAffinityStore } from './project-affinity-store.js'; +import { fetchStitchTool, isSafeToMove } from './transport.js'; +import type { + AccountReference, + AccountStatus, + AccountStrategy, + QuotaSnapshot, + ResolvedAccount, + Sleep, + StitchRequestExecutor, +} from './types.js'; +export interface AccountPoolOptions { + provider?: CredentialProvider; + store?: ProjectAffinityStore; + strategy?: AccountStrategy; + fetchImpl?: typeof fetch; + url?: string; + sleep?: Sleep; + retryBaseMs?: number; + maxServerRetries?: number; + cooldownMs?: number; + now?: () => number; +} + +export interface AccountPoolStatus { + configured: boolean; + legacy: boolean; + strategy: AccountStrategy; + stateFile: string; + affinity: Record; + accounts: AccountStatus[]; + quota?: QuotaSnapshot; +} +interface SuccessfulProjectResult { + account: ResolvedAccount; + result: unknown; +} + +export class AccountPoolError extends Error { + readonly code: string; + readonly classification?: ErrorClassification; + + constructor(code: string, message: string, classification?: ErrorClassification) { + super(message); + this.name = 'AccountPoolError'; + this.code = code; + this.classification = classification; + } +} + +export class AccountPool { + private readonly provider: CredentialProvider; + private readonly store: ProjectAffinityStore; + private readonly health: AccountHealthTracker; + private readonly selector: AccountSelector; + private readonly fetchImpl?: typeof fetch; + private readonly url?: string; + private readonly sleep: Sleep; + private readonly retryBaseMs: number; + private readonly maxServerRetries: number; + private readonly projectLocks = new Map>(); + private readonly referencesById: Record; + + private constructor(options: { + provider: CredentialProvider; + store: ProjectAffinityStore; + health: AccountHealthTracker; + selector: AccountSelector; + fetchImpl?: typeof fetch; + url?: string; + sleep: Sleep; + retryBaseMs: number; + maxServerRetries: number; + }) { + this.provider = options.provider; + this.store = options.store; + this.health = options.health; + this.selector = options.selector; + this.fetchImpl = options.fetchImpl; + this.url = options.url; + this.sleep = options.sleep; + this.retryBaseMs = options.retryBaseMs; + this.maxServerRetries = options.maxServerRetries; + this.referencesById = Object.fromEntries(this.provider.references.map((ref) => [ref.id, ref])); + } + + static async create(options: AccountPoolOptions = {}): Promise { + const provider = options.provider ?? new CredentialProvider(); + const store = options.store ?? new ProjectAffinityStore(); + await store.load(); + const strategy = options.strategy ?? readConfiguredStrategy() ?? store.strategy ?? 'least_used'; + store.setStrategy(strategy); + + const health = new AccountHealthTracker( + provider.references.map((ref) => ref.id), + store.snapshot().accounts, + { + now: options.now, + cooldownMs: options.cooldownMs, + onChange: (accountId, state) => store.setAccountState(accountId, state), + }, + ); + const selector = new AccountSelector(strategy); + return new AccountPool({ + provider, + store, + health, + selector, + fetchImpl: options.fetchImpl, + url: options.url, + sleep: options.sleep ?? defaultSleep, + retryBaseMs: options.retryBaseMs ?? 250, + maxServerRetries: options.maxServerRetries ?? 2, + }); + } + + get hasConfiguration(): boolean { + return this.provider.hasConfiguration; + } + + get isLegacy(): boolean { + return this.provider.isLegacy; + } + + get primaryCredential(): ResolvedAccount | undefined { + return this.selectAvailable(new Set()); + } + + get primaryApiKey(): string | undefined { + return this.primaryCredential?.apiKey; + } + + get accountIds(): string[] { + return this.provider.references.map((ref) => ref.id); + } + + async call( + toolName: string, + args: Record = {}, + executor: StitchRequestExecutor = this.defaultExecutor, + ): Promise { + if (toolName === 'list_projects') return this.listProjects(executor); + + const projectId = extractProjectId(args); + if (projectId) { + return this.withProjectLock(projectId, () => this.callForProject(toolName, args, projectId, executor)); + } + return this.callWithoutProject(toolName, args, executor); + } + + async testAccount(accountId: string, executor: StitchRequestExecutor = this.defaultExecutor): Promise { + const account = this.provider.resolve(accountId); + if (!account) throw new AccountPoolError('ACCOUNT_NOT_CONFIGURED', `Stitch account is not configured: ${accountId}`); + return this.executeOnAccount(account, 'list_projects', {}, executor); + } + + recordStartupAttempt(accountId: string): void { + this.health.recordRequest(accountId); + } + + recordStartupSuccess(accountId: string): void { + this.health.recordSuccess(accountId); + } + + recordStartupFailure(accountId: string, error: unknown): ErrorClassification { + const classification = classifyError(error); + this.health.recordFailure(accountId, classification.kind, classification.retryAfterMs); + return classification; + } + + startupCredential(excluded: Set): ResolvedAccount | undefined { + return this.selectAvailable(excluded); + } + + async updateQuota(accountId: string, snapshot: Omit): Promise { + this.ensureReference(accountId); + this.store.setQuota({ ...snapshot, accountId, source: 'stitch-web' }); + await this.store.flush(); + } + + async enableAccount(accountId: string): Promise { + this.ensureReference(accountId); + this.health.enable(accountId); + await this.store.flush(); + } + + async disableAccount(accountId: string): Promise { + this.ensureReference(accountId); + this.health.disable(accountId); + await this.store.flush(); + } + + async reset(): Promise { + this.health.reset(); + await this.store.reset(); + this.selector.setStrategy(readConfiguredStrategy() ?? 'least_used'); + this.store.setStrategy(this.selector.getStrategy()); + await this.store.flush(); + } + + async flush(): Promise { + await this.store.flush(); + } + + status(): AccountPoolStatus { + const accounts = this.provider.references.map((ref) => { + const configured = this.provider.resolve(ref.id) !== undefined; + return this.health.status(ref.id, ref.envVar, ref.legacy === true, configured); + }); + const status: AccountPoolStatus = { + configured: this.hasConfiguration, + legacy: this.isLegacy, + strategy: this.selector.getStrategy(), + stateFile: this.store.filePath, + affinity: this.store.affinities(), + accounts, + }; + const quota = this.store.quota(); + if (quota) status.quota = quota; + return status; + } + + private async callForProject( + toolName: string, + args: Record, + projectId: string, + executor: StitchRequestExecutor, + ): Promise { + const affinityAccountId = this.store.get(projectId); + if (affinityAccountId) { + const account = this.provider.resolve(affinityAccountId); + if (!account || !this.health.isAvailable(affinityAccountId)) { + throw new AccountPoolError('AFFINITY_ACCOUNT_UNAVAILABLE', `The account owning project ${projectId} is unavailable`); + } + try { + return await this.executeOnAccount(account, toolName, args, executor); + } catch (error) { + throw this.toPoolError(error); + } + } + + return this.callWithFailover(toolName, args, projectId, isSafeToMove(toolName), executor); + } + + private async callWithoutProject( + toolName: string, + args: Record, + executor: StitchRequestExecutor, + ): Promise { + return this.callWithFailover(toolName, args, undefined, isSafeToMove(toolName), executor); + } + + private async callWithFailover( + toolName: string, + args: Record, + projectId: string | undefined, + allowFailover: boolean, + executor: StitchRequestExecutor, + ): Promise { + const excluded = new Set(); + let lastError: unknown; + + while (true) { + const account = this.selectAvailable(excluded); + if (!account) { + if (lastError) throw this.toPoolError(lastError); + throw new AccountPoolError('NO_HEALTHY_ACCOUNTS', 'No healthy Stitch accounts available'); + } + excluded.add(account.id); + + try { + const result = await this.executeOnAccount(account, toolName, args, executor); + if (projectId) this.store.set(projectId, account.id); + return result; + } catch (error) { + lastError = error; + const classification = classifyError(error); + if (!allowFailover || !canFailover(classification)) throw this.toPoolError(error, classification); + } + } + } + + private async listProjects(executor: StitchRequestExecutor): Promise { + const accounts = this.availableAccounts(); + if (accounts.length === 0) throw new AccountPoolError('NO_HEALTHY_ACCOUNTS', 'No healthy Stitch accounts available'); + + const settled = await Promise.allSettled( + accounts.map(async (account) => ({ account, result: await this.executeOnAccount(account, 'list_projects', {}, executor) })), + ); + const successful: SuccessfulProjectResult[] = []; + let lastError: unknown; + for (const outcome of settled) { + if (outcome.status === 'fulfilled') successful.push(outcome.value); + else lastError = outcome.reason; + } + if (successful.length === 0) throw this.toPoolError(lastError ?? new Error('All Stitch accounts failed')); + + for (const entry of successful) { + for (const project of extractProjects(entry.result)) { + const projectId = extractProjectId(project, true); + if (projectId && !this.store.has(projectId)) this.store.set(projectId, entry.account.id); + } + } + return aggregateProjectResults(successful.map((entry) => entry.result)); + } + + private async executeOnAccount( + account: ResolvedAccount, + toolName: string, + args: Record, + executor: StitchRequestExecutor, + ): Promise { + let serverAttempt = 0; + while (true) { + this.health.recordRequest(account.id); + try { + const result = await executor(account, toolName, args); + this.health.recordSuccess(account.id); + return result; + } catch (error) { + const classification = classifyError(error); + if (classification.kind === 'server' && serverAttempt < this.maxServerRetries) { + const exponentialDelay = this.retryBaseMs * 2 ** serverAttempt; + const delay = Math.max(exponentialDelay, classification.retryAfterMs ?? 0); + serverAttempt += 1; + await this.sleep(delay); + continue; + } + this.health.recordFailure(account.id, classification.kind, classification.retryAfterMs); + throw error; + } + } + } + + private selectAvailable(excluded: Set): ResolvedAccount | undefined { + const accounts = this.availableAccounts(excluded); + if (accounts.length === 0) return undefined; + const candidates: SelectableAccount[] = accounts.map((account) => { + const status = this.statusFor(account.id); + return { id: account.id, usageCount: status.usageCount, lastUsedAt: status.lastUsedAt }; + }); + const selected = this.selector.select(candidates); + return accounts.find((account) => account.id === selected.id); + } + + private availableAccounts(excluded: Set = new Set()): ResolvedAccount[] { + return this.provider.resolveAll().filter((account) => !excluded.has(account.id) && this.health.isAvailable(account.id)); + } + + private statusFor(accountId: string): AccountStatus { + const reference = this.referencesById[accountId]; + if (!reference) throw new AccountPoolError('UNKNOWN_ACCOUNT', `Unknown Stitch account: ${accountId}`); + return this.health.status(accountId, reference.envVar, reference.legacy === true, this.provider.resolve(accountId) !== undefined); + } + + private ensureReference(accountId: string): void { + if (!this.referencesById[accountId]) throw new AccountPoolError('UNKNOWN_ACCOUNT', `Unknown Stitch account: ${accountId}`); + } + + private async withProjectLock(projectId: string, operation: () => Promise): Promise { + const previous = this.projectLocks.get(projectId) ?? Promise.resolve(); + const current = previous.catch(() => undefined).then(operation); + this.projectLocks.set(projectId, current); + try { + return await current; + } finally { + if (this.projectLocks.get(projectId) === current) this.projectLocks.delete(projectId); + } + } + + private readonly defaultExecutor: StitchRequestExecutor = async (account, toolName, args) => fetchStitchTool({ + account, + toolName, + args, + url: this.url, + fetchImpl: this.fetchImpl, + }); + + private toPoolError(error: unknown, classification = classifyError(error)): AccountPoolError { + const message = classification.kind === 'other' ? 'Stitch request failed' : `Stitch request failed (${classification.kind})`; + return new AccountPoolError(classification.kind.toUpperCase(), message, classification); + } +} + +export async function createAccountPoolFromEnvironment(options: AccountPoolOptions = {}): Promise { + const provider = options.provider ?? new CredentialProvider(); + if (!provider.hasConfiguration) return undefined; + return AccountPool.create({ ...options, provider }); +} + +export function extractProjectId(value: unknown, includeGenericId = false): string | undefined { + if (!value || typeof value !== 'object') return undefined; + const record = value as Record; + const keys = includeGenericId ? ['projectId', 'project_id', 'project', 'id'] : ['projectId', 'project_id', 'project']; + for (const key of keys) { + const candidate = record[key]; + if (typeof candidate === 'string' && candidate.length > 0) return stripProjectPrefix(candidate); + } + const name = record.name; + if (typeof name === 'string' && name.startsWith('projects/')) return stripProjectPrefix(name); + return undefined; +} + +export function aggregateProjectResults(results: unknown[]): unknown { + const projects = deduplicateProjects(results.flatMap((result) => extractProjects(result))); + const first = results[0]; + const location = first === undefined ? undefined : findProjectsPath(first); + if (!location) return { projects }; + const clone = cloneJson(first); + setPath(clone, location, projects); + return clone; +} + +export function extractProjects(value: unknown): Record[] { + if (Array.isArray(value)) { + return value.filter((project): project is Record => !!project && typeof project === 'object'); + } + if (!value || typeof value !== 'object') return []; + const record = value as Record; + if (Array.isArray(record.projects)) { + return record.projects.filter((project): project is Record => !!project && typeof project === 'object'); + } + for (const child of Object.values(record)) { + const nested = extractProjects(child); + if (nested.length > 0) return nested; + } + return []; +} + +function findProjectsPath(value: unknown, path: string[] = []): string[] | undefined { + if (!value || typeof value !== 'object') return undefined; + const record = value as Record; + if (Array.isArray(record.projects)) return [...path, 'projects']; + for (const [key, child] of Object.entries(record)) { + const nested = findProjectsPath(child, [...path, key]); + if (nested) return nested; + } + return undefined; +} + +function deduplicateProjects(projects: Record[]): Record[] { + const seen = new Set(); + return projects.filter((project) => { + const key = projectKey(project); + if (seen.has(key)) return false; + seen.add(key); + return true; + }); +} + +function projectKey(project: Record): string { + return extractProjectId(project, true) ?? String(project.id ?? project.name ?? JSON.stringify(project)); +} + +function setPath(value: unknown, path: string[], replacement: unknown): void { + if (!value || typeof value !== 'object' || path.length === 0) return; + let cursor = value as Record; + for (const segment of path.slice(0, -1)) { + const next = cursor[segment]; + if (!next || typeof next !== 'object') return; + cursor = next as Record; + } + const lastSegment = path[path.length - 1]; + if (!lastSegment) return; + cursor[lastSegment] = replacement; +} + +function cloneJson(value: T): T { + return JSON.parse(JSON.stringify(value)) as T; +} + +function stripProjectPrefix(value: string): string { + return value.startsWith('projects/') ? value.slice('projects/'.length) : value; +} + +function canFailover(classification: ErrorClassification): boolean { + return classification.kind === 'unauthorized' || classification.kind === 'rate_limited' || classification.kind === 'server'; +} + +function readConfiguredStrategy(): AccountStrategy | undefined { + const value = process.env.STITCH_ACCOUNT_STRATEGY; + if (value === undefined || value.trim() === '') return undefined; + if (value === 'least_used' || value === 'round_robin') return value; + throw new Error('STITCH_ACCOUNT_STRATEGY must be least_used or round_robin'); +} + +async function defaultSleep(milliseconds: number): Promise { + await new Promise((resolve) => setTimeout(resolve, milliseconds)); +} diff --git a/src/services/stitch-pool/account-selector.ts b/src/services/stitch-pool/account-selector.ts new file mode 100644 index 0000000..2958c0b --- /dev/null +++ b/src/services/stitch-pool/account-selector.ts @@ -0,0 +1,53 @@ +import type { AccountStrategy } from './types.js'; + +export interface SelectableAccount { + id: string; + usageCount: number; + lastUsedAt?: number; +} + +export class AccountSelector { + private strategy: AccountStrategy; + private cursor = 0; + + constructor(strategy: AccountStrategy = 'least_used') { + this.strategy = strategy; + } + + setStrategy(strategy: AccountStrategy): void { + this.strategy = strategy; + this.cursor = 0; + } + + getStrategy(): AccountStrategy { + return this.strategy; + } + + select(candidates: SelectableAccount[]): SelectableAccount { + if (candidates.length === 0) throw new Error('No healthy Stitch accounts available'); + if (this.strategy === 'least_used') return this.selectLeastUsed(candidates); + return this.selectRoundRobin(candidates); + } + + private selectLeastUsed(candidates: SelectableAccount[]): SelectableAccount { + const selected = [...candidates].sort((left, right) => { + const usageDelta = left.usageCount - right.usageCount; + if (usageDelta !== 0) return usageDelta; + const leftLastUsed = left.lastUsedAt ?? 0; + const rightLastUsed = right.lastUsedAt ?? 0; + const timeDelta = leftLastUsed - rightLastUsed; + return timeDelta !== 0 ? timeDelta : left.id.localeCompare(right.id); + })[0]; + if (!selected) throw new Error('No healthy Stitch accounts available'); + return selected; + } + + private selectRoundRobin(candidates: SelectableAccount[]): SelectableAccount { + const ordered = [...candidates].sort((left, right) => left.id.localeCompare(right.id)); + const index = this.cursor % ordered.length; + this.cursor = (this.cursor + 1) % ordered.length; + const selected = ordered[index]; + if (!selected) throw new Error('No healthy Stitch accounts available'); + return selected; + } +} diff --git a/src/services/stitch-pool/credential-provider.ts b/src/services/stitch-pool/credential-provider.ts new file mode 100644 index 0000000..a6cba57 --- /dev/null +++ b/src/services/stitch-pool/credential-provider.ts @@ -0,0 +1,126 @@ +import type { AccountReference, ResolvedAccount } from './types.js'; + +const ACCOUNT_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$/; +const ENV_VAR_PATTERN = /^[A-Za-z_][A-Za-z0-9_]*$/; + +export class CredentialProvider { + private readonly env: NodeJS.ProcessEnv; + private readonly refs: AccountReference[]; + + constructor(env: NodeJS.ProcessEnv = process.env) { + this.env = env; + this.refs = parseAccountReferences(env.STITCH_ACCOUNTS); + if (this.refs.length === 0 && env.STITCH_API_KEY) { + this.refs.push({ id: 'default', envVar: 'STITCH_API_KEY', legacy: true }); + } + } + + get references(): AccountReference[] { + return this.refs.map((ref) => ({ ...ref })); + } + + get isLegacy(): boolean { + return this.refs.length === 1 && this.refs[0]?.legacy === true; + } + + get hasConfiguration(): boolean { + return this.refs.length > 0; + } + + resolve(id: string): ResolvedAccount | undefined { + const ref = this.refs.find((candidate) => candidate.id === id); + const apiKey = ref ? this.env[ref.envVar]?.trim() : undefined; + return ref && apiKey ? { ...ref, apiKey } : undefined; + } + + resolveAll(): ResolvedAccount[] { + return this.refs + .map((ref) => this.resolve(ref.id)) + .filter((account): account is ResolvedAccount => account !== undefined); + } + + resolvePrimary(): ResolvedAccount | undefined { + return this.resolveAll()[0]; + } +} + +export function parseAccountReferences(raw: string | undefined): AccountReference[] { + if (!raw?.trim()) return []; + + const value = raw.trim(); + const parsed = value.startsWith('[') || value.startsWith('{') ? parseStructuredReferences(value) : parseCompactReferences(value); + const seen = new Set(); + + return parsed.map((ref) => { + if (!ACCOUNT_ID_PATTERN.test(ref.id)) { + throw new Error(`Invalid Stitch account id: ${ref.id}`); + } + if (!ENV_VAR_PATTERN.test(ref.envVar)) { + throw new Error(`Stitch account ${ref.id} must reference an environment variable`); + } + if (seen.has(ref.id)) { + throw new Error(`Duplicate Stitch account id: ${ref.id}`); + } + seen.add(ref.id); + return { id: ref.id, envVar: ref.envVar }; + }); +} + +function parseStructuredReferences(raw: string): AccountReference[] { + let parsed: unknown; + try { + parsed = JSON.parse(raw); + } catch { + throw new Error('STITCH_ACCOUNTS must be valid JSON or id:ENV_VAR pairs'); + } + + if (Array.isArray(parsed)) { + return parsed.map((entry, index) => parseReferenceEntry(entry, `STITCH_ACCOUNTS[${index}]`)); + } + + if (parsed && typeof parsed === 'object') { + const record = parsed as Record; + if (Array.isArray(record.accounts)) { + return record.accounts.map((entry, index) => parseReferenceEntry(entry, `STITCH_ACCOUNTS.accounts[${index}]`)); + } + return Object.entries(record).map(([id, envVar]) => { + if (typeof envVar !== 'string') { + throw new Error(`Stitch account ${id} must reference an environment variable`); + } + return { id, envVar }; + }); + } + + throw new Error('STITCH_ACCOUNTS must be a JSON array or object'); +} + +function parseReferenceEntry(entry: unknown, location: string): AccountReference { + if (!entry || typeof entry !== 'object' || Array.isArray(entry)) { + throw new Error(`${location} must contain { id, env }`); + } + + const record = entry as Record; + for (const [key, value] of Object.entries(record)) { + if (!['id', 'env', 'envVar', 'keyEnv'].includes(key) || typeof value !== 'string') { + throw new Error(`${location} accepts only id and environment-variable references, not plaintext secrets`); + } + } + + const id = typeof record.id === 'string' ? record.id : undefined; + const envVar = [record.env, record.envVar, record.keyEnv].find((value): value is string => typeof value === 'string'); + if (!id || !envVar) { + throw new Error(`${location} must contain id and env`); + } + return { id, envVar }; +} + +function parseCompactReferences(raw: string): AccountReference[] { + return raw.split(',').filter(Boolean).map((entry) => { + const separator = entry.indexOf('=') >= 0 ? '=' : ':'; + const [id, envVar] = entry.split(separator, 2).map((part) => part.trim()); + if (!id || !envVar) { + throw new Error('STITCH_ACCOUNTS must use id:ENV_VAR pairs'); + } + return { id, envVar }; + }); +} diff --git a/src/services/stitch-pool/errors.ts b/src/services/stitch-pool/errors.ts new file mode 100644 index 0000000..dbcefcd --- /dev/null +++ b/src/services/stitch-pool/errors.ts @@ -0,0 +1,101 @@ +export type ErrorKind = 'unauthorized' | 'rate_limited' | 'forbidden' | 'server' | 'other'; + +export interface ErrorClassification { + kind: ErrorKind; + status?: number; + code?: string | number; + message: string; + retryAfterMs?: number; +} + +export class StitchRequestError extends Error { + readonly status?: number; + readonly code?: string | number; + readonly retryAfterMs?: number; + readonly details?: string; + + constructor(input: { + message: string; + status?: number; + code?: string | number; + retryAfterMs?: number; + details?: string; + }) { + super(input.message); + this.name = 'StitchRequestError'; + this.status = input.status; + this.code = input.code; + this.retryAfterMs = input.retryAfterMs; + this.details = input.details; + } +} + +export function classifyError(error: unknown): ErrorClassification { + const record = error && typeof error === 'object' ? error as Record : {}; + const status = readNumber(record.status) ?? + readNumber(record.statusCode) ?? + readNumber(record.httpStatus) ?? + readStatusFromMessage(error instanceof Error ? error.message : undefined); + const code = readCode(record.code) ?? readCode(record.errorCode); + const message = error instanceof Error + ? error.message + : typeof record.message === 'string' + ? record.message + : String(error); + const normalized = `${String(code ?? '')} ${message}`.toUpperCase(); + const retryAfterMs = readNumber(record.retryAfterMs); + + if ( + status === 401 || + code === 16 || + normalized.includes('UNAUTHENTICATED') || + normalized.includes('INVALID_API_KEY') + ) { + return { kind: 'unauthorized', status, code, message, retryAfterMs }; + } + if ( + status === 429 || + code === 8 || + normalized.includes('RESOURCE_EXHAUSTED') || + normalized.includes('TOO MANY REQUESTS') || + normalized.includes('QUOTA_EXCEEDED') + ) { + return { kind: 'rate_limited', status, code, message, retryAfterMs }; + } + if (status === 403 || code === 7 || normalized.includes('PERMISSION_DENIED') || normalized.includes('FORBIDDEN')) { + return { kind: 'forbidden', status, code, message, retryAfterMs }; + } + if ( + (status !== undefined && status >= 500 && status <= 599) || + code === 13 || + code === 14 || + normalized.includes('INTERNAL') || + normalized.includes('UNAVAILABLE') + ) { + return { kind: 'server', status, code, message, retryAfterMs }; + } + return { kind: 'other', status, code, message, retryAfterMs }; +} + +export function parseRetryAfter(value: string | null | undefined, now = Date.now()): number | undefined { + if (!value) return undefined; + const seconds = Number(value.trim()); + if (Number.isFinite(seconds) && seconds >= 0) return Math.round(seconds * 1000); + const timestamp = Date.parse(value); + if (!Number.isNaN(timestamp)) return Math.max(0, timestamp - now); + return undefined; +} + +function readNumber(value: unknown): number | undefined { + return typeof value === 'number' && Number.isFinite(value) ? value : undefined; +} + +function readCode(value: unknown): string | number | undefined { + return typeof value === 'string' || typeof value === 'number' ? value : undefined; +} + +function readStatusFromMessage(message: string | undefined): number | undefined { + if (!message) return undefined; + const match = message.match(/\b(?:status|HTTP|error)\D{0,4}([45]\d{2})\b/i); + return match?.[1] ? Number(match[1]) : undefined; +} diff --git a/src/services/stitch-pool/index.ts b/src/services/stitch-pool/index.ts new file mode 100644 index 0000000..8624063 --- /dev/null +++ b/src/services/stitch-pool/index.ts @@ -0,0 +1,8 @@ +export { AccountHealthTracker } from './account-health-tracker.js'; +export { AccountSelector } from './account-selector.js'; +export { AccountPool, AccountPoolError, aggregateProjectResults, createAccountPoolFromEnvironment, extractProjectId, extractProjects } from './account-pool.js'; +export { classifyError, parseRetryAfter, StitchRequestError } from './errors.js'; +export { CredentialProvider, parseAccountReferences } from './credential-provider.js'; +export { ProjectAffinityStore } from './project-affinity-store.js'; +export { fetchStitchTool, isSafeToMove, redactSecrets } from './transport.js'; +export type { AccountStatus, AccountReference, AccountStrategy, PersistedAccountState, PersistedPoolState, ResolvedAccount, StitchRequestExecutor } from './types.js'; diff --git a/src/services/stitch-pool/project-affinity-store.ts b/src/services/stitch-pool/project-affinity-store.ts new file mode 100644 index 0000000..d7b8857 --- /dev/null +++ b/src/services/stitch-pool/project-affinity-store.ts @@ -0,0 +1,192 @@ +import { mkdir, readFile, rename, rm, writeFile } from 'node:fs/promises'; +import { dirname, join } from 'node:path'; +import { getStitchDir } from '../../platform/detector.js'; +import type { AccountStrategy, PersistedAccountState, PersistedPoolState, QuotaSnapshot } from './types.js'; + +const EMPTY_STATE: PersistedPoolState = { + version: 1, + accounts: {}, + affinity: {}, +}; + +export interface ProjectAffinityStoreOptions { + filePath?: string; +} + +export class ProjectAffinityStore { + private state: PersistedPoolState = cloneState(EMPTY_STATE); + private writeChain: Promise = Promise.resolve(); + private dirty = false; + + constructor(private readonly options: ProjectAffinityStoreOptions = {}) {} + + get filePath(): string { + return this.options.filePath ?? process.env.STITCH_POOL_STATE_FILE ?? join(getStitchDir(), 'account-pool.json'); + } + + async load(): Promise { + try { + const content = await readFile(this.filePath, 'utf8'); + this.state = normalizeState(JSON.parse(content)); + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== 'ENOENT') throw error; + this.state = cloneState(EMPTY_STATE); + } + } + + get(projectId: string): string | undefined { + return this.state.affinity[projectId]; + } + + has(projectId: string): boolean { + return this.state.affinity[projectId] !== undefined; + } + + set(projectId: string, accountId: string): void { + if (!projectId || !accountId) return; + if (this.state.affinity[projectId] === accountId) return; + this.state.affinity[projectId] = accountId; + this.queueWrite(); + } + + delete(projectId: string): void { + if (!(projectId in this.state.affinity)) return; + delete this.state.affinity[projectId]; + this.queueWrite(); + } + + affinities(): Record { + return { ...this.state.affinity }; + } + + getAccountState(accountId: string): PersistedAccountState { + return { ...(this.state.accounts[accountId] ?? {}) }; + } + + setAccountState(accountId: string, accountState: PersistedAccountState): void { + this.state.accounts[accountId] = sanitizeAccountState(accountState); + this.queueWrite(); + } + + setStrategy(strategy: AccountStrategy): void { + if (this.state.strategy === strategy) return; + this.state.strategy = strategy; + this.queueWrite(); + } + + get strategy(): AccountStrategy | undefined { + return this.state.strategy; + } + quota(): QuotaSnapshot | undefined { + return this.state.quota ? { ...this.state.quota } : undefined; + } + + setQuota(quota: QuotaSnapshot): void { + this.state.quota = sanitizeQuota(quota); + this.queueWrite(); + } + + snapshot(): PersistedPoolState { + return cloneState(this.state); + } + + async reset(): Promise { + this.state = cloneState(EMPTY_STATE); + this.queueWrite(); + await this.flush(); + } + + async flush(): Promise { + await this.writeChain; + } + + private queueWrite(): void { + this.dirty = true; + this.writeChain = this.writeChain + .catch(() => undefined) + .then(async () => { + if (!this.dirty) return; + this.dirty = false; + await this.writeState(); + }); + } + + private async writeState(): Promise { + const path = this.filePath; + const temporaryPath = `${path}.${process.pid}.tmp`; + await mkdir(dirname(path), { recursive: true }); + await writeFile(temporaryPath, `${JSON.stringify(this.state, null, 2)}\n`, { encoding: 'utf8', mode: 0o600 }); + try { + await rename(temporaryPath, path); + } catch { + await writeFile(path, `${JSON.stringify(this.state, null, 2)}\n`, { encoding: 'utf8', mode: 0o600 }); + await rm(temporaryPath, { force: true }); + } + } +} + +function cloneState(state: PersistedPoolState): PersistedPoolState { + return { + version: 1, + strategy: state.strategy, + accounts: Object.fromEntries( + Object.entries(state.accounts).map(([id, accountState]) => [id, sanitizeAccountState(accountState)]), + ), + affinity: { ...state.affinity }, + quota: state.quota ? { ...state.quota } : undefined, + }; +} + +function normalizeState(value: unknown): PersistedPoolState { + if (!value || typeof value !== 'object') return cloneState(EMPTY_STATE); + const record = value as Record; + const accounts: Record = {}; + if (record.accounts && typeof record.accounts === 'object') { + for (const [id, accountState] of Object.entries(record.accounts as Record)) { + if (accountState && typeof accountState === 'object') { + accounts[id] = sanitizeAccountState(accountState as PersistedAccountState); + } + } + } + + const affinity: Record = {}; + if (record.affinity && typeof record.affinity === 'object') { + for (const [projectId, accountId] of Object.entries(record.affinity as Record)) { + if (typeof accountId === 'string' && accountId.length > 0) affinity[projectId] = accountId; + } + } + + const strategy = record.strategy === 'round_robin' || record.strategy === 'least_used' ? record.strategy : undefined; + const quota = record.quota && typeof record.quota === 'object' + ? sanitizeQuota(record.quota as QuotaSnapshot) + : undefined; + return { version: 1, strategy, accounts, affinity, quota }; +} + +function sanitizeAccountState(value: PersistedAccountState): PersistedAccountState { + const sanitized: PersistedAccountState = {}; + if (typeof value.enabled === 'boolean') sanitized.enabled = value.enabled; + if (typeof value.disabledReason === 'string') sanitized.disabledReason = value.disabledReason.slice(0, 120); + if (typeof value.cooldownUntil === 'number' && Number.isFinite(value.cooldownUntil)) sanitized.cooldownUntil = value.cooldownUntil; + if (typeof value.usageCount === 'number' && Number.isFinite(value.usageCount)) sanitized.usageCount = Math.max(0, Math.floor(value.usageCount)); + if (typeof value.lastUsedAt === 'number' && Number.isFinite(value.lastUsedAt)) sanitized.lastUsedAt = value.lastUsedAt; + if (typeof value.consecutiveFailures === 'number' && Number.isFinite(value.consecutiveFailures)) { + sanitized.consecutiveFailures = Math.max(0, Math.floor(value.consecutiveFailures)); + } + return sanitized; +} + +function sanitizeQuota(value: QuotaSnapshot): QuotaSnapshot { + const snapshot: QuotaSnapshot = { + source: 'stitch-web', + observedAt: Number.isFinite(value.observedAt) ? value.observedAt : Date.now(), + }; + if (typeof value.accountId === 'string' && value.accountId.length > 0) snapshot.accountId = value.accountId.slice(0, 64); + if (typeof value.used === 'number' && Number.isFinite(value.used) && value.used >= 0) snapshot.used = value.used; + if (typeof value.allocated === 'number' && Number.isFinite(value.allocated) && value.allocated >= 0) snapshot.allocated = value.allocated; + if (typeof value.imageProUsed === 'number' && Number.isFinite(value.imageProUsed) && value.imageProUsed >= 0) snapshot.imageProUsed = value.imageProUsed; + if (typeof value.imageProAllocated === 'number' && Number.isFinite(value.imageProAllocated) && value.imageProAllocated >= 0) { + snapshot.imageProAllocated = value.imageProAllocated; + } + return snapshot; +} diff --git a/src/services/stitch-pool/transport.ts b/src/services/stitch-pool/transport.ts new file mode 100644 index 0000000..f5211fc --- /dev/null +++ b/src/services/stitch-pool/transport.ts @@ -0,0 +1,132 @@ +import { classifyError, parseRetryAfter, StitchRequestError } from './errors.js'; +import type { ResolvedAccount } from './types.js'; + +const DEFAULT_STITCH_URL = 'https://stitch.googleapis.com/mcp'; +let requestSequence = 0; + +export interface StitchTransportOptions { + account: ResolvedAccount; + toolName: string; + args: Record; + url?: string; + fetchImpl?: typeof fetch; +} + +export async function fetchStitchTool(options: StitchTransportOptions): Promise { + const fetchImpl = options.fetchImpl ?? globalThis.fetch; + const url = options.url ?? process.env.STITCH_MCP_URL ?? process.env.STITCH_HOST ?? DEFAULT_STITCH_URL; + const response = await fetchImpl(url, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Accept: 'application/json, text/event-stream', + 'X-Goog-Api-Key': options.account.apiKey, + }, + body: JSON.stringify({ + jsonrpc: '2.0', + method: 'tools/call', + params: { name: options.toolName, arguments: options.args }, + id: `${Date.now()}-${requestSequence++}`, + }), + }); + + const body = await readBody(response); + const errorPayload = extractRpcError(body); + if (!response.ok || errorPayload) { + const status = response.ok ? readErrorStatus(errorPayload) : response.status; + const message = errorPayload?.message ?? readErrorMessage(body) ?? `Stitch API request failed with status ${response.status}`; + throw new StitchRequestError({ + message, + status, + code: errorPayload?.code, + retryAfterMs: getRetryAfter(response), + details: typeof body === 'string' ? body : undefined, + }); + } + + if (isRpcEnvelope(body) && Object.hasOwn(body, 'result')) return body.result; + return body; +} + +export function isSafeToMove(toolName: string): boolean { + const normalized = toolName.toLowerCase(); + return ( + normalized === 'list_projects' || + normalized.startsWith('list_') || + normalized.startsWith('get_') || + normalized.startsWith('read_') || + normalized.startsWith('search_') || + normalized.startsWith('check_') + ); +} + +export function redactSecrets(value: unknown, secrets: string[] = []): unknown { + if (typeof value === 'string') { + return secrets.reduce((result, secret) => (secret ? result.split(secret).join('[REDACTED]') : result), value); + } + if (Array.isArray(value)) return value.map((entry) => redactSecrets(entry, secrets)); + if (!value || typeof value !== 'object') return value; + + const result: Record = {}; + for (const [key, entry] of Object.entries(value as Record)) { + if (/(?:api.?key|token|authorization|secret|password)/i.test(key)) { + result[key] = '[REDACTED]'; + } else { + result[key] = redactSecrets(entry, secrets); + } + } + return result; +} + +async function readBody(response: Response): Promise { + if (typeof response.text === 'function') { + const text = await response.text(); + if (!text.trim()) return {}; + try { + return JSON.parse(text); + } catch { + const event = text.split('\n').find((line) => line.startsWith('data:'))?.slice(5).trim(); + if (event) { + try { + return JSON.parse(event); + } catch { + return text; + } + } + return text; + } + } + return response.json(); +} + +function extractRpcError(body: unknown): { code?: string | number; message?: string } | undefined { + if (!body || typeof body !== 'object') return undefined; + const error = (body as Record).error; + if (!error || typeof error !== 'object') return undefined; + const record = error as Record; + return { + code: typeof record.code === 'string' || typeof record.code === 'number' ? record.code : undefined, + message: typeof record.message === 'string' ? record.message : undefined, + }; +} + +function isRpcEnvelope(body: unknown): body is { result?: unknown } { + return !!body && typeof body === 'object' && ('result' in body || 'error' in body); +} + +function readErrorStatus(error: { code?: string | number; message?: string } | undefined): number | undefined { + if (typeof error?.code === 'number' && error.code >= 100 && error.code <= 599) return error.code; + const classification = classifyError(error?.code === undefined ? error?.message : `${error.code}`); + return classification.status; +} + +function readErrorMessage(body: unknown): string | undefined { + if (typeof body === 'string') return body.slice(0, 500); + if (!body || typeof body !== 'object') return undefined; + const message = (body as Record).message; + return typeof message === 'string' ? message : undefined; +} + +function getRetryAfter(response: Response): number | undefined { + return parseRetryAfter(response.headers?.get('Retry-After')); +} diff --git a/src/services/stitch-pool/types.ts b/src/services/stitch-pool/types.ts new file mode 100644 index 0000000..21c9c68 --- /dev/null +++ b/src/services/stitch-pool/types.ts @@ -0,0 +1,59 @@ +export type AccountStrategy = 'least_used' | 'round_robin'; + +export interface AccountReference { + id: string; + envVar: string; + legacy?: boolean; +} + +export interface ResolvedAccount extends AccountReference { + apiKey: string; +} + +export interface PersistedAccountState { + enabled?: boolean; + disabledReason?: string; + cooldownUntil?: number; + usageCount?: number; + lastUsedAt?: number; + consecutiveFailures?: number; +} + +export interface QuotaSnapshot { + accountId?: string; + source: 'stitch-web'; + used?: number; + allocated?: number; + imageProUsed?: number; + imageProAllocated?: number; + observedAt: number; +} + +export interface PersistedPoolState { + version: 1; + strategy?: AccountStrategy; + accounts: Record; + affinity: Record; + quota?: QuotaSnapshot; +} + +export interface AccountStatus { + id: string; + envVar: string; + legacy: boolean; + configured: boolean; + enabled: boolean; + cooldownUntil?: number; + usageCount: number; + lastUsedAt?: number; + consecutiveFailures: number; + disabledReason?: string; +} + +export type StitchRequestExecutor = ( + account: ResolvedAccount, + toolName: string, + args: Record, +) => Promise; + +export type Sleep = (milliseconds: number) => Promise; diff --git a/tests/services/stitch-pool/account-pool.test.ts b/tests/services/stitch-pool/account-pool.test.ts new file mode 100644 index 0000000..616c3f2 --- /dev/null +++ b/tests/services/stitch-pool/account-pool.test.ts @@ -0,0 +1,273 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +import { mkdtemp, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { AccountHealthTracker } from '../../../src/services/stitch-pool/account-health-tracker.js'; +import { AccountSelector } from '../../../src/services/stitch-pool/account-selector.js'; +import { + AccountPool, + aggregateProjectResults, + extractProjectId, +} from '../../../src/services/stitch-pool/account-pool.js'; +import { classifyError, parseRetryAfter } from '../../../src/services/stitch-pool/errors.js'; +import { CredentialProvider } from '../../../src/services/stitch-pool/credential-provider.js'; +import { fetchStitchTool, redactSecrets } from '../../../src/services/stitch-pool/transport.js'; +import { ProjectAffinityStore } from '../../../src/services/stitch-pool/project-affinity-store.js'; +import type { ResolvedAccount } from '../../../src/services/stitch-pool/types.js'; + +const temporaryDirectories: string[] = []; + +afterEach(async () => { + await Promise.all(temporaryDirectories.splice(0).map((directory) => rm(directory, { recursive: true, force: true }))); +}); + +describe('CredentialProvider', () => { + test('resolves env references without accepting plaintext keys', () => { + const provider = new CredentialProvider({ + STITCH_ACCOUNTS: JSON.stringify([ + { id: 'personal', env: 'STITCH_KEY_PERSONAL' }, + { id: 'work', env: 'STITCH_KEY_WORK' }, + ]), + STITCH_KEY_PERSONAL: 'test-key-personal', + STITCH_KEY_WORK: 'test-key-work', + }); + + expect(provider.resolveAll().map(({ id, apiKey }) => ({ id, apiKey }))).toEqual([ + { id: 'personal', apiKey: 'test-key-personal' }, + { id: 'work', apiKey: 'test-key-work' }, + ]); + expect(() => new CredentialProvider({ STITCH_ACCOUNTS: '[{"id":"x","apiKey":"inline-key"}]' })).toThrow(/plaintext/); + }); + + test('keeps legacy STITCH_API_KEY mode as one account', () => { + const provider = new CredentialProvider({ STITCH_API_KEY: 'legacy-key' }); + expect(provider.isLegacy).toBe(true); + expect(provider.resolvePrimary()).toMatchObject({ id: 'default', envVar: 'STITCH_API_KEY', apiKey: 'legacy-key' }); + }); +}); + +describe('AccountSelector', () => { + test('least_used selects the lowest usage account', () => { + const selector = new AccountSelector('least_used'); + expect(selector.select([ + { id: 'a', usageCount: 3 }, + { id: 'b', usageCount: 1 }, + ]).id).toBe('b'); + }); + + test('round_robin is opt-in and advances only when selection is needed', () => { + const selector = new AccountSelector('round_robin'); + const candidates = [{ id: 'b', usageCount: 0 }, { id: 'a', usageCount: 0 }]; + expect([selector.select(candidates).id, selector.select(candidates).id, selector.select(candidates).id]).toEqual(['a', 'b', 'a']); + }); +}); + +describe('error and health policy', () => { + test('classifies quota, auth, forbidden, and server failures', () => { + expect(classifyError({ status: 429, message: 'quota' }).kind).toBe('rate_limited'); + expect(classifyError({ message: 'RESOURCE_EXHAUSTED' }).kind).toBe('rate_limited'); + expect(classifyError({ status: 401, message: 'no' }).kind).toBe('unauthorized'); + expect(classifyError({ status: 403, message: 'project permission denied' }).kind).toBe('forbidden'); + expect(classifyError({ status: 503, message: 'upstream' }).kind).toBe('server'); + }); + + test('parses Retry-After seconds and HTTP dates', () => { + expect(parseRetryAfter('2')).toBe(2_000); + expect(parseRetryAfter('1970-01-01T00:00:03.000Z', 1_000)).toBe(2_000); + }); + + test('cooldown expires while unauthorized credentials stay disabled', () => { + let now = 1_000; + const tracker = new AccountHealthTracker(['a'], {}, { now: () => now, cooldownMs: 100 }); + tracker.recordFailure('a', 'rate_limited'); + expect(tracker.isAvailable('a')).toBe(false); + now += 101; + expect(tracker.isAvailable('a')).toBe(true); + + tracker.recordFailure('a', 'unauthorized'); + expect(tracker.isAvailable('a')).toBe(false); + tracker.enable('a'); + expect(tracker.isAvailable('a')).toBe(true); + }); + + test('classifies JSON-RPC RESOURCE_EXHAUSTED and preserves Retry-After', async () => { + const account: ResolvedAccount = { id: 'a', envVar: 'KEY_A', apiKey: 'test-key-a' }; + const fetchImpl = async () => new Response(JSON.stringify({ + jsonrpc: '2.0', + error: { code: 'RESOURCE_EXHAUSTED', message: 'quota exhausted' }, + }), { status: 200, headers: { 'Retry-After': '3' } }); + await expect(fetchStitchTool({ + account, + toolName: 'list_projects', + args: {}, + fetchImpl: fetchImpl as unknown as typeof fetch, + })).rejects.toMatchObject({ retryAfterMs: 3_000 }); + }); + + test('redacts credentials from captured payloads', () => { + expect(redactSecrets({ apiKey: 'test-key-a', accessToken: 'test-token' }, ['test-key-a'])).toEqual({ + apiKey: '[REDACTED]', + accessToken: '[REDACTED]', + }); + }); +}); + +describe('ProjectAffinityStore', () => { + test('persists affinity and account state across store instances', async () => { + const directory = await mkdtemp(join(tmpdir(), 'stitch-pool-')); + temporaryDirectories.push(directory); + const filePath = join(directory, 'state.json'); + const first = new ProjectAffinityStore({ filePath }); + await first.load(); + first.set('project-1', 'work'); + first.setAccountState('work', { enabled: false, disabledReason: 'unauthorized', usageCount: 4 }); + first.setQuota({ + accountId: 'work', + source: 'stitch-web', + used: 3, + allocated: 400, + imageProUsed: 1, + imageProAllocated: 15, + observedAt: 123, + }); + await first.flush(); + + const second = new ProjectAffinityStore({ filePath }); + await second.load(); + expect(second.get('project-1')).toBe('work'); + expect(second.getAccountState('work')).toMatchObject({ enabled: false, disabledReason: 'unauthorized', usageCount: 4 }); + expect(second.quota()).toMatchObject({ accountId: 'work', used: 3, allocated: 400, imageProAllocated: 15 }); + }); +}); + +describe('AccountPool', () => { + async function makePool(strategy: 'least_used' | 'round_robin' = 'least_used') { + const directory = await mkdtemp(join(tmpdir(), 'stitch-pool-')); + temporaryDirectories.push(directory); + const provider = new CredentialProvider({ + STITCH_ACCOUNTS: 'a:KEY_A,b:KEY_B', + KEY_A: 'test-key-a', + KEY_B: 'test-key-b', + }); + return AccountPool.create({ + provider, + store: new ProjectAffinityStore({ filePath: join(directory, 'state.json') }), + strategy, + retryBaseMs: 0, + maxServerRetries: 2, + cooldownMs: 1_000, + sleep: async () => undefined, + }); + } + + test('aggregates and deduplicates projects while building affinity', async () => { + const result = aggregateProjectResults([ + { structuredContent: { projects: [{ projectId: 'one' }, { projectId: 'two' }] } }, + { structuredContent: { projects: [{ projectId: 'two' }, { projectId: 'three' }] } }, + ]) as { structuredContent: { projects: Array<{ projectId: string }> } }; + expect(result.structuredContent.projects.map((project) => project.projectId)).toEqual(['one', 'two', 'three']); + expect(extractProjectId({ name: 'projects/three' })).toBe('three'); + }); + + test('aggregates list_projects from every healthy account and records ownership', async () => { + const pool = await makePool(); + const result = await pool.call('list_projects', {}, async (account) => ({ + structuredContent: { + projects: account.id === 'a' + ? [{ projectId: 'one' }, { projectId: 'shared' }] + : [{ projectId: 'shared' }, { projectId: 'two' }], + }, + })) as { structuredContent: { projects: Array<{ projectId: string }> } }; + + expect(result.structuredContent.projects.map((project) => project.projectId)).toEqual(['one', 'shared', 'two']); + expect(pool.status().affinity).toEqual({ one: 'a', shared: 'a', two: 'b' }); + }); + + test('pins known projects and fails over safe reads after quota errors', async () => { + const pool = await makePool(); + const calls: string[] = []; + const executor = async (account: ResolvedAccount, toolName: string, args: Record) => { + calls.push(`${account.id}:${toolName}:${String(args.projectId ?? '')}`); + if (account.id === 'a' && args.projectId === 'new-project') { + throw Object.assign(new Error('quota'), { status: 429 }); + } + return { projectId: args.projectId ?? 'created' }; + }; + + await pool.call('get_screen', { projectId: 'new-project' }, executor); + await pool.call('get_screen', { projectId: 'new-project' }, executor); + expect(calls).toEqual(['a:get_screen:new-project', 'b:get_screen:new-project', 'b:get_screen:new-project']); + }); + + test('disables an account after 401 and rotates safe reads', async () => { + const pool = await makePool(); + const calls: string[] = []; + const executor = async (account: ResolvedAccount) => { + calls.push(account.id); + if (account.id === 'a') throw Object.assign(new Error('unauthorized'), { status: 401 }); + return { ok: true }; + }; + + await pool.call('get_screen', { projectId: 'auth-project' }, executor); + expect(calls).toEqual(['a', 'b']); + expect(pool.status().accounts.find((account) => account.id === 'a')?.enabled).toBe(false); + }); + + test('retries server failures on the same account before safe failover', async () => { + const pool = await makePool(); + const calls: string[] = []; + const executor = async (account: ResolvedAccount) => { + calls.push(account.id); + if (account.id === 'a') throw Object.assign(new Error('server'), { status: 503 }); + return { ok: true }; + }; + + await pool.call('get_screen', { projectId: 'server-project' }, executor); + expect(calls).toEqual(['a', 'a', 'a', 'b']); + }); + + test('does not move a mutation to another account after quota', async () => { + const pool = await makePool(); + const calls: string[] = []; + const executor = async (account: ResolvedAccount) => { + calls.push(account.id); + throw Object.assign(new Error('quota'), { status: 429 }); + }; + + await expect(pool.call('create_project', { title: 'new' }, executor)).rejects.toThrow(); + expect(calls).toEqual(['a']); + }); + + + test('does not disable an account for project-specific forbidden responses', async () => { + const pool = await makePool(); + await expect(pool.call('get_screen', { projectId: 'forbidden-project' }, async () => { + throw Object.assign(new Error('project permission denied'), { status: 403 }); + })).rejects.toThrow(); + expect(pool.status().accounts.find((account) => account.id === 'a')?.enabled).toBe(true); + }); + test('serializes first access to a new project under concurrency', async () => { + const pool = await makePool('round_robin'); + const calls: string[] = []; + const executor = async (account: ResolvedAccount) => { + calls.push(account.id); + return { ok: true }; + }; + + await Promise.all(Array.from({ length: 5 }, () => pool.call('get_screen', { projectId: 'concurrent' }, executor))); + expect(new Set(calls).size).toBe(1); + expect(pool.status().affinity.concurrent).toBe(calls[0]); + }); + + test('keeps legacy single-account behavior', async () => { + const provider = new CredentialProvider({ STITCH_API_KEY: 'legacy-key' }); + const pool = await AccountPool.create({ provider, sleep: async () => undefined }); + const seen: string[] = []; + await pool.call('list_projects', {}, async (account) => { + seen.push(account.id); + return { projects: [] }; + }); + expect(pool.isLegacy).toBe(true); + expect(seen).toEqual(['default']); + }); +});