Skip to content

Add rate limiting and cost guardrails for Bedrock AI calls - #41

Merged
llinsss merged 2 commits into
DogStark:mainfrom
odiliflash-ops:feature/bedrock-rate-limiting
Aug 31, 2026
Merged

Add rate limiting and cost guardrails for Bedrock AI calls#41
llinsss merged 2 commits into
DogStark:mainfrom
odiliflash-ops:feature/bedrock-rate-limiting

Conversation

@odiliflash-ops

Copy link
Copy Markdown

Closes #7

Problem

/api/v1/hint and /api/v1/story both default use_bedrock=True with zero rate limiting, so a retry-happy frontend (or a malicious caller) could trigger unbounded Bedrock invocations and an unexpected AWS bill.

Solution

New module agent/bedrock_guardrails.py — the single place that caps how much Bedrock can be spent, per caller and globally:

  • Per-principal token bucket (BEDROCK_RATE_LIMIT_PER_MINUTE, default 10/min): keyed by student ID when the request names one (/story), else by account, else by client IP (anonymous /hint callers). Only requests that would actually reach Bedrock consume allowance (attempt-1 hints only; use_bedrock=False never counts). Thread-safe; buckets are memory-bounded with stale eviction.
    • Documented choice: exceeding this limit returns 429 with Retry-After, not a silent template fallback — silent fallback hides abuse from operators and gives retry-happy clients no signal to back off.
  • Global daily/monthly budgets (BEDROCK_DAILY_BUDGET default 1000/day, BEDROCK_MONTHLY_BUDGET default 20000/month, UTC windows): once either is exhausted, every Bedrock-backed request hard-cutovers to the deterministic template fallback until the window resets. The cutover is logged exactly once per exhausted window. Budget slots are reserved at dispatch time (just before invoke_model), so failed/retried provider calls cannot bypass the cap — deliberately conservative for billing.
  • Admin observability endpoint GET /api/v1/admin/bedrock-usage (requires admin role): daily/monthly usage vs limits, exhaustion flag, configured rate limit, tracked-principal count — aggregates only, never identifiers or raw principals.
  • All limits are env-configurable: negative value disables a limit, 0 blocks everything; invalid values fail loudly. Counters are in-memory per process (documented in README for horizontal deployments).

Acceptance criteria

  • Exceeding the per-student rate limit returns 429 with a Retry-After header (choice documented above and in README).
  • Exceeding the global budget forces template-only behavior for all requests until the window resets; test simulates budget exhaustion through the API and verifies recovery after rollover, plus the provider mock proving no further invoke_model happens.
  • Usage counters observable via the admin endpoint (auth matrix tested: 401/403/200).
  • Limits configurable via environment variables with documented sane defaults.

Testing

  • New tests/test_bedrock_cost_guardrails.py (28 tests): config parsing/defaults, token-bucket unit behavior, per-principal HTTP 429s (story per-student isolation, hint-by-IP for anonymous callers), global-budget cutover + once-per-window logging + daily/monthly reset recovery via the real API surface, admin endpoint auth matrix/payload shape, and an identifier-leak guard on the snapshot.
  • AWS is never contacted (invoke_model always mocked), matching the existing test conventions.
  • Full suite: 227 passed; ruff check . clean; mypy clean; coverage gate (80%) still met at ~86%.
  • Verified end-to-end against a live uvicorn server: third burst request returns HTTP/1.1 429 with retry-after, budget counters advance and are visible on the admin endpoint.

odiliemmanuel and others added 2 commits August 24, 2026 17:27
/api/v1/hint and /api/v1/story defaulted use_bedrock=True with no cap, so a
retry-happy frontend (or anonymous caller) could trigger unbounded Bedrock
invocations and an unexpected AWS bill.

- per-principal token bucket (agent/bedrock_guardrails.py): N calls/minute
  keyed by student ID when the request names one, else account, else client
  IP; exceeding it returns 429 with Retry-After rather than silently falling
  back to templates. Only requests that would actually reach Bedrock consume
  allowance.
- global daily/monthly budgets with hard cutover: once exhausted, every
  Bedrock-backed feature serves the template-only fallback until the window
  resets; the cutover is logged exactly once per exhausted window. Budget
  slots are reserved at dispatch time so failed/retried provider calls
  cannot bypass the cap.
- GET /api/v1/admin/bedrock-usage (admin role) exposes daily/monthly usage
  against limits, exhaustion state, configured rate limit, and tracked
  principal count - aggregates only, no identifiers.
- limits configurable via BEDROCK_RATE_LIMIT_PER_MINUTE,
  BEDROCK_DAILY_BUDGET, and BEDROCK_MONTHLY_BUDGET (documented defaults;
  negative disables a limit, zero blocks all).

Closes DogStark#7
@llinsss
llinsss merged commit 7d974e5 into DogStark:main Aug 31, 2026
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.

Add rate limiting and cost guardrails for Bedrock AI calls

3 participants