From 5fca0a919625491750adf3b6efbddc03c66d273a Mon Sep 17 00:00:00 2001 From: Hardik Bhatia Date: Sat, 26 Sep 2026 21:00:49 +0530 Subject: [PATCH] feat(jobs): persist encrypted work with leases quotas and retention --- apps/gateway/src/app.ts | 89 ++++++++------------- docs/deployment.md | 30 ++++++++ docs/openapi.yaml | 65 +++++++++++++--- packages/storage/src/index.ts | 11 +++ packages/storage/src/jobs.test.ts | 34 +++++++++ packages/storage/src/jobs.ts | 123 ++++++++++++++++++++++++++++++ 6 files changed, 284 insertions(+), 68 deletions(-) create mode 100644 docs/deployment.md create mode 100644 packages/storage/src/jobs.test.ts create mode 100644 packages/storage/src/jobs.ts diff --git a/apps/gateway/src/app.ts b/apps/gateway/src/app.ts index 853b658..2b3d926 100644 --- a/apps/gateway/src/app.ts +++ b/apps/gateway/src/app.ts @@ -34,7 +34,7 @@ import { type StoredIntegration, } from "@pyro/contracts"; import { ConcurrentQueue, QueueFullError } from "@pyro/queue"; -import { CachedDocument, decryptText, openDatabase, PolicyStore, revisionOf, type PolicyRecord, policyHash } from "@pyro/storage"; +import { CachedDocument, decryptText, openDatabase, PolicyStore, revisionOf, type PolicyRecord, policyHash, DurableJobs, DurableWorker, consumeQuota, JobCapacityError, JobConflict } from "@pyro/storage"; import { decisionDeliveries, DeliveryWorker } from "@pyro/integrations"; import type { GatewayConfig } from "./config.js"; import { GatewayMetrics } from "./metrics.js"; @@ -43,15 +43,6 @@ interface SecretFile { typesafeApiKey?: StoredSecret; } -interface JobRecord { - id: string; - appId: string; - status: "queued" | "running" | "complete" | "failed"; - createdAt: string; - completedAt?: string; - decision?: ClassificationDecision; - error?: string; -} function localRuleDecision( id: string, @@ -203,8 +194,7 @@ export async function buildGateway(config: GatewayConfig): Promise(); - const rateLimits = new Map(); + const jobs = new DurableJobs(database, config.controlPlaneSecret, "classification_jobs", Math.min(config.queueMaxDepth, 200)); const circuit = { consecutiveFailures: 0, openUntil: 0 }; const authenticate = async (candidate: string | undefined): Promise => { @@ -221,18 +211,15 @@ export async function buildGateway(config: GatewayConfig): Promise candidate.id === (apiKey.appId ?? "default")); if (!firewallApp) return reply.code(403).send({ error: "This API key is not assigned to an application." }); if (!firewallApp.enabled) return reply.code(403).send({ error: "This application is disabled." }); - const rateLimit = apiKey.rateLimitPerMinute ?? firewallApp.rateLimitPerMinute; - if (rateLimit) { - const window = Math.floor(Date.now() / 60_000); - const current = rateLimits.get(apiKey.id); - const count = current?.window === window ? current.count + 1 : 1; - rateLimits.set(apiKey.id, { window, count }); - reply.header("X-RateLimit-Limit", rateLimit); - reply.header("X-RateLimit-Remaining", Math.max(0, rateLimit - count)); - if (count > rateLimit) { - reply.header("Retry-After", 60 - Math.floor((Date.now() % 60_000) / 1_000)); - return reply.code(429).send({ error: "Application rate limit exceeded." }); - } + const limits = [ + ...(firewallApp.rateLimitPerMinute ? [{ id: `app:${firewallApp.id}`, limit: firewallApp.rateLimitPerMinute }] : []), + ...(apiKey.rateLimitPerMinute ? [{ id: `key:${apiKey.id}`, limit: apiKey.rateLimitPerMinute }] : []), + ]; + if (limits.length && request.method === "POST") { + const quota = await consumeQuota(database, limits); + reply.header("X-RateLimit-Limit", Math.min(...limits.map((l) => l.limit))); + reply.header("X-RateLimit-Remaining", quota.remaining); + if (!quota.allowed) { reply.header("Retry-After", 60 - Math.floor((Date.now() % 60_000) / 1_000)); return reply.code(429).send({ error: "Application rate limit exceeded." }); } } request.apiKey = apiKey; request.firewallApp = firewallApp; @@ -436,6 +423,7 @@ export async function buildGateway(config: GatewayConfig): Promise Date.now() ? "open" : circuit.consecutiveFailures > 0 ? "degraded" : "closed", consecutiveFailures: circuit.consecutiveFailures, @@ -512,46 +500,31 @@ export async function buildGateway(config: GatewayConfig): Promise profile.maxInputChars) return reply.code(413).send({ error: "Input is too large." }); - const id = identity.decisionId; - jobs.set(id, { id, appId: firewallApp.id, status: "queued", createdAt: new Date().toISOString() }); try { - const pending = queue.submit(async (queueMs) => { - const current = jobs.get(id); - if (current) current.status = "running"; - return classify(id, envelope, profile, queueMs, firewallApp, apiKey, identity.traceId, identity.requestId); - }, id); - void pending - .then((result) => { - jobs.set(id, { - ...jobs.get(id)!, - status: "complete", - completedAt: new Date().toISOString(), - decision: result.value, - }); - }) - .catch((error: unknown) => { - jobs.set(id, { - ...jobs.get(id)!, - status: "failed", - completedAt: new Date().toISOString(), - error: error instanceof Error ? error.message : "Unknown failure", - }); - }); - return reply.code(202).send({ id, status: "queued", statusUrl: `/v1/jobs/${id}` }); + const job = await jobs.enqueue({ id: identity.decisionId, appId: firewallApp.id, + fingerprint: hash(JSON.stringify(envelope)), idempotencyKey: request.headers["idempotency-key"] as string | undefined, + input: { envelope, profile, firewallApp, apiKey, identity } }); + return reply.code(202).send({ id: job.id, status: job.status, statusUrl: `/v1/jobs/${job.id}` }); } catch (error) { - jobs.delete(id); - if (error instanceof QueueFullError) return reply.code(429).send({ error: error.message }); + if (error instanceof JobCapacityError || error instanceof JobConflict) return reply.code(error.statusCode).send({ error: error.message }); throw error; } }); app.get<{ Params: { id: string } }>("/v1/jobs/:id", { preHandler: requireApiKey }, async (request, reply) => { - const job = jobs.get(request.params.id); + const job = await jobs.get(request.params.id, request.firewallApp!.id); if (!job) return reply.code(404).send({ error: "Job not found or expired." }); - if (job.appId !== request.firewallApp!.id) return reply.code(404).send({ error: "Job not found or expired." }); - return job; + return { id: job.id, appId: job.appId, status: job.status, createdAt: job.createdAt, completedAt: job.completedAt, expiresAt: job.expiresAt, decision: job.result, error: job.error }; }); + const jobWorker = new DurableWorker(jobs, config.queueConcurrency, async (job) => { + const { envelope, profile, firewallApp, apiKey, identity } = jobs.input<{ envelope: ClassificationEnvelope; profile: Profile; firewallApp: AppRecord; apiKey: ApiKeyRecord; identity: ReturnType }>(job); + const existing = await eventsStore.findById(job.id); + if (existing) { const { inputPreview, inputHash, appRulesSnapshot, ...decision } = existing; return decision; } + return classify(job.id, envelope, profile, Date.now() - Date.parse(job.createdAt), firewallApp, apiKey, identity.traceId, identity.requestId); + }, (error) => app.log.error({ err: error }, "Classification worker failed")); + jobWorker.start(); + app.get("/v1/events", { websocket: true }, (socket: WebSocket) => { let authenticatedAppId: string | undefined; let authenticating = false; @@ -594,16 +567,14 @@ export async function buildGateway(config: GatewayConfig): Promise { - const cutoff = Date.now() - 10 * 60_000; - for (const [id, job] of jobs) { - const timestamp = new Date(job.completedAt ?? job.createdAt).getTime(); - if (timestamp < cutoff) jobs.delete(id); - } + void database.prune(new Date(Date.now() - retentionDays * 86_400_000).toISOString()).catch((error) => app.log.error(error, "Retention failed")); }, 60_000); cleanup.unref(); app.addHook("onClose", async () => { clearInterval(cleanup); + await jobWorker.stop(); await deliveryWorker.stop(); await database.close(); }); diff --git a/docs/deployment.md b/docs/deployment.md new file mode 100644 index 0000000..c8f87fc --- /dev/null +++ b/docs/deployment.md @@ -0,0 +1,30 @@ + +## Durable jobs and retention + +`POST /v1/jobs` persists an encrypted input and an immutable policy/application +snapshot before returning 202. Send a unique `Idempotency-Key` (1–128 printable +characters) for a logical submission. Reusing it with different input returns +409; matching retries return the original job for its remaining lifetime. +Workers lease jobs for 30 seconds and renew while running. Crashed workers are +replaced after lease expiry. Provider execution is **at least once**: a crash +after an upstream call can incur another call and charge. Pyro never executes +your tools. Use the decision ID for downstream deduplication. + +Pending inputs expire after 15 minutes and are erased on completion/failure. +Expired pending work becomes an explicit failed result. Results and idempotency +records expire 10 minutes after completion. Encrypt database backups and retain +`CONTROL_PLANE_SECRET` separately; losing it makes queued inputs unreadable. + +The beta queue uses a PostgreSQL document lock across all replicas, is capped at +200 outstanding jobs (or the lower `QUEUE_MAX_DEPTH`), and schedules applications +in turn. This bounded implementation is intended for pilot traffic. Measure +throughput on your hardware before expanding it. Queue age, counts and capacity +are exposed in `/v1/health`. Application and key quotas share database counters; +both limits apply, and job polling does not consume classification quota. + +`EVENT_RETENTION_DAYS` defaults to 30 (minimum 1). The gateway prunes old events +and terminal webhook deliveries every minute. Durable job inputs use the much +shorter deadlines above, independent of `persistInputs`. Labels, caller metadata, +and opt-in previews are event data: never put credentials in them. Review and +evaluation retention are documented with those features. Audit and policy +history are retained until an administrator removes the deployment database. diff --git a/docs/openapi.yaml b/docs/openapi.yaml index ea37b84..1d7c127 100644 --- a/docs/openapi.yaml +++ b/docs/openapi.yaml @@ -2,7 +2,9 @@ openapi: 3.1.0 info: title: Pyro API version: 0.2.0 - description: Typed prompt-injection classification for text, chats, and arbitrary JSON payloads. Each bearer key is scoped to one configured application. + description: Typed prompt-injection classification for text, chats, and + arbitrary JSON payloads. Each bearer key is scoped to one configured + application. servers: - url: http://localhost:8080 security: @@ -32,9 +34,17 @@ paths: schema: $ref: "#/components/schemas/Decision" "401": { description: Invalid API key } - "403": { description: Application is disabled or the key/app cannot use the requested policy } + "403": + { + description: Application is disabled or the key/app cannot use the requested + policy + } "413": { description: Input exceeds the selected profile limit } - "429": { description: Queue capacity or effective application/API-key rate limit has been reached } + "429": + { + description: Queue capacity or effective application/API-key rate limit has been + reached + } x-cli-command: classify x-cli-input: classification parameters: @@ -70,6 +80,17 @@ paths: in: header schema: type: string + - name: Idempotency-Key + in: header + required: false + schema: + type: string + maxLength: 128 + description: Durable encrypted jobs survive gateway restarts. Optional + Idempotency-Key header deduplicates matching requests per application + until the result expires. Pending inputs expire after 15 minutes; + results expire 10 minutes after completion. Upstream execution is at + least once after crashes. /v1/jobs/{id}: get: summary: Read an asynchronous job @@ -93,7 +114,8 @@ paths: /v1/events: get: summary: Open a real-time classification event stream - description: Upgrade to WebSocket, then send `{"type":"auth","apiKey":"..."}` within five seconds. Events are scoped to the API key's application. + description: Upgrade to WebSocket, then send `{"type":"auth","apiKey":"..."}` + within five seconds. Events are scoped to the API key's application. operationId: streamClassificationEvents x-websocket: true responses: @@ -122,7 +144,11 @@ paths: security: [] summary: Prometheus metrics responses: - "200": { description: Prometheus exposition format, content: { text/plain: { schema: { type: string } } } } + "200": + { + description: Prometheus exposition format, + content: { text/plain: { schema: { type: string } } } + } operationId: getMetrics x-cli-command: metrics /: @@ -157,7 +183,8 @@ components: labels: type: object maxProperties: 20 - description: Searchable request context such as a session URL, tenant, or environment. + description: Searchable request context such as a session URL, tenant, or + environment. propertyNames: { pattern: "^[A-Za-z0-9_.-]{1,64}$" } additionalProperties: { type: string, maxLength: 2048 } DetectorResult: @@ -170,10 +197,29 @@ components: weightedProbability: { type: number, minimum: 0, maximum: 1 } Decision: type: object - required: [ id, createdAt, profileId, verdict, action, risk, confidence, reason, detectors, model, provider, latencyMs, queueMs ] + required: + [ + id, + createdAt, + profileId, + verdict, + action, + risk, + confidence, + reason, + detectors, + model, + provider, + latencyMs, + queueMs + ] properties: id: { type: string, description: Server-generated decision identifier } - requestId: { type: string, description: Valid incoming X-Request-Id or a generated request identifier } + requestId: + { + type: string, + description: Valid incoming X-Request-Id or a generated request identifier + } createdAt: { type: string, format: date-time } traceId: { type: string, description: W3C trace identifier } profileId: { type: string } @@ -200,7 +246,8 @@ components: totalMs: { type: number } usage: type: object - description: Provider token usage and actual billed cost when reported. Local decisions report zero cost. + description: Provider token usage and actual billed cost when reported. Local + decisions report zero cost. properties: inputTokens: { type: integer, minimum: 0 } outputTokens: { type: integer, minimum: 0 } diff --git a/packages/storage/src/index.ts b/packages/storage/src/index.ts index 6999d85..d44cb74 100644 --- a/packages/storage/src/index.ts +++ b/packages/storage/src/index.ts @@ -112,6 +112,7 @@ export interface Database { document(key: string, fallback: () => T): DocumentStore; readonly events: EventStore; readonly deliveries: DeliveryStore; + prune(before: string): Promise; ping(): Promise; close(): Promise; } @@ -576,6 +577,11 @@ class PostgresDatabase implements Database { return new PostgresDocument(this.pool, key, fallback); } + async prune(before: string): Promise { + await this.pool.query("DELETE FROM pyro_deliveries WHERE created_at < $1 AND status IN ('delivered', 'failed')", [before]); + await this.pool.query("DELETE FROM pyro_events WHERE created_at < $1", [before]); + } + async ping(): Promise { await this.pool.query("SELECT 1"); } @@ -767,6 +773,9 @@ class MemoryDatabase implements Database { return new MemoryDocument(this.documents, key, fallback); } + async prune(before: string): Promise { + this.eventRows.splice(0, this.eventRows.length, ...this.eventRows.filter((e) => e.createdAt >= before)); + } async ping(): Promise {} async close(): Promise {} } @@ -858,3 +867,5 @@ export function decryptText(value: StoredSecret, secret: string): string { } export * from "./policies.js"; + +export * from "./jobs.js"; diff --git a/packages/storage/src/jobs.test.ts b/packages/storage/src/jobs.test.ts new file mode 100644 index 0000000..08d3776 --- /dev/null +++ b/packages/storage/src/jobs.test.ts @@ -0,0 +1,34 @@ +import assert from "node:assert/strict"; +import test from "node:test"; +import { randomUUID } from "node:crypto"; +import { openDatabase, DurableJobs, consumeQuota } from "./index.js"; +for (const kind of ["memory", "postgres"] as const) test(`${kind}: durable recovery, fencing, idempotency, fairness, encryption and shared quotas`, { skip: kind === "postgres" && !process.env.TEST_DATABASE_URL }, async () => { + const db = await openDatabase(kind === "memory" ? `memory://${randomUUID()}` : process.env.TEST_DATABASE_URL!); + try { + const name = `test_jobs_${randomUUID().replaceAll("-", "")}`; + const first = new DurableJobs(db, "a-secret-longer-than-sixteen", name, 10, 50); + const second = new DurableJobs(db, "a-secret-longer-than-sixteen", name, 10, 50); + const id = randomUUID(); const appId = randomUUID(); + const data = { id, appId, input: { prompt: "sensitive-example" }, fingerprint: "one", idempotencyKey: "one" }; + assert.equal((await first.enqueue(data)).id, (await second.enqueue({ ...data, id: randomUUID() })).id); + await assert.rejects(second.enqueue({ ...data, fingerprint: "different" })); + assert.ok(!JSON.stringify(await db.document(name, () => ({})).read()).includes("sensitive-example")); + const stale = (await first.claim())!; + await new Promise((r) => setTimeout(r, 70)); + const recovered = (await second.claim())!; + assert.equal(recovered.id, id); assert.equal(recovered.attempts, 2); + assert.deepEqual(second.input(recovered), data.input); + assert.equal(await first.finish(stale, { stale: true }), false); + assert.equal(await second.finish(recovered, { ok: true }), true); + assert.equal((await first.get(id, appId))!.input, undefined); + assert.equal(await first.get(id, "other-app"), undefined); + await first.enqueue({ id: "a1", appId: "a", input: 1, fingerprint: "a1" }); + await first.enqueue({ id: "a2", appId: "a", input: 2, fingerprint: "a2" }); + await first.enqueue({ id: "b1", appId: "b", input: 3, fingerprint: "b1" }); + const a = (await first.claim())!, b = (await second.claim())!; + assert.notEqual(a.appId, b.appId, "fair scheduling gives each application a turn"); + const limits = [{ id: `app:${appId}`, limit: 3 }]; + const quota = await Promise.all(Array.from({ length: 10 }, () => consumeQuota(db, limits))); + assert.equal(quota.filter((q) => q.allowed).length, 3); + } finally { await db.close(); } +}); diff --git a/packages/storage/src/jobs.ts b/packages/storage/src/jobs.ts new file mode 100644 index 0000000..91d3d67 --- /dev/null +++ b/packages/storage/src/jobs.ts @@ -0,0 +1,123 @@ +import { randomUUID } from "node:crypto"; +import type { StoredSecret } from "@pyro/contracts"; +import { decryptText, encryptText, type Database, type DocumentStore } from "./index.js"; + +export interface DurableJob { + id: string; appId: string; status: "queued" | "running" | "complete" | "failed" | "cancelled"; + createdAt: string; expiresAt: string; completedAt?: string; + input?: StoredSecret; fingerprint: string; idempotencyKey?: string; + leaseToken?: string; leaseUntil?: number; attempts: number; + result?: unknown; error?: string; +} +interface QueueState { jobs: DurableJob[]; served: Record } +export class JobConflict extends Error { readonly statusCode = 409; } +export class JobCapacityError extends Error { readonly statusCode = 429; } + +// A bounded queue uses the same transactional document lock in every replica. +// This beta implementation deliberately favors correctness over high throughput. +export class DurableJobs { + private readonly store: DocumentStore; + constructor(database: Database, readonly secret: string, name = "classification_jobs", readonly maxDepth = 200, readonly leaseMs = 30_000) { + this.store = database.document(name, () => ({ jobs: [], served: {} })); + } + async enqueue(job: { id: string; appId: string; input: unknown; fingerprint: string; idempotencyKey?: string }): Promise { + if (job.idempotencyKey && (job.idempotencyKey.length > 128 || !/^[\x21-\x7e]+$/.test(job.idempotencyKey))) throw new JobConflict("Idempotency-Key must contain 1–128 printable characters."); + let selected!: DurableJob; + await this.store.update((state) => { + this.expire(state); + const prior = job.idempotencyKey ? state.jobs.find((j) => j.appId === job.appId && j.idempotencyKey === job.idempotencyKey) : undefined; + if (prior) { + if (prior.fingerprint !== job.fingerprint) throw new JobConflict("Idempotency-Key was already used with a different request."); + selected = prior; return state; + } + if (state.jobs.filter((j) => j.status === "queued" || j.status === "running").length >= this.maxDepth) throw new JobCapacityError("Durable queue is full. Retry later with the same idempotency key."); + selected = { id: job.id, appId: job.appId, fingerprint: job.fingerprint, idempotencyKey: job.idempotencyKey, input: encryptText(JSON.stringify(job.input), this.secret), status: "queued", createdAt: new Date().toISOString(), expiresAt: new Date(Date.now() + 15 * 60_000).toISOString(), attempts: 0 }; + state.jobs.push(selected); return state; + }); + return selected; + } + private expire(state: QueueState, now = Date.now()) { + state.jobs = state.jobs.filter((j) => !j.completedAt || Date.parse(j.expiresAt) > now); + for (const job of state.jobs) if (!job.completedAt && Date.parse(job.expiresAt) <= now) { + job.status = "failed"; job.error = "Job deadline exceeded."; job.completedAt = new Date(now).toISOString(); + job.expiresAt = new Date(now + 10 * 60_000).toISOString(); delete job.input; delete job.leaseToken; + } + const apps = new Set(state.jobs.map((j) => j.appId)); + for (const id of Object.keys(state.served)) if (!apps.has(id)) delete state.served[id]; + } + async claim(now = Date.now()): Promise { + let claimed: DurableJob | undefined; + await this.store.update((state) => { + this.expire(state, now); + const due = state.jobs.filter((j) => j.status === "queued" || j.status === "running" && (j.leaseUntil ?? 0) <= now); + due.sort((a, b) => (state.served[a.appId] ?? 0) - (state.served[b.appId] ?? 0) || a.createdAt.localeCompare(b.createdAt)); + const job = due[0]; + if (job) { job.status = "running"; job.leaseToken = randomUUID(); job.leaseUntil = now + this.leaseMs; job.attempts++; state.served[job.appId] = now; claimed = structuredClone(job); } + return state; + }); + return claimed; + } + input(job: DurableJob): T { return JSON.parse(decryptText(job.input!, this.secret)) as T; } + async heartbeat(job: DurableJob): Promise { + await this.store.update((state) => { const current = state.jobs.find((j) => j.id === job.id && j.status === "running" && j.leaseToken === job.leaseToken); if (current && (current.leaseUntil ?? 0) > Date.now()) current.leaseUntil = Date.now() + this.leaseMs; return state; }); + } + async finish(job: DurableJob, result?: unknown, error?: string): Promise { + let applied = false; + await this.store.update((state) => { + const current = state.jobs.find((j) => j.id === job.id && j.status === "running" && j.leaseToken === job.leaseToken && (j.leaseUntil ?? 0) > Date.now()); + if (current) { current.status = error ? "failed" : "complete"; current.result = result; current.error = error; current.completedAt = new Date().toISOString(); current.expiresAt = new Date(Date.now() + 10 * 60_000).toISOString(); delete current.input; delete current.leaseToken; delete current.leaseUntil; applied = true; } + return state; + }); + return applied; + } + async get(id: string, appId: string): Promise { + const state = await this.store.update((s) => { this.expire(s); return s; }); + return state.jobs.find((j) => j.id === id && j.appId === appId); + } + async stats() { + const state = await this.store.update((s) => { this.expire(s); return s; }); + const pending = state.jobs.filter((j) => !j.completedAt); + return { queued: pending.filter((j) => j.status === "queued").length, running: pending.filter((j) => j.status === "running").length, failed: state.jobs.filter((j) => j.status === "failed").length, oldestAgeMs: pending.length ? Date.now() - Math.min(...pending.map((j) => Date.parse(j.createdAt))) : 0, capacity: this.maxDepth }; + } +} + +export class DurableWorker { + private timer?: ReturnType; + private stopped = false; + private ticking = false; + private readonly active = new Set>(); + constructor(private readonly jobs: DurableJobs, private readonly concurrency: number, private readonly handler: (job: DurableJob) => Promise, private readonly onError: (error: unknown) => void) {} + start() { this.timer = setInterval(() => void this.tick(), 100); this.timer.unref(); } + private async tick() { + if (this.ticking || this.stopped) return; + this.ticking = true; + try { + while (!this.stopped && this.active.size < this.concurrency) { + const job = await this.jobs.claim(); if (!job) break; + const pending = this.run(job); this.active.add(pending); void pending.finally(() => this.active.delete(pending)); + } + } catch (e) { this.onError(e); } finally { this.ticking = false; } + } + private async run(job: DurableJob) { + const heartbeat = setInterval(() => void this.jobs.heartbeat(job).catch(this.onError), Math.max(10, this.jobs.leaseMs / 3)); + heartbeat.unref(); + try { await this.jobs.finish(job, await this.handler(job)); } + catch (e) { try { await this.jobs.finish(job, undefined, e instanceof Error ? e.message : "Job failed."); } catch (storageError) { this.onError(storageError); } } + finally { clearInterval(heartbeat); } + } + async stop() { this.stopped = true; clearInterval(this.timer); while (this.ticking) await new Promise((r) => setTimeout(r, 10)); await Promise.allSettled([...this.active]); } +} + +export async function consumeQuota(database: Database, limits: Array<{ id: string; limit: number }>, now = Date.now()): Promise<{ allowed: boolean; remaining: number }> { + const store = database.document>("rate_limits", () => ({})); + const window = Math.floor(now / 60_000); + let remaining = Infinity; + let allowed = true; + await store.update((state) => { + for (const [id, row] of Object.entries(state)) if (row.window !== window) delete state[id]; + for (const { id, limit } of limits) if ((state[id]?.count ?? 0) >= limit) allowed = false; + for (const { id, limit } of limits) { const row = state[id] ?? { window, count: 0 }; if (allowed) row.count++; state[id] = row; remaining = Math.min(remaining, Math.max(0, limit - row.count)); } + return state; + }); + return { allowed, remaining }; +}