Add cross-process TPM rate limiter

ts/packages/benchmarks/README.AUTOGEN.md

@@ -3,7 +3,7 @@
 
 <!-- AUTOGEN:DOCS:START -->
 
-<!-- AUTOGEN:DOCS:HASH:sha256=d94460dd448816fb1ae22304eb0522c647ecdc03c3c08d2738ad914689282f8d -->
+<!-- AUTOGEN:DOCS:HASH:sha256=6add7220cd2366adea51b91f59e3a7c5a4cd7f6ba61ebef9308335d80c2e4027 -->
 <!-- AUTOGEN:DOCS:SOURCE: ./README.md (hand-written documentation; this file is the AI-generated companion) -->
 
 # @typeagent/benchmarks — AI-generated documentation
@@ -49,13 +49,13 @@ _None._
 - [./src/core/model-prices.generated.json](./src/core/model-prices.generated.json)
 - [./src/core/paths.ts](./src/core/paths.ts)
 - [./src/core/prices.ts](./src/core/prices.ts)
+- [./src/core/rateLimiter.ts](./src/core/rateLimiter.ts)
 - [./src/core/types.ts](./src/core/types.ts)
 - [./src/translationBench/action-parameters-grader.generated.json](./src/translationBench/action-parameters-grader.generated.json)
-- [./src/translationBench/catalog.generated.json](./src/translationBench/catalog.generated.json)
-- _…and 32 more under `./src/`._
+- _…and 34 more under `./src/`._
 
 ---
 
-_Auto-generated against commit `d917a1fbedd52940d73141b0152b6ae508258b54` on `2026-08-12T23:28:04.947Z` by `docs-generate.yml`. Links validated at that commit; the working tree may have drifted by up to 24h. Re-run `pnpm --filter @typeagent/benchmarks docs:verify-links` to spot-check._
+_Auto-generated against commit `d9cf714f7d151120013855a722e7583cbf2c30d7` on `2026-08-13T01:26:47.704Z` by `docs-generate.yml`. Links validated at that commit; the working tree may have drifted by up to 24h. Re-run `pnpm --filter @typeagent/benchmarks docs:verify-links` to spot-check._
 
 <!-- AUTOGEN:DOCS:END -->

ts/packages/benchmarks/src/core/rateLimiter.ts

@@ -0,0 +1,318 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import fs from "node:fs";
+import path from "node:path";
+import { randomUUID } from "node:crypto";
+import { DatabaseSync, type StatementSync } from "node:sqlite";
+
+const WINDOW_MS = 60_000;
+const MAX_SLEEP_MS = 1_000;
+// Long enough for multi-minute TB translates + retries; pending claims older
+// than this are treated as abandoned (process crash) and purged.
+const STALE_MS = 30 * 60_000;
+const BUSY_TIMEOUT_MS = 15_000;
+const SQLITE_BUSY = 5;
+const OPEN_MAX_ATTEMPTS = 50;
+const OPEN_RETRY_MIN_MS = 20;
+const OPEN_RETRY_JITTER_MS = 30;
+
+export interface RateLimiterOptions {
+    dbPath: string;
+    estTokensPerCall?: number;
+    maxWaitMs?: number;
+    onWait?: (model: string, waitedMs: number, waitMs: number) => void;
+}
+
+export interface RateLimiter {
+    disabledFor(model: string): boolean;
+    run<T>(
+        model: string,
+        est: number | undefined,
+        fn: () => Promise<{ result: T; actualTokens: number | undefined }>,
+    ): Promise<T>;
+    close(): void;
+}
+
+export type TpmLimits = Readonly<Record<string, number>>;
+
+interface Reservation {
+    id: string | undefined;
+    waitMs: number;
+}
+
+interface ClaimRow {
+    created_at: number;
+    tokens: number;
+}
+
+function sleep(ms: number): Promise<void> {
+    return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+function isBusyError(error: unknown): boolean {
+    return (
+        typeof error === "object" &&
+        error !== null &&
+        (error as { errcode?: number }).errcode === SQLITE_BUSY
+    );
+}
+
+function openDatabase(dbPath: string): DatabaseSync {
+    fs.mkdirSync(path.dirname(dbPath), { recursive: true });
+    let lastError: unknown;
+    for (let attempt = 0; attempt < OPEN_MAX_ATTEMPTS; attempt++) {
+        let db: DatabaseSync | undefined;
+        try {
+            db = new DatabaseSync(dbPath);
+            db.exec(`PRAGMA busy_timeout = ${BUSY_TIMEOUT_MS}`);
+            db.exec("PRAGMA journal_mode = WAL");
+            db.exec("PRAGMA synchronous = NORMAL");
+            db.exec(
+                "CREATE TABLE IF NOT EXISTS claims (" +
+                    "id TEXT PRIMARY KEY, " +
+                    "model TEXT NOT NULL, " +
+                    "tokens REAL NOT NULL, " +
+                    "created_at INTEGER NOT NULL, " +
+                    "pending INTEGER NOT NULL)",
+            );
+            db.exec(
+                "CREATE INDEX IF NOT EXISTS idx_claims_model_time " +
+                    "ON claims (model, created_at)",
+            );
+            return db;
+        } catch (error) {
+            lastError = error;
+            if (db !== undefined) {
+                try {
+                    db.close();
+                } catch {
+                    // no-op
+                }
+            }
+            if (!isBusyError(error)) {
+                throw error;
+            }
+            const until =
+                Date.now() +
+                OPEN_RETRY_MIN_MS +
+                Math.floor(Math.random() * OPEN_RETRY_JITTER_MS);
+            // Yield the event loop instead of a tight spin-wait.
+            const sab = new SharedArrayBuffer(4);
+            Atomics.wait(
+                new Int32Array(sab),
+                0,
+                0,
+                Math.max(1, until - Date.now()),
+            );
+        }
+    }
+    throw lastError;
+}
+
+class Ledger {
+    private readonly insertStmt: StatementSync;
+    private readonly settleStmt: StatementSync;
+    private readonly insertSettledStmt: StatementSync;
+    private readonly purgeExpiredStmt: StatementSync;
+    private readonly purgeStaleStmt: StatementSync;
+    private readonly usedStmt: StatementSync;
+    private readonly oldestStmt: StatementSync;
+
+    constructor(
+        private readonly db: DatabaseSync,
+        private readonly tpmLimits: TpmLimits,
+    ) {
+        this.insertStmt = db.prepare(
+            "INSERT INTO claims (id, model, tokens, created_at, pending) " +
+                "VALUES (?, ?, ?, ?, 1)",
+        );
+        this.settleStmt = db.prepare(
+            "UPDATE claims SET tokens = ?, pending = 0 WHERE id = ?",
+        );
+        this.insertSettledStmt = db.prepare(
+            "INSERT OR REPLACE INTO claims " +
+                "(id, model, tokens, created_at, pending) VALUES (?, ?, ?, ?, 0)",
+        );
+        this.purgeExpiredStmt = db.prepare(
+            "DELETE FROM claims WHERE pending = 0 AND created_at <= ?",
+        );
+        this.purgeStaleStmt = db.prepare(
+            "DELETE FROM claims WHERE pending = 1 AND created_at <= ?",
+        );
+        this.usedStmt = db.prepare(
+            "SELECT COALESCE(SUM(tokens), 0) AS used " +
+                "FROM claims WHERE model = ? AND created_at > ?",
+        );
+        this.oldestStmt = db.prepare(
+            "SELECT created_at, tokens FROM claims " +
+                "WHERE model = ? AND created_at > ? ORDER BY created_at ASC",
+        );
+    }
+
+    private transaction<T>(fn: () => T): T {
+        this.db.exec("BEGIN IMMEDIATE");
+        try {
+            const out = fn();
+            this.db.exec("COMMIT");
+            return out;
+        } catch (error) {
+            try {
+                this.db.exec("ROLLBACK");
+            } catch {
+                // no-op
+            }
+            throw error;
+        }
+    }
+
+    private waitForCapacity(
+        model: string,
+        limit: number,
+        need: number,
+        now: number,
+    ): number {
+        const excess = need - limit;
+        let freed = 0;
+        const rows = this.oldestStmt.all(
+            model,
+            now - WINDOW_MS,
+        ) as unknown as ClaimRow[];
+        for (const row of rows) {
+            freed += row.tokens;
+            if (freed >= excess) {
+                return Math.max(5, row.created_at + WINDOW_MS - now);
+            }
+        }
+        return Math.max(5, WINDOW_MS);
+    }
+
+    reserve(model: string, cost: number): Reservation {
+        const limit = this.tpmLimits[model];
+        const need = Math.min(cost, limit);
+        return this.transaction(() => {
+            const now = Date.now();
+            this.purgeExpiredStmt.run(now - WINDOW_MS);
+            this.purgeStaleStmt.run(now - STALE_MS);
+            const { used } = this.usedStmt.get(model, now - WINDOW_MS) as {
+                used: number;
+            };
+            if (used + need <= limit) {
+                const id = randomUUID();
+                this.insertStmt.run(id, model, need, now);
+                return { id, waitMs: 0 };
+            }
+            return {
+                id: undefined,
+                waitMs: this.waitForCapacity(model, limit, used + need, now),
+            };
+        });
+    }
+
+    settle(id: string, model: string, actualCost: number): void {
+        this.transaction(() => {
+            const result = this.settleStmt.run(actualCost, id);
+            if (result.changes === 0) {
+                this.insertSettledStmt.run(id, model, actualCost, Date.now());
+            }
+        });
+    }
+}
+
+export function createRateLimiter(
+    limits: TpmLimits,
+    options: RateLimiterOptions,
+): RateLimiter {
+    const tpmLimits: Record<string, number> = {};
+    for (const [model, tpm] of Object.entries(limits)) {
+        if (Number.isFinite(tpm) && tpm > 0) {
+            tpmLimits[model] = tpm;
+        }
+    }
+
+    let db: DatabaseSync | undefined;
+    let ledger: Ledger | undefined;
+    if (Object.keys(tpmLimits).length > 0) {
+        db = openDatabase(options.dbPath);
+        ledger = new Ledger(db, tpmLimits);
+    }
+
+    async function admit(model: string, estCost: number): Promise<string> {
+        const activeLedger = ledger as Ledger;
+        const startedAt = Date.now();
+        for (;;) {
+            const reservation = activeLedger.reserve(model, estCost);
+            if (reservation.id !== undefined) {
+                return reservation.id;
+            }
+            const waited = Date.now() - startedAt;
+            if (
+                options.maxWaitMs !== undefined &&
+                waited >= options.maxWaitMs
+            ) {
+                throw new Error(
+                    `rate limiter: exceeded max wait ${options.maxWaitMs}ms for ${model}`,
+                );
+            }
+            options.onWait?.(model, waited, reservation.waitMs);
+            await sleep(Math.min(reservation.waitMs, MAX_SLEEP_MS));
+        }
+    }
+
+    async function run<T>(
+        model: string,
+        est: number | undefined,
+        fn: () => Promise<{ result: T; actualTokens: number | undefined }>,
+    ): Promise<T> {
+        if (ledger === undefined || tpmLimits[model] === undefined) {
+            return (await fn()).result;
+        }
+
+        const estCost =
+            est !== undefined && Number.isFinite(est) && est > 0
+                ? est
+                : options.estTokensPerCall;
+        if (estCost === undefined || !(estCost > 0)) {
+            throw new Error(
+                `rate limiter: no positive token estimate for ${model}`,
+            );
+        }
+
+        const id = await admit(model, estCost);
+        let actual = estCost;
+        try {
+            const out = await fn();
+            actual =
+                out.actualTokens !== undefined &&
+                Number.isFinite(out.actualTokens) &&
+                out.actualTokens > 0
+                    ? out.actualTokens
+                    : estCost;
+            return out.result;
+        } finally {
+            try {
+                (ledger as Ledger).settle(id, model, actual);
+            } catch (error) {
+                const message =
+                    error instanceof Error ? error.message : String(error);
+                console.error(
+                    `[rate-limit] settle failed model=${model} id=${id} actual=${actual}: ${message}`,
+                );
+            }
+        }
+    }
+
+    return {
+        disabledFor(model: string): boolean {
+            return tpmLimits[model] === undefined;
+        },
+        close(): void {
+            if (db !== undefined) {
+                db.close();
+                db = undefined;
+                ledger = undefined;
+            }
+        },
+        run,
+    };
+}

ts/packages/benchmarks/src/index.ts

@@ -4,4 +4,5 @@
 export * from "./core/paths.js";
 export * from "./core/types.js";
 export * from "./core/prices.js";
+export * from "./core/rateLimiter.js";
 export * from "./translationBench/index.js";

ts/packages/benchmarks/src/translationBench/index.ts

@@ -2,4 +2,5 @@
 // Licensed under the MIT License.
 
 export * from "./catalog.js";
+export * from "./runConfig.js";
 export * from "./synthesizer/index.js";