diff --git a/.changeset/cooperative-rate-limiting.md b/.changeset/cooperative-rate-limiting.md new file mode 100644 index 000000000..524a58109 --- /dev/null +++ b/.changeset/cooperative-rate-limiting.md @@ -0,0 +1,11 @@ +--- +"@bradygaster/squad-sdk": minor +--- + +feat: Cooperative rate limiting with predictive circuit breaker + +Added cooperative rate limiting patterns for multi-agent deployments: +- Traffic Light, Predictive Circuit Breaker, Priority Retry Windows +- Cooperative Token Pool for shared quota management + +Closes #515 diff --git a/.squad-templates/cooperative-rate-limiting.md b/.squad-templates/cooperative-rate-limiting.md new file mode 100644 index 000000000..bf56ef122 --- /dev/null +++ b/.squad-templates/cooperative-rate-limiting.md @@ -0,0 +1,229 @@ +# Cooperative Rate Limiting for Multi-Agent Deployments + +> Coordinate API quota across multiple Ralph instances to prevent cascading failures. + +## Problem + +The [circuit breaker template](ralph-circuit-breaker.md) handles single-instance rate limiting well. But when multiple Ralphs run across machines (or pods on K8s), each instance independently hits API limits: + +- **No coordination** — 5 Ralphs each think they have full API quota +- **Thundering herd** — All Ralphs retry simultaneously after rate limit resets +- **Priority inversion** — Low-priority work exhausts quota before critical work runs +- **Reactive only** — Circuit opens AFTER 429, wasting the failed request + +## Solution: 6-Pattern Architecture + +These patterns layer on top of the existing circuit breaker. Each is independent — adopt one or all. + +### Pattern 1: Traffic Light (RAAS — Rate-Aware Agent Scheduling) + +Map GitHub API `X-RateLimit-Remaining` to traffic light states: + +| State | Remaining % | Behavior | +|-------|------------|----------| +| 🟢 GREEN | >20% | Normal operation | +| 🟡 AMBER | 5–20% | Only P0 agents proceed | +| 🔴 RED | <5% | Block all except emergency P0 | + +```typescript +type TrafficLight = 'green' | 'amber' | 'red'; + +function getTrafficLight(remaining: number, limit: number): TrafficLight { + const pct = remaining / limit; + if (pct > 0.20) return 'green'; + if (pct > 0.05) return 'amber'; + return 'red'; +} + +function shouldProceed(light: TrafficLight, agentPriority: number): boolean { + if (light === 'green') return true; + if (light === 'amber') return agentPriority === 0; // P0 only + return false; // RED — block all +} +``` + +### Pattern 2: Cooperative Token Pool (CMARP) + +A shared JSON file (`~/.squad/rate-pool.json`) distributes API quota: + +```json +{ + "totalLimit": 5000, + "resetAt": "2026-03-22T20:00:00Z", + "allocations": { + "picard": { "priority": 0, "allocated": 2000, "used": 450, "leaseExpiry": "2026-03-22T19:55:00Z" }, + "data": { "priority": 1, "allocated": 1750, "used": 200, "leaseExpiry": "2026-03-22T19:55:00Z" }, + "ralph": { "priority": 2, "allocated": 1250, "used": 100, "leaseExpiry": "2026-03-22T19:55:00Z" } + } +} +``` + +**Rules:** +- P0 agents (Lead) get 40% of quota +- P1 agents (specialists) get 35% +- P2 agents (Ralph, Scribe) get 25% +- Stale leases (>5 minutes without heartbeat) are auto-recovered +- Each agent checks their remaining allocation before making API calls + +```typescript +interface RatePoolAllocation { + priority: number; + allocated: number; + used: number; + leaseExpiry: string; +} + +interface RatePool { + totalLimit: number; + resetAt: string; + allocations: Record; +} + +function canUseQuota(pool: RatePool, agentName: string): boolean { + const alloc = pool.allocations[agentName]; + if (!alloc) return true; // Unknown agent — allow (graceful) + + // Reclaim stale leases from crashed agents + const now = new Date(); + for (const [name, a] of Object.entries(pool.allocations)) { + if (new Date(a.leaseExpiry) < now && name !== agentName) { + a.allocated = 0; // Reclaim + } + } + + return alloc.used < alloc.allocated; +} +``` + +### Pattern 3: Predictive Circuit Breaker (PCB) + +Opens the circuit BEFORE getting a 429 by predicting when quota will run out: + +```typescript +interface RateSample { + timestamp: number; // Date.now() + remaining: number; // from X-RateLimit-Remaining header +} + +class PredictiveCircuitBreaker { + private samples: RateSample[] = []; + private readonly maxSamples = 10; + private readonly warningThresholdSeconds = 120; + + addSample(remaining: number): void { + this.samples.push({ timestamp: Date.now(), remaining }); + if (this.samples.length > this.maxSamples) { + this.samples.shift(); + } + } + + /** Predict seconds until quota exhaustion using linear regression */ + predictExhaustion(): number | null { + if (this.samples.length < 3) return null; + + const n = this.samples.length; + const first = this.samples[0]; + const last = this.samples[n - 1]; + + const elapsedMs = last.timestamp - first.timestamp; + if (elapsedMs === 0) return null; + + const consumedPerMs = (first.remaining - last.remaining) / elapsedMs; + if (consumedPerMs <= 0) return null; // Not consuming — safe + + const msUntilExhausted = last.remaining / consumedPerMs; + return msUntilExhausted / 1000; + } + + shouldOpen(): boolean { + const eta = this.predictExhaustion(); + if (eta === null) return false; + return eta < this.warningThresholdSeconds; + } +} +``` + +### Pattern 4: Priority Retry Windows (PWJG) + +Non-overlapping jitter windows prevent thundering herd: + +| Priority | Retry Window | Description | +|----------|-------------|-------------| +| P0 (Lead) | 500ms–5s | Recovers first | +| P1 (Specialists) | 2s–30s | Moderate delay | +| P2 (Ralph/Scribe) | 5s–60s | Most patient | + +```typescript +function getRetryDelay(priority: number, attempt: number): number { + const windows: Record = { + 0: [500, 5000], // P0: 500ms–5s + 1: [2000, 30000], // P1: 2s–30s + 2: [5000, 60000], // P2: 5s–60s + }; + + const [min, max] = windows[priority] ?? windows[2]; + const base = Math.min(min * Math.pow(2, attempt), max); + const jitter = Math.random() * base * 0.5; + return base + jitter; +} +``` + +### Pattern 5: Resource Epoch Tracker (RET) + +Heartbeat-based lease system for multi-machine deployments: + +```typescript +interface ResourceLease { + agent: string; + machine: string; + leaseStart: string; + leaseExpiry: string; // Typically 5 minutes from now + allocated: number; +} + +// Each agent renews its lease every 2 minutes +// If lease expires (agent crashed), allocation is reclaimed +``` + +### Pattern 6: Cascade Dependency Detector (CDD) + +Track downstream failures and apply backpressure: + +``` +Agent A (rate limited) → Agent B (waiting for A) → Agent C (waiting for B) + ↑ Backpressure signal: "don't start new work" +``` + +When a dependency is rate-limited, upstream agents should pause new work rather than queuing requests that will fail. + +## Kubernetes Integration + +On K8s, cooperative rate limiting can use KEDA to scale pods based on API quota: + +```yaml +apiVersion: keda.sh/v1alpha1 +kind: ScaledObject +spec: + scaleTargetRef: + name: ralph-deployment + triggers: + - type: external + metadata: + scalerAddress: keda-copilot-scaler:6000 + # Scaler returns 0 when rate limited → pods scale to zero +``` + +See [keda-copilot-scaler](https://github.com/tamirdresher/keda-copilot-scaler) for a complete implementation. + +## Quick Start + +1. **Minimum viable:** Adopt Pattern 1 (Traffic Light) — read `X-RateLimit-Remaining` from API responses +2. **Multi-machine:** Add Pattern 2 (Cooperative Pool) — shared `rate-pool.json` +3. **Production:** Add Pattern 3 (Predictive CB) — prevent 429s entirely +4. **Kubernetes:** Add KEDA scaler for automatic pod scaling + +## References + +- [Circuit Breaker Template](ralph-circuit-breaker.md) — Foundation patterns +- [Squad on AKS](https://github.com/tamirdresher/squad-on-aks) — Production K8s deployment +- [KEDA Copilot Scaler](https://github.com/tamirdresher/keda-copilot-scaler) — Custom KEDA external scaler diff --git a/packages/squad-sdk/package.json b/packages/squad-sdk/package.json index 5fc8902e3..bb50a2aa2 100644 --- a/packages/squad-sdk/package.json +++ b/packages/squad-sdk/package.json @@ -78,6 +78,10 @@ "types": "./dist/ralph/triage.d.ts", "import": "./dist/ralph/triage.js" }, + "./ralph/rate-limiting": { + "types": "./dist/ralph/rate-limiting.d.ts", + "import": "./dist/ralph/rate-limiting.js" + }, "./casting": { "types": "./dist/casting/index.d.ts", "import": "./dist/casting/index.js" diff --git a/packages/squad-sdk/src/ralph/index.ts b/packages/squad-sdk/src/ralph/index.ts index f4ffa692b..cb946b201 100644 --- a/packages/squad-sdk/src/ralph/index.ts +++ b/packages/squad-sdk/src/ralph/index.ts @@ -183,3 +183,5 @@ export class RalphMonitor { this.eventBus = null; } } + +export { getTrafficLight, shouldProceed, getRetryDelay, PredictiveCircuitBreaker, canUseQuota, loadRatePool, type RatePool, type RatePoolAllocation, type RateSample, type TrafficLight, type AgentPriority } from './rate-limiting.js'; diff --git a/packages/squad-sdk/src/ralph/rate-limiting.ts b/packages/squad-sdk/src/ralph/rate-limiting.ts new file mode 100644 index 000000000..be97c9bca --- /dev/null +++ b/packages/squad-sdk/src/ralph/rate-limiting.ts @@ -0,0 +1,209 @@ +/** + * Predictive Circuit Breaker — Rate Limit Protection + * + * Opens the circuit BEFORE getting a 429 by predicting when + * API quota will be exhausted. Prediction uses the observed + * consumption rate between the oldest and newest samples + * (first-to-last delta), not a full least-squares regression. + * + * @see https://github.com/bradygaster/squad/issues/515 + */ + +import { readFile } from 'node:fs/promises'; +import { existsSync } from 'node:fs'; +import path from 'node:path'; +import os from 'node:os'; + +/** A rate limit sample from API response headers */ +export interface RateSample { + timestamp: number; // Date.now() + remaining: number; // from X-RateLimit-Remaining header + limit: number; // from X-RateLimit-Limit header +} + +/** Traffic light state for rate-aware scheduling */ +export type TrafficLight = 'green' | 'amber' | 'red'; + +/** Agent priority for quota allocation */ +export type AgentPriority = 0 | 1 | 2; + +/** Priority-based retry windows (ms) */ +const RETRY_WINDOWS: Record = { + 0: [500, 5_000], // P0 (Lead): 500ms–5s + 1: [2_000, 30_000], // P1 (Specialists): 2s–30s + 2: [5_000, 60_000], // P2 (Ralph/Scribe): 5s–60s +}; + +/** + * Determine traffic light from rate limit headers. + */ +export function getTrafficLight(remaining: number, limit: number): TrafficLight { + if (limit === 0) return 'red'; + const pct = remaining / limit; + if (pct > 0.20) return 'green'; + if (pct > 0.05) return 'amber'; + return 'red'; +} + +/** + * Check if an agent should proceed based on traffic light and priority. + * - GREEN: all agents proceed + * - AMBER: only P0 agents proceed + * - RED: no agents proceed + */ +export function shouldProceed(light: TrafficLight, priority: AgentPriority): boolean { + if (light === 'green') return true; + if (light === 'amber') return priority === 0; + return false; +} + +/** + * Get retry delay with priority-based jitter windows. + * Higher priority agents retry sooner with smaller windows. + */ +export function getRetryDelay(priority: AgentPriority, attempt: number): number { + const [min, max] = RETRY_WINDOWS[priority] ?? RETRY_WINDOWS[2]; + const base = Math.min(min * Math.pow(2, attempt), max); + const jitter = Math.random() * base * 0.5; + return Math.round(base + jitter); +} + +/** + * Predictive circuit breaker that opens BEFORE rate limit errors. + * + * Tracks the last N rate limit samples and uses linear regression + * to predict when quota will be exhausted. If predicted ETA is + * below the warning threshold, the circuit opens preemptively. + */ +export class PredictiveCircuitBreaker { + private samples: RateSample[] = []; + private readonly maxSamples: number; + private readonly warningThresholdSeconds: number; + + constructor(options?: { maxSamples?: number; warningThresholdSeconds?: number }) { + this.maxSamples = options?.maxSamples ?? 10; + this.warningThresholdSeconds = options?.warningThresholdSeconds ?? 120; + } + + /** Record a rate limit sample from API response headers */ + addSample(remaining: number, limit: number): void { + this.samples.push({ timestamp: Date.now(), remaining, limit }); + if (this.samples.length > this.maxSamples) { + this.samples.shift(); + } + } + + /** Get all recorded samples (for testing/debugging) */ + getSamples(): readonly RateSample[] { + return this.samples; + } + + /** + * Predict seconds until quota exhaustion using the observed + * consumption rate between the first and last recorded samples. + * Returns null if insufficient data or quota is not being consumed. + */ + predictExhaustion(): number | null { + if (this.samples.length < 3) return null; + + const n = this.samples.length; + const first = this.samples[0]!; + const last = this.samples[n - 1]!; + + const elapsedMs = last.timestamp - first.timestamp; + if (elapsedMs === 0) return null; + + const consumed = first.remaining - last.remaining; + if (consumed <= 0) return null; // Not consuming or recovering + + const consumedPerMs = consumed / elapsedMs; + const msUntilExhausted = last.remaining / consumedPerMs; + return msUntilExhausted / 1000; + } + + /** + * Should the circuit open preemptively? + * Returns true when predicted ETA to exhaustion is below threshold. + */ + shouldOpen(): boolean { + const eta = this.predictExhaustion(); + if (eta === null) return false; + return eta < this.warningThresholdSeconds; + } + + /** Reset all samples (e.g., after rate limit window resets) */ + reset(): void { + this.samples = []; + } +} + +/** Rate pool allocation for cooperative multi-agent quota management */ +export interface RatePoolAllocation { + priority: AgentPriority; + allocated: number; + used: number; + leaseExpiry: string; +} + +/** Shared rate pool state */ +export interface RatePool { + totalLimit: number; + resetAt: string; + allocations: Record; +} + +/** + * Check if an agent has remaining quota in the cooperative pool. + * Pure read — no side effects. + */ +export function canUseQuota(pool: RatePool, agentName: string): boolean { + const alloc = pool.allocations[agentName]; + if (!alloc) return true; // Unknown agent — allow gracefully + + return alloc.used < alloc.allocated; +} + +/** + * Consume one unit of quota for an agent. + * Also reclaims stale leases from other crashed agents so their + * unused allocation is freed for the pool. + * Call this after canUseQuota() confirms there is quota available. + */ +export function consumeQuota(pool: RatePool, agentName: string): void { + const alloc = pool.allocations[agentName]; + if (!alloc) return; // Unknown agent — nothing to consume + + // Reclaim stale leases from other agents + const now = new Date(); + for (const [name, a] of Object.entries(pool.allocations)) { + if (new Date(a.leaseExpiry) < now && name !== agentName) { + a.allocated = 0; + } + } + + alloc.used++; +} + +/** + * Load rate pool state from the shared file. + */ +export async function loadRatePool(teamRoot?: string): Promise { + const candidates: string[] = []; + + if (teamRoot) { + candidates.push(path.join(teamRoot, '.squad', 'rate-pool.json')); + } + candidates.push(path.join(os.homedir(), '.squad', 'rate-pool.json')); + + for (const candidate of candidates) { + if (existsSync(candidate)) { + try { + const raw = await readFile(candidate, 'utf8'); + return JSON.parse(raw) as RatePool; + } catch { + // Malformed — skip + } + } + } + return null; +} diff --git a/test/rate-limiting.test.ts b/test/rate-limiting.test.ts new file mode 100644 index 000000000..78b472c83 --- /dev/null +++ b/test/rate-limiting.test.ts @@ -0,0 +1,279 @@ +import { describe, it, expect, beforeEach } from 'vitest'; +import { + getTrafficLight, + shouldProceed, + getRetryDelay, + PredictiveCircuitBreaker, + canUseQuota, + consumeQuota, + type RatePool, + type AgentPriority, +} from '@bradygaster/squad-sdk/ralph/rate-limiting'; + +describe('getTrafficLight', () => { + it('returns green when >20% remaining', () => { + expect(getTrafficLight(1500, 5000)).toBe('green'); + expect(getTrafficLight(1001, 5000)).toBe('green'); + }); + + it('returns amber when 5-20% remaining', () => { + expect(getTrafficLight(1000, 5000)).toBe('amber'); + expect(getTrafficLight(251, 5000)).toBe('amber'); + }); + + it('returns red when <5% remaining', () => { + expect(getTrafficLight(250, 5000)).toBe('red'); + expect(getTrafficLight(0, 5000)).toBe('red'); + }); + + it('returns red when limit is 0', () => { + expect(getTrafficLight(0, 0)).toBe('red'); + }); +}); + +describe('shouldProceed', () => { + it('allows all agents on green', () => { + expect(shouldProceed('green', 0)).toBe(true); + expect(shouldProceed('green', 1)).toBe(true); + expect(shouldProceed('green', 2)).toBe(true); + }); + + it('allows only P0 on amber', () => { + expect(shouldProceed('amber', 0)).toBe(true); + expect(shouldProceed('amber', 1)).toBe(false); + expect(shouldProceed('amber', 2)).toBe(false); + }); + + it('blocks all on red', () => { + expect(shouldProceed('red', 0)).toBe(false); + expect(shouldProceed('red', 1)).toBe(false); + expect(shouldProceed('red', 2)).toBe(false); + }); +}); + +describe('getRetryDelay', () => { + it('returns delay within P0 window', () => { + const delay = getRetryDelay(0, 0); + expect(delay).toBeGreaterThanOrEqual(500); + expect(delay).toBeLessThanOrEqual(7500); // 5000 + 50% jitter + }); + + it('returns delay within P2 window', () => { + const delay = getRetryDelay(2, 0); + expect(delay).toBeGreaterThanOrEqual(5000); + expect(delay).toBeLessThanOrEqual(90000); // 60000 + 50% jitter + }); + + it('increases with attempt count', () => { + const d0 = getRetryDelay(1, 0); + const d3 = getRetryDelay(1, 3); + // d3 should generally be larger (exponential backoff) + // But jitter makes this probabilistic, so we just check it runs + expect(d0).toBeGreaterThan(0); + expect(d3).toBeGreaterThan(0); + }); + + it('caps at max window', () => { + const delay = getRetryDelay(0, 100); // Very high attempt + expect(delay).toBeLessThanOrEqual(7500); // 5000 + 50% jitter max + }); +}); + +describe('PredictiveCircuitBreaker', () => { + let cb: PredictiveCircuitBreaker; + + beforeEach(() => { + cb = new PredictiveCircuitBreaker({ maxSamples: 5, warningThresholdSeconds: 120 }); + }); + + it('does not open with insufficient samples', () => { + cb.addSample(5000, 5000); + cb.addSample(4900, 5000); + expect(cb.shouldOpen()).toBe(false); // Only 2 samples + }); + + it('does not open when quota is not being consumed', () => { + cb.addSample(4000, 5000); + cb.addSample(4000, 5000); + cb.addSample(4000, 5000); + expect(cb.shouldOpen()).toBe(false); + }); + + it('opens when quota is being consumed rapidly', () => { + // Simulate rapid consumption: 1000 → 100 in ~60 seconds + const now = Date.now(); + cb = new PredictiveCircuitBreaker({ warningThresholdSeconds: 120 }); + // Manually set samples with timestamps + (cb as any).samples = [ + { timestamp: now - 60000, remaining: 1000, limit: 5000 }, + { timestamp: now - 40000, remaining: 700, limit: 5000 }, + { timestamp: now - 20000, remaining: 400, limit: 5000 }, + { timestamp: now, remaining: 100, limit: 5000 }, + ]; + + const eta = cb.predictExhaustion(); + expect(eta).not.toBeNull(); + expect(eta!).toBeLessThan(120); // Should predict exhaustion within 120s + expect(cb.shouldOpen()).toBe(true); + }); + + it('does not open when consumption rate is slow', () => { + const now = Date.now(); + (cb as any).samples = [ + { timestamp: now - 3600000, remaining: 5000, limit: 5000 }, // 1 hour ago + { timestamp: now - 1800000, remaining: 4900, limit: 5000 }, + { timestamp: now, remaining: 4800, limit: 5000 }, + ]; + + const eta = cb.predictExhaustion(); + expect(eta).not.toBeNull(); + expect(eta!).toBeGreaterThan(120); + expect(cb.shouldOpen()).toBe(false); + }); + + it('resets samples', () => { + cb.addSample(1000, 5000); + cb.addSample(500, 5000); + cb.addSample(100, 5000); + cb.reset(); + expect(cb.getSamples()).toHaveLength(0); + expect(cb.predictExhaustion()).toBeNull(); + }); + + it('limits to maxSamples', () => { + for (let i = 0; i < 10; i++) { + cb.addSample(5000 - i * 100, 5000); + } + expect(cb.getSamples()).toHaveLength(5); // maxSamples = 5 + }); +}); + +describe('canUseQuota', () => { + const futureExpiry = new Date(Date.now() + 300000).toISOString(); // 5 min from now + const pastExpiry = new Date(Date.now() - 60000).toISOString(); // 1 min ago + + it('allows unknown agents', () => { + const pool: RatePool = { + totalLimit: 5000, + resetAt: futureExpiry, + allocations: {}, + }; + expect(canUseQuota(pool, 'unknown-agent')).toBe(true); + }); + + it('allows agent with remaining quota', () => { + const pool: RatePool = { + totalLimit: 5000, + resetAt: futureExpiry, + allocations: { + picard: { priority: 0, allocated: 2000, used: 500, leaseExpiry: futureExpiry }, + }, + }; + expect(canUseQuota(pool, 'picard')).toBe(true); + }); + + it('blocks agent that exhausted quota', () => { + const pool: RatePool = { + totalLimit: 5000, + resetAt: futureExpiry, + allocations: { + ralph: { priority: 2, allocated: 1000, used: 1000, leaseExpiry: futureExpiry }, + }, + }; + expect(canUseQuota(pool, 'ralph')).toBe(false); + }); + + it('is pure — does not reclaim stale leases from other agents', () => { + const pool: RatePool = { + totalLimit: 5000, + resetAt: futureExpiry, + allocations: { + picard: { priority: 0, allocated: 2000, used: 500, leaseExpiry: pastExpiry }, // Stale + ralph: { priority: 2, allocated: 1000, used: 500, leaseExpiry: futureExpiry }, + }, + }; + canUseQuota(pool, 'ralph'); + expect(pool.allocations.picard!.allocated).toBe(2000); // NOT reclaimed — pure read + }); + + it('is pure — does not modify own stale lease', () => { + const pool: RatePool = { + totalLimit: 5000, + resetAt: futureExpiry, + allocations: { + ralph: { priority: 2, allocated: 1000, used: 500, leaseExpiry: pastExpiry }, + }, + }; + expect(canUseQuota(pool, 'ralph')).toBe(true); + expect(pool.allocations.ralph!.allocated).toBe(1000); // Unchanged + expect(pool.allocations.ralph!.used).toBe(500); // Unchanged + }); +}); + +describe('consumeQuota', () => { + const futureExpiry = new Date(Date.now() + 300000).toISOString(); + const pastExpiry = new Date(Date.now() - 60000).toISOString(); + + it('increments used count for the calling agent', () => { + const pool: RatePool = { + totalLimit: 5000, + resetAt: futureExpiry, + allocations: { + ralph: { priority: 2, allocated: 1000, used: 5, leaseExpiry: futureExpiry }, + }, + }; + consumeQuota(pool, 'ralph'); + expect(pool.allocations.ralph!.used).toBe(6); + }); + + it('reclaims stale leases from other agents', () => { + const pool: RatePool = { + totalLimit: 5000, + resetAt: futureExpiry, + allocations: { + picard: { priority: 0, allocated: 2000, used: 500, leaseExpiry: pastExpiry }, // Stale + ralph: { priority: 2, allocated: 1000, used: 500, leaseExpiry: futureExpiry }, + }, + }; + consumeQuota(pool, 'ralph'); + expect(pool.allocations.picard!.allocated).toBe(0); // Reclaimed + expect(pool.allocations.ralph!.used).toBe(501); // Incremented + }); + + it('does not reclaim own stale lease', () => { + const pool: RatePool = { + totalLimit: 5000, + resetAt: futureExpiry, + allocations: { + ralph: { priority: 2, allocated: 1000, used: 500, leaseExpiry: pastExpiry }, + }, + }; + consumeQuota(pool, 'ralph'); + expect(pool.allocations.ralph!.allocated).toBe(1000); // Not reclaimed + expect(pool.allocations.ralph!.used).toBe(501); // Incremented + }); + + it('is a no-op for unknown agents', () => { + const pool: RatePool = { + totalLimit: 5000, + resetAt: futureExpiry, + allocations: {}, + }; + expect(() => consumeQuota(pool, 'ghost')).not.toThrow(); + }); + + it('canUseQuota + consumeQuota work correctly together', () => { + const pool: RatePool = { + totalLimit: 5000, + resetAt: futureExpiry, + allocations: { + ralph: { priority: 2, allocated: 1000, used: 999, leaseExpiry: futureExpiry }, + }, + }; + // One quota left + expect(canUseQuota(pool, 'ralph')).toBe(true); + consumeQuota(pool, 'ralph'); + // Now exhausted + expect(canUseQuota(pool, 'ralph')).toBe(false); + }); +});