Add rate limiting and cost guardrails for Bedrock AI calls - #41
Merged
llinsss merged 2 commits intoAug 31, 2026
Merged
Conversation
/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
3 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #7
Problem
/api/v1/hintand/api/v1/storyboth defaultuse_bedrock=Truewith 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: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/hintcallers). Only requests that would actually reach Bedrock consume allowance (attempt-1 hints only;use_bedrock=Falsenever counts). Thread-safe; buckets are memory-bounded with stale eviction.Retry-After, not a silent template fallback — silent fallback hides abuse from operators and gives retry-happy clients no signal to back off.BEDROCK_DAILY_BUDGETdefault 1000/day,BEDROCK_MONTHLY_BUDGETdefault 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 beforeinvoke_model), so failed/retried provider calls cannot bypass the cap — deliberately conservative for billing.GET /api/v1/admin/bedrock-usage(requiresadminrole): daily/monthly usage vs limits, exhaustion flag, configured rate limit, tracked-principal count — aggregates only, never identifiers or raw principals.0blocks everything; invalid values fail loudly. Counters are in-memory per process (documented in README for horizontal deployments).Acceptance criteria
Retry-Afterheader (choice documented above and in README).invoke_modelhappens.Testing
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.invoke_modelalways mocked), matching the existing test conventions.ruff check .clean;mypyclean; coverage gate (80%) still met at ~86%.uvicornserver: third burst request returnsHTTP/1.1 429withretry-after, budget counters advance and are visible on the admin endpoint.