Skip to content

feat: implement AI request budget tracking, token usage limits, concu…#326

Open
A-one-tech wants to merge 1 commit into
Toolbox-Lab:mainfrom
A-one-tech:feature/ai-request-dedup-budget-cache
Open

feat: implement AI request budget tracking, token usage limits, concu…#326
A-one-tech wants to merge 1 commit into
Toolbox-Lab:mainfrom
A-one-tech:feature/ai-request-dedup-budget-cache

Conversation

@A-one-tech

Copy link
Copy Markdown

#Closes
#833

What & Why

A burst of identical POST /api/ai/explain requests — from multiple browser tabs,
client retries, or automated test harnesses — previously triggered independent LLM
calls, wasting tokens and money. There was also no mechanism to cap total spend,
limit concurrency, or invalidate stale explanations when prompts change.

This PR introduces three interlocking subsystems that close all of those gaps:

Subsystem What it does
Request deduplication Identical requests within a configurable window are coalesced — only one LLM call is made, all waiters receive the same result
Token budget enforcement Pre-flight checks against hourly/daily counters and a concurrency semaphore block expensive requests before they reach the API
Cache invalidation AI responses are Redis-cached with TTL expiry, prompt-version tagging (automatic on prompt change), and explicit per-tx / bulk invalidation endpoints

Request Flow

POST /api/ai/explain
        │
        ▼
┌──────────────┐  HIT  ┌──────────────────┐
│ Cache lookup │──────▶│ Return (cached)  │
└──────────────┘       └──────────────────┘
        │ MISS
        ▼
┌──────────────────┐  LOCKED  ┌────────────────────┐
│ Dedup lock       │─────────▶│ waitForResult()    │
│ SET NX PX {ttl} │          │ poll → return result│
└──────────────────┘          └────────────────────┘
        │ ACQUIRED
        ▼
┌──────────────────┐
│ Budget pre-flight│──EXCEEDED──▶ HTTP 429 + reset_at
│ hourly / daily   │
└──────────────────┘
        │ OK
        ▼
┌──────────────────┐
│ Concurrency gate │──AT LIMIT──▶ HTTP 429
│ INCR semaphore   │
└──────────────────┘
        │ OK
        ▼
┌──────────────────┐
│  LLM API call    │  (any OpenAI-compatible endpoint)
└──────────────────┘
        │
        ▼
┌────────────────────────────────────┐
│ Record usage · Cache · Notify dedup│
└────────────────────────────────────┘
        │
        ▼
    AIExplanation { cached: false }

Files Changed

New — apps/server/src/redis/

client.ts — Singleton ioredis client with typed helpers used throughout the new modules.

Helper Redis op
cacheSet(key, val, ttl?) SET [EX]
cacheGet(key) GET
cacheDel(...keys) DEL
acquireLock(key, val, ttlMs) SET NX PX
incrBy / decrBy INCRBY / DECRBY
getNumber(key) GET → number (0 on miss)
setExpiry(key, ttl) EXPIRE
scanKeys(pattern) SCAN loop

New — apps/server/src/dedup/

dedup.ts — SHA-256 fingerprinting and lock-based request coalescing.

fingerprint = SHA-256( txHash.lower + ":" + network.lower + ":" + requestType.lower )

Lock key  : grat:dedup:{fingerprint}   SET NX PX {dedupWindowMs}
Result key: grat:result:{fingerprint}  SET EX {dedupWindowMs/1000 + 2}
  • tryAcquire(fp){ acquired, fingerprint, requestId }
  • publishResult(fp, json) → writes result for waiters to poll
  • waitForResult(fp, pollMs?) → polls at 250 ms intervals, returns null on timeout

New — apps/server/src/budget/

tracker.ts — Auto-expiring Redis counters, pre-flight assertions, concurrency semaphore.

Redis keys
  grat:budget:hourly     INCRBY · EXPIRE → top of next hour
  grat:budget:daily      INCRBY · EXPIRE → next UTC midnight
  grat:budget:concurrent INCR on start  · DECR in finally
  • assertBudgetAvailable() — throws BudgetExceededError or ConcurrencyLimitError
  • getBudgetStatus() — non-throwing read of current counters
  • recordUsage(tokens) — called after every successful LLM response

errors.ts — Custom Error subclasses caught at the route layer and serialised into structured HTTP 429 bodies.


New — apps/server/src/ai/

types.ts — TypeScript interfaces shared across the subsystem:

interface AIExplanation {
  summary: string
  detailed_explanation: string
  suggested_fixes: string[]
  tokens_used: number
  cached: boolean        // true on cache hit or dedup coalesce
  request_id: string
}

interface BudgetStatus {
  allowed: boolean
  hourly_remaining: number
  daily_remaining: number
  concurrent_active: number
  concurrent_limit: number
  hourly_reset_at: string  // ISO 8601
  daily_reset_at: string
}

prompts.ts — Versioned templates. PROMPT_VERSION_HASH (first 12 chars of SHA-256(PROMPT_VERSION)) is embedded in every cache key — bumping PROMPT_VERSION auto-invalidates all stale cached explanations on next deploy.

client.tsexplainError(reportJson, txHash, network) orchestrates the full pipeline: cache → dedup → budget → concurrency → fetch → record → cache → publish.
Calls an OpenAI-compatible /chat/completions endpoint with response_format: json_object. Provider is swapped by changing AI_API_BASE_URL — no code changes required.


New — apps/server/src/cache/

manager.ts — Cache management layer.

Key schema: grat:cache:{type}:{promptVersionHash}:{fingerprint}

TTLs  decode / explain → CACHE_DECODE_TTL  (default 24 h)
      profile           → CACHE_PROFILE_TTL (default  1 h)
Function What it does
setCache(type, fp, val) Write with type TTL + version tag
getCache(type, fp) Read; null on miss or version mismatch
invalidateByTxHash(txHash) SCAN + DEL all entries for a transaction
invalidateEntry(type, fp) DEL across all prompt versions
clearByType(type) Bulk DEL one entry type
clearAll() Nuke entire grat:cache:* namespace

New — apps/server/src/routes/ai.ts

Four new Fastify endpoints registered at /api:

Method Path Description
POST /api/ai/explain Submit a DiagnosticReport, receive AIExplanation. Full pipeline.
GET /api/ai/budget Read current token usage and remaining budget.
DELETE /api/ai/cache/:txHash Invalidate cached explanations for one transaction.
DELETE /api/ai/cache Clear the entire AI cache (admin).

Budget/concurrency errors surface as:

// HTTP 429 — budget exceeded
{
  "error": "AI token budget exceeded: hourly limit reached. Resets at 2026-07-16T23:00:00.000Z.",
  "hourly_remaining": 0,
  "daily_remaining": 41230,
  "reset_at": "2026-07-16T23:00:00.000Z"
}

// HTTP 429 — concurrency limit
{
  "error": "AI concurrency limit reached: 5/5 requests in-flight.",
  "concurrent_active": 5,
  "concurrent_limit": 5
}

Modified — existing files

File Change
apps/server/src/config.ts Added ai, budget, cache config blocks — all env-var overridable
apps/server/src/index.ts Registered aiRoutes at /api prefix
apps/server/src/workers/replay.worker.ts Optionally calls explainError() after replay completes; AI failure is non-fatal
apps/server/tsconfig.json Added "types": ["node"], "lib": ["ES2022","DOM"]; excluded __tests__/ from build output
apps/server/package.json Added test + test:watch scripts; vitest in devDependencies

Environment Variables

Only AI_API_KEY is required. Omitting it leaves all existing routes unchanged.

# AI provider
AI_API_KEY=sk-...
AI_API_BASE_URL=https://api.openai.com/v1   # swap for Groq, Ollama, Anthropic compat, etc.
AI_MODEL=gpt-4o-mini
AI_MAX_TOKENS_PER_REQUEST=2048
AI_TEMPERATURE=0.2
AI_TIMEOUT_MS=30000

# Token budget
AI_HOURLY_TOKEN_LIMIT=100000
AI_DAILY_TOKEN_LIMIT=1000000
AI_MAX_CONCURRENT=5

# Cache TTLs (seconds)
CACHE_DECODE_TTL=86400   # 24 h — decode & explain entries
CACHE_PROFILE_TTL=3600   # 1 h  — profile entries
DEDUP_WINDOW_MS=10000    # 10 s — deduplication lock window

Tests

All Redis and LLM calls are fully mocked — npm test requires no external services.

 ✓  src/__tests__/cache.test.ts    (8 tests)   31 ms
 ✓  src/__tests__/budget.test.ts  (14 tests)   26 ms
 ✓  src/__tests__/dedup.test.ts    (9 tests)  ~2.4 s

 Test Files  3 passed (3)
      Tests  31 passed (31)
   Duration  ~3 s
Suite What is tested
cache.test.ts Round-trip, miss returns null, type namespace isolation, entry invalidation across prompt versions, txHash-scoped invalidation, type bulk clear, full clear, prompt version hash embedded in key
budget.test.ts Status under/over hourly limit, over daily limit, at concurrency cap; assertBudgetAvailable throws correct error type; counter accumulation; expiry set on first write; semaphore increment/decrement/floor-at-zero
dedup.test.ts Fingerprint stability, case-insensitivity, SHA-256 output format, lock acquisition, duplicate rejection, unique request IDs, result publish/poll coalescing, timeout path
cd apps/server
npm test            # single run
npm run test:watch  # watch mode

Acceptance Criteria

  • Repeated identical AI requests reuse cached or deduped results safelycached: true in response; dedup lock coalesces concurrent waiters onto the single in-flight call
  • Budget enforcement blocks expensive requests once configured thresholds are exceeded — HTTP 429 with reset_at timestamp; hourly and daily counters enforced independently
  • Cache behaviour is documented and covered by tests — 8 dedicated cache tests; JSDoc on every public function in cache/manager.ts

Reviewer Checklist

  • AI_API_KEY added to staging/production secrets before merging
  • Redis instance is reachable (already required by the existing BullMQ queue)
  • Confirm acceptable default budget limits (AI_HOURLY_TOKEN_LIMIT, AI_DAILY_TOKEN_LIMIT) for your environment
  • To invalidate all cached AI explanations after a future prompt change: bump PROMPT_VERSION in apps/server/src/ai/prompts.ts — no manual cache flush required
  • No breaking changes — all pre-existing routes (/health, /replay, /session) are untouched; new routes are purely additive
  • ioredis (already in dependencies) is the only new runtime dependency; vitest is the only new devDependency

…rrency control, and Redis-backed caching system
@codeZe-us
codeZe-us self-requested a review July 16, 2026 22:36
@codeZe-us

Copy link
Copy Markdown
Contributor

@A-one-tech please resolve workflow issues

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants