feat: implement stitch account pooling system with load balancing, health tracking, and transport management

.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

README.md

@@ -127,12 +127,16 @@ stitch-mcp <command>
 | `serve -p <id>` | Preview project screens locally |
 | `screens -p <id>` | Browse screens in terminal |
 | `view` | Interactive resource browser |
+| `dashboard` | Run the local account and quota dashboard |
 | **Build** | |
 | `site -p <id>` | 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 |

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"
   },

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<string, Record<string, never>> = {
+  name: 'accounts',
+  description: 'Manage Stitch API-key accounts',
+  arguments: '<action> [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;
+    }
+  },
+};

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<void> {
+    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 <id>|disable <id>');
+    }
+  }
+
+  private async testAccounts(pool: AccountPool): Promise<void> {
+    const results: Array<Record<string, unknown>> = [];
+    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} <id>`);
+    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));
+  }
+}

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<undefined, DashboardOptions> = {
+  name: 'dashboard',
+  description: 'Run the local Stitch account and quota dashboard',
+  options: [
+    { flags: '--port <number>', description: 'Dashboard port', fn: (value) => Number.parseInt(value, 10), defaultValue: 4173 },
+    { flags: '--host <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;
+    }
+  },
+};

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<DashboardResult> {
+    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<Response> {
+    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<Response> {
+    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<string, unknown>;
+      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);
+})();`;
+}