diff --git a/.gitignore b/.gitignore
index 10c5091..7cc8ee6 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,6 +4,7 @@ Somzilla.md
# Python
__pycache__/
*.pyc
+.venv/
# Test/lint/type-check tool caches and coverage output
.pytest_cache/
diff --git a/README.md b/README.md
index bfabd83..04cc90c 100644
--- a/README.md
+++ b/README.md
@@ -122,6 +122,11 @@ export BEDROCK_READ_TIMEOUT_SECONDS=8
export BEDROCK_MAX_RETRIES=2
export AI_SAFETY_POLICY_VERSION=2026-08-18
export ENABLE_GENERATIVE_FEATURES=1 # set to 0 to disable Bedrock hints/stories entirely
+
+# Optional cost guardrails, defaults shown
+export BEDROCK_RATE_LIMIT_PER_MINUTE=10 # per-student/account/IP; -1 disables, 0 blocks all
+export BEDROCK_DAILY_BUDGET=1000 # global invocations per UTC day; -1 = unlimited
+export BEDROCK_MONTHLY_BUDGET=20000 # global invocations per UTC month; -1 = unlimited
```
### AI safety guardrails
@@ -136,6 +141,48 @@ anything that fails is discarded in favor of the same deterministic template
fallback used when Bedrock is unavailable. See `agent/ai_safety.py` for the
full set of checks and `PRIVACY.md` for the privacy posture.
+### Bedrock cost guardrails
+
+`agent/bedrock_guardrails.py` caps how much Bedrock can be spent, per caller
+and globally (`agent/hint_generator.py` and `agent/story_mode.py` consult it;
+counters are in-memory and reset on process restart):
+
+- **Per-principal rate limit** — a token bucket of
+ `BEDROCK_RATE_LIMIT_PER_MINUTE` requests per minute, keyed by student ID
+ when the request names one (e.g. `/story`), else by account, else by client
+ IP (anonymous `/hint` callers). Only requests that would actually reach
+ Bedrock consume allowance: later hint attempts are deterministic reveals,
+ and `use_bedrock=False` never counts. Exceeding it returns **429** with a
+ `Retry-After` header rather than silently falling back to templates —
+ silent fallback hides abuse from operators and gives retry-happy clients no
+ signal to back off.
+- **Global daily/monthly budgets** — `BEDROCK_DAILY_BUDGET` /
+ `BEDROCK_MONTHLY_BUDGET` cap total invocations per UTC day/month across all
+ callers. A slot is reserved at dispatch time (just before `invoke_model`),
+ so failed or retried provider calls cannot bypass the cap — deliberately
+ conservative, since under-counting is what produces surprise bills. Once a
+ budget is exhausted, **every** Bedrock-backed request hard-cutovers to the
+ template-only fallback until the window resets; the cutover is logged once
+ per exhausted window.
+- **Observability** — `GET /api/v1/admin/bedrock-usage` (requires the `admin`
+ role) reports daily/monthly usage against limits, whether a budget is
+ currently exhausted, the configured rate limit, and how many principals are
+ tracked. Aggregates only — never student identifiers or raw principals.
+
+```json
+{
+ "daily": { "window": "2026-08-24", "used": 12, "limit": 1000 },
+ "monthly": { "window": "2026-08", "used": 340, "limit": 20000 },
+ "budget_exhausted": false,
+ "rate_limit": { "per_minute": 10 },
+ "tracked_principals": 2
+}
+```
+
+Counters are per-process. A single-process deployment gets exact enforcement;
+horizontally scaled deployments should front these endpoints with a shared
+limiter (or treat each replica's budget as a shard of the true limit).
+
### Privacy/retention configuration
```bash
@@ -170,8 +217,9 @@ teacher's classroom is simply that account's set of students. A caller can only
read or write students it owns; anything else returns `403`. Requests with no
key or an unknown key return `401`.
-Curriculum-management (`/api/v1/word-bank/*`) and experiment-reporting
-(`/api/v1/experiments/*`) routes are not scoped to a student set at all —
+Curriculum-management (`/api/v1/word-bank/*`), experiment-reporting
+(`/api/v1/experiments/*`), and Bedrock-usage reporting
+(`/api/v1/admin/bedrock-usage`) routes are not scoped to a student set at all —
they require an account with the `admin` role (or, for experiment reporting,
the `researcher` role) rather than any parent/teacher account.
@@ -420,6 +468,18 @@ POST /api/v1/experiments/report/export?retention_days=30
{ "exported_file": "experiment_report.json" }
```
+### Get Bedrock Usage
+
+Requires an `admin` account. Returns the current Bedrock cost-guardrail
+counters — daily/monthly budget consumption against limits, whether a budget
+is exhausted, the per-principal rate limit, and the number of tracked
+principals (aggregates only, no identifiers). See
+[Bedrock cost guardrails](#bedrock-cost-guardrails).
+
+```http
+GET /api/v1/admin/bedrock-usage
+```
+
---
## Structured Logging
diff --git a/agent/auth.py b/agent/auth.py
index ca18b2e..7c3b9d5 100644
--- a/agent/auth.py
+++ b/agent/auth.py
@@ -69,6 +69,18 @@ def require_account(
return account
+def resolve_account_from_key(raw_key: str | None) -> Account | None:
+ """Resolve a raw bearer key to an Account, or None if unknown/empty.
+
+ Used where authentication is optional (e.g. attributing anonymous
+ Bedrock rate-limit buckets to an account when a key is supplied) and
+ must not raise 401 the way require_account does.
+ """
+ if not raw_key:
+ return None
+ return _load_registry().get(_hash_key(raw_key))
+
+
def authorize_student(account: Account, student_id: str):
"""Ensure the account is allowed to act on this student, or reject with 403.
diff --git a/agent/bedrock_guardrails.py b/agent/bedrock_guardrails.py
new file mode 100644
index 0000000..e5666ed
--- /dev/null
+++ b/agent/bedrock_guardrails.py
@@ -0,0 +1,291 @@
+"""Cost guardrails for AWS Bedrock usage (rate limiting + global budgets).
+
+Both `/api/v1/hint` and `/api/v1/story` default `use_bedrock=True`, and a
+retry-happy frontend (or an anonymous caller — the hint route is public until
+auth lands there) could otherwise trigger unbounded `invoke_model` calls and
+an unexpected AWS bill. This module is the single place that:
+
+- enforces a per-principal token-bucket rate limit (N Bedrock-backed requests
+ per minute; "principal" is the authenticated account or student when the
+ request carries one, else the client IP) — enforced at the API layer, which
+ answers HTTP 429 with a Retry-After header;
+- maintains global daily and monthly invocation budgets with a hard cutover:
+ once either budget is exhausted, every Bedrock-backed feature silently
+ degrades to its deterministic template fallback until the window resets,
+ and the cutover is logged exactly once per exhausted window. Budgets are
+ consumed at dispatch time (just before `invoke_model`), so failed or
+ retried provider calls cannot bypass the cap;
+- exposes a usage snapshot for the admin endpoint
+ (`GET /api/v1/admin/bedrock-usage`) without leaking any student identifier.
+
+All limits are environment-configurable (see README, "Bedrock cost
+guardrails"). Counters are in-memory: they bound spend per process lifetime
+per calendar window and reset on restart, which is the right trade-off for a
+single-process deployment; horizontal deployments should front this with a
+shared limiter.
+"""
+
+import math
+import os
+import threading
+import time
+from datetime import UTC, datetime
+
+from agent.log_config import get_logger
+
+logger = get_logger(__name__)
+
+DEFAULT_RATE_LIMIT_PER_MINUTE = 10
+DEFAULT_DAILY_BUDGET = 1000
+DEFAULT_MONTHLY_BUDGET = 20000
+
+RATE_LIMIT_WINDOW_SECONDS = 60.0
+
+# Safety valve against unbounded memory growth from per-IP buckets: when more
+# than this many principals are tracked, stale/idle buckets are pruned on the
+# next acquire.
+_MAX_TRACKED_PRINCIPALS = 10_000
+
+
+class RateLimitExceededError(Exception):
+ """Raised when a principal exceeds the per-minute Bedrock rate limit."""
+
+ def __init__(self, retry_after_seconds: int):
+ self.retry_after_seconds = retry_after_seconds
+ super().__init__(
+ f"Bedrock rate limit exceeded; retry after {retry_after_seconds}s."
+ )
+
+
+class _TokenBucket:
+ """Classic token bucket: burst up to capacity, refilling linearly."""
+
+ __slots__ = ("capacity", "tokens", "updated_at")
+
+ def __init__(self, capacity: float, now: float):
+ self.capacity = capacity
+ self.tokens = capacity
+ self.updated_at = now
+
+ def try_acquire(self, now: float) -> tuple[bool, int]:
+ """Attempt to take one token; return (ok, retry_after_seconds)."""
+ if self.capacity <= 0:
+ # Emergency-stop configuration (BEDROCK_RATE_LIMIT_PER_MINUTE=0):
+ # nothing is ever allowed; suggest retrying after a full window.
+ return False, int(RATE_LIMIT_WINDOW_SECONDS)
+ refill_per_second = self.capacity / RATE_LIMIT_WINDOW_SECONDS
+ self.tokens = min(
+ self.capacity, self.tokens + (now - self.updated_at) * refill_per_second
+ )
+ self.updated_at = now
+ if self.tokens >= 1.0:
+ self.tokens -= 1.0
+ return True, 0
+ deficit = 1.0 - self.tokens
+ return False, max(1, math.ceil(deficit / refill_per_second))
+
+
+_lock = threading.Lock()
+_buckets: dict[str, _TokenBucket] = {}
+
+
+# --------------------------------------------------------------------------
+# Configuration (environment variables)
+# --------------------------------------------------------------------------
+
+def _int_env(name: str, default: int) -> int:
+ """Parse a limit env var; negative values mean 'unlimited/disabled'."""
+ raw = os.getenv(name, str(default))
+ try:
+ return int(raw)
+ except ValueError as exc:
+ raise ValueError(f"{name} must be an integer.") from exc
+
+
+def get_rate_limit_per_minute() -> int:
+ """Max Bedrock-backed requests per principal per minute (-1 disables)."""
+ return _int_env("BEDROCK_RATE_LIMIT_PER_MINUTE", DEFAULT_RATE_LIMIT_PER_MINUTE)
+
+
+def get_daily_budget() -> int:
+ """Global Bedrock invocations allowed per UTC day (-1 = unlimited)."""
+ return _int_env("BEDROCK_DAILY_BUDGET", DEFAULT_DAILY_BUDGET)
+
+
+def get_monthly_budget() -> int:
+ """Global Bedrock invocations allowed per UTC month (-1 = unlimited)."""
+ return _int_env("BEDROCK_MONTHLY_BUDGET", DEFAULT_MONTHLY_BUDGET)
+
+
+# --------------------------------------------------------------------------
+# Per-principal rate limiting (API layer → HTTP 429)
+# --------------------------------------------------------------------------
+
+def _prune_stale_buckets(now: float) -> None:
+ """Bound memory: drop idle buckets when tracking too many principals."""
+ if len(_buckets) <= _MAX_TRACKED_PRINCIPALS:
+ return
+ stale_cutoff = now - 2 * RATE_LIMIT_WINDOW_SECONDS
+ stale = [key for key, b in _buckets.items() if b.updated_at < stale_cutoff]
+ for key in stale:
+ del _buckets[key]
+ if len(_buckets) > _MAX_TRACKED_PRINCIPALS:
+ # Still full of fresh buckets (heavy traffic): evict oldest first.
+ ordered = sorted(_buckets.items(), key=lambda item: item[1].updated_at)
+ excess = len(_buckets) - int(_MAX_TRACKED_PRINCIPALS * 0.9)
+ for key, _ in ordered[:excess]:
+ del _buckets[key]
+
+
+def acquire_request_slot(principal: str) -> None:
+ """Charge one Bedrock-backed request to `principal`'s token bucket.
+
+ Raises RateLimitExceededError (with a Retry-After hint) when the
+ principal's per-minute allowance is spent. A non-positive configured
+ limit disables enforcement entirely.
+ """
+ limit = get_rate_limit_per_minute()
+ if limit < 0:
+ return
+ now = time.monotonic()
+ with _lock:
+ bucket = _buckets.get(principal)
+ if bucket is None or bucket.capacity != limit:
+ # Rebuild when missing or when the configured limit changed.
+ bucket = _TokenBucket(float(max(limit, 0)), now)
+ _buckets[principal] = bucket
+ ok, retry_after = bucket.try_acquire(now)
+ _prune_stale_buckets(now)
+ if not ok:
+ logger.warning(
+ "Bedrock rate limit exceeded for a principal",
+ extra={
+ "source_module": __name__,
+ "source_function": "acquire_request_slot",
+ "provider_outcome": "rate_limited",
+ "retry_after_seconds": retry_after,
+ },
+ )
+ raise RateLimitExceededError(retry_after)
+
+
+# --------------------------------------------------------------------------
+# Global budgets (generator layer → template-only fallback)
+# --------------------------------------------------------------------------
+
+# Counters are keyed implicitly: they only ever hold the *current* UTC
+# day/month, so a window reset is just "the date changed". The cutover-log
+# markers work the same way and are dropped when their window rolls over so
+# each exhausted window logs exactly once.
+_daily_count = 0
+_monthly_count = 0
+_current_day = ""
+_current_month = ""
+_cutover_logged_windows: set[str] = set()
+
+
+def _utc_day() -> str:
+ return datetime.now(UTC).strftime("%Y-%m-%d")
+
+
+def _utc_month() -> str:
+ return datetime.now(UTC).strftime("%Y-%m")
+
+
+def try_consume_budget(feature: str) -> bool:
+ """Atomically reserve one Bedrock invocation against the global budgets.
+
+ Returns True when the caller may proceed with `invoke_model`; returns
+ False once either the daily or monthly budget is exhausted, logging the
+ hard cutover to template-only mode once per exhausted window. The slot is
+ consumed before the provider call so retries/failures cannot bypass the
+ cap (deliberately conservative for billing).
+ """
+ global _daily_count, _monthly_count, _current_day, _current_month
+
+ daily_limit = get_daily_budget()
+ monthly_limit = get_monthly_budget()
+
+ with _lock:
+ today = _utc_day()
+ if today != _current_day:
+ _current_day = today
+ _daily_count = 0
+ _cutover_logged_windows.discard("daily")
+ month = _utc_month()
+ if month != _current_month:
+ _current_month = month
+ _monthly_count = 0
+ _cutover_logged_windows.discard("monthly")
+
+ over_daily = 0 <= daily_limit <= _daily_count
+ over_monthly = 0 <= monthly_limit <= _monthly_count
+ if over_daily or over_monthly:
+ window = "daily" if over_daily else "monthly"
+ if window not in _cutover_logged_windows:
+ _cutover_logged_windows.add(window)
+ logger.warning(
+ "Global Bedrock budget exhausted; serving template-only "
+ "fallback until the window resets",
+ extra={
+ "source_module": __name__,
+ "source_function": "try_consume_budget",
+ "feature": feature,
+ "budget_window": window,
+ "budget_limit": daily_limit if over_daily else monthly_limit,
+ "budget_used": _daily_count if over_daily else _monthly_count,
+ "outcome": "budget_exhausted",
+ },
+ )
+ return False
+
+ # Count even when a limit is disabled (-1) so usage stays observable.
+ _daily_count += 1
+ _monthly_count += 1
+ return True
+
+
+def usage_snapshot() -> dict:
+ """Current guardrail state for the admin endpoint.
+
+ Contains only aggregates — never student identifiers or raw principals.
+ """
+ daily_limit = get_daily_budget()
+ monthly_limit = get_monthly_budget()
+ rate_limit = get_rate_limit_per_minute()
+ with _lock:
+ daily_used = _daily_count
+ monthly_used = _monthly_count
+ day = _current_day or _utc_day()
+ month = _current_month or _utc_month()
+ principals = len(_buckets)
+ return {
+ "daily": {
+ "window": day,
+ "used": daily_used,
+ "limit": None if daily_limit < 0 else daily_limit,
+ },
+ "monthly": {
+ "window": month,
+ "used": monthly_used,
+ "limit": None if monthly_limit < 0 else monthly_limit,
+ },
+ "budget_exhausted": (
+ (0 <= daily_limit <= daily_used) or (0 <= monthly_limit <= monthly_used)
+ ),
+ "rate_limit": {"per_minute": None if rate_limit < 0 else rate_limit},
+ "tracked_principals": principals,
+ }
+
+
+def reset_state() -> None:
+ """Drop all counters and buckets so tests start from a clean slate."""
+ global _buckets, _daily_count, _monthly_count, _current_day, _current_month
+
+ with _lock:
+ _buckets = {}
+ _daily_count = 0
+ _monthly_count = 0
+ _current_day = ""
+ _current_month = ""
+ _cutover_logged_windows.clear()
diff --git a/agent/hint_generator.py b/agent/hint_generator.py
index c7b01d8..3a9b5e3 100644
--- a/agent/hint_generator.py
+++ b/agent/hint_generator.py
@@ -14,6 +14,7 @@
validate_theme,
validate_word,
)
+from agent.bedrock_guardrails import try_consume_budget
from agent.log_config import get_logger
logger = get_logger(__name__)
@@ -102,6 +103,10 @@ def _parse_structured_hint(raw_text: str) -> str:
def _bedrock_hint(word: str, theme: str) -> str | None:
+ if not try_consume_budget("hint"):
+ # Global budget exhausted: hard cutover to the template fallback
+ # (the cutover itself is logged once per window by the guardrail).
+ return None
try:
client = boto3.client("bedrock-runtime", config=bedrock_client_config())
prompt = f"{word}\n{theme}\nGive the hint now."
diff --git a/agent/story_mode.py b/agent/story_mode.py
index ae20599..c5d2b3b 100644
--- a/agent/story_mode.py
+++ b/agent/story_mode.py
@@ -13,6 +13,7 @@
validate_story_output,
validate_words_for_generation,
)
+from agent.bedrock_guardrails import try_consume_budget
from agent.log_config import get_logger
logger = get_logger(__name__)
@@ -76,6 +77,10 @@ def _parse_structured_story(raw_text: str) -> str:
def _bedrock_story(words: list) -> str | None:
+ if not try_consume_budget("story"):
+ # Global budget exhausted: hard cutover to the template fallback
+ # (the cutover itself is logged once per window by the guardrail).
+ return None
word_count = len(words)
try:
client = boto3.client("bedrock-runtime", config=bedrock_client_config())
diff --git a/api/routes.py b/api/routes.py
index 1e548e0..0ae9b18 100644
--- a/api/routes.py
+++ b/api/routes.py
@@ -2,7 +2,7 @@
from datetime import datetime
from typing import Literal
-from fastapi import APIRouter, Depends, HTTPException, Query, status
+from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from pydantic import BaseModel, ConfigDict, Field
from agent.ai_safety import UnsafeContentError, validate_theme, validate_word
@@ -12,6 +12,12 @@
require_account,
require_admin,
require_researcher,
+ resolve_account_from_key,
+)
+from agent.bedrock_guardrails import (
+ RateLimitExceededError,
+ acquire_request_slot,
+ usage_snapshot,
)
from agent.diagnostic import get_next_diagnostic_question, submit_diagnostic_answer
from agent.hint_generator import get_encouragement, get_hint
@@ -153,35 +159,50 @@ def _word_bank_http_error(exc: WordBankError) -> HTTPException:
return HTTPException(status_code=422, detail=str(exc))
-def _validate_attempt_payload(req: AttemptRequest) -> tuple[str, str, list[str]]:
- """Reject attempts that reference content outside the canonical curriculum.
+def _bearer_key(request: Request) -> str | None:
+ header = request.headers.get("Authorization", "")
+ if not header.startswith("Bearer "):
+ return None
+ key = header.removeprefix("Bearer ").strip()
+ return key or None
+
+
+def _bedrock_principal(
+ request: Request,
+ student_id: str | None = None,
+) -> str:
+ """Identity a Bedrock rate-limit charge is attributed to.
- Attempts write directly into the learning profile (``words``,
- ``theme_preferences``, ``phonics_struggles``), which feeds the recommender,
- difficulty calibration, and reports. A client could otherwise fabricate
- arbitrary words, themes, or phonics tags and permanently pollute that data.
- This boundary check pins every attempt to a canonical curriculum word with
- its real theme and phonics tags, and returns the normalized values to store
- so the profile never accumulates duplicate casing/spelling variants.
+ Per-student when the request names one (finest grain), else the
+ authenticated account, else the client IP for anonymous callers such
+ as /hint. Raw principals never leave this module: only aggregate
+ counts are observable via the admin endpoint.
+ """
+ if student_id is not None:
+ return f"student:{student_id}"
+ account = resolve_account_from_key(_bearer_key(request))
+ if account is not None:
+ return f"account:{account.account_id}"
+ client = request.client
+ return f"ip:{client.host if client else 'unknown'}"
+
+
+def _enforce_bedrock_rate_limit(principal: str) -> None:
+ """Apply the per-principal token bucket to a Bedrock-backed request.
+
+ Design choice (documented in README): exceeding the per-principal rate
+ limit returns 429 with a Retry-After hint rather than silently falling
+ back to templates — silent fallback hides abuse from operators and
+ gives retry-happy clients no signal to back off.
"""
try:
- word = validate_word(req.word)
- theme = validate_theme(req.theme)
- entry = get_word_entry(word) # validate_word guarantees membership
- except (UnsafeContentError, WordNotFoundError) as exc:
- raise HTTPException(status_code=422, detail=str(exc)) from exc
- if theme != entry["theme"]:
- raise HTTPException(
- status_code=422,
- detail=f"theme '{theme}' does not match the curriculum entry for word '{word}'.",
- )
- unknown_tags = sorted(set(req.phonics_tags) - set(entry["phonics"]))
- if unknown_tags:
+ acquire_request_slot(principal)
+ except RateLimitExceededError as exc:
raise HTTPException(
- status_code=422,
- detail=f"phonics_tags contains tags not present on word '{word}': {unknown_tags}",
- )
- return word, theme, req.phonics_tags
+ status_code=status.HTTP_429_TOO_MANY_REQUESTS,
+ detail=str(exc),
+ headers={"Retry-After": str(exc.retry_after_seconds)},
+ ) from exc
# --- Endpoints ---
@@ -262,7 +283,11 @@ def get_recommendations(req: RecommendRequest, account: Account = Depends(requir
@router.post("/hint")
-def get_word_hint(req: HintRequest):
+def get_word_hint(req: HintRequest, request: Request):
+ # Only attempt-1 hints consult Bedrock (later attempts are deterministic
+ # first-letter reveals), so only those consume rate-limit allowance.
+ if req.use_bedrock and req.attempt_number == 1:
+ _enforce_bedrock_rate_limit(_bedrock_principal(request))
hint = get_hint(req.word, req.theme, req.attempt_number, req.use_bedrock)
if req.use_bedrock:
is_fallback = hint.startswith(("It's a", "It belongs to"))
@@ -282,8 +307,10 @@ def get_word_hint(req: HintRequest):
@router.post("/story")
-def create_story(req: StoryRequest, account: Account = Depends(require_account)):
+def create_story(req: StoryRequest, request: Request, account: Account = Depends(require_account)):
authorize_student(account, req.student_id)
+ if req.use_bedrock:
+ _enforce_bedrock_rate_limit(_bedrock_principal(request, req.student_id))
# Story requests carry a student ID and therefore use the same consent gate.
# The student ID is never forwarded to generate_story: it is a persistent
# child identifier and the story generator has no use for a display name.
@@ -455,6 +482,17 @@ def delete_curriculum_word(word: str, account: Account = Depends(require_admin))
return {"deleted": True, "word": deleted["word"]}
+@router.get("/admin/bedrock-usage")
+def get_bedrock_usage(account: Account = Depends(require_admin)):
+ """Current Bedrock cost-guardrail counters (admin only).
+
+ Aggregates only: daily/monthly budget consumption against their limits,
+ the configured per-principal rate limit, and how many principals are
+ tracked. No student identifiers or raw principals are exposed.
+ """
+ return usage_snapshot()
+
+
@router.post("/onboarding/diagnostic/next")
def get_next_question(req: DiagnosticNextRequest, account: Account = Depends(require_account)):
"""Retrieve the next onboarding question, enforcing consent before storage."""
diff --git a/tests/test_bedrock_cost_guardrails.py b/tests/test_bedrock_cost_guardrails.py
new file mode 100644
index 0000000..0f52172
--- /dev/null
+++ b/tests/test_bedrock_cost_guardrails.py
@@ -0,0 +1,538 @@
+"""Tests for Bedrock rate limiting and cost guardrails (issue #7).
+
+SCOPE: `/api/v1/hint` and `/api/v1/story` both default `use_bedrock=True`,
+so a retry-happy or malicious caller could trigger unbounded Bedrock
+invocations. These tests cover:
+
+- the per-principal token bucket (`agent.bedrock_guardrails`) and its HTTP
+ surface — exceeding it returns 429 with Retry-After, not a silent fallback;
+- the global daily/monthly budget with hard cutover to template-only
+ fallback once exceeded (cutover logged exactly once per window), including
+ recovery after the window resets;
+- observability of usage counters via GET /api/v1/admin/bedrock-usage
+ (admin-only);
+- environment-variable configurability of every limit.
+
+AWS is never contacted: `invoke_model` is always mocked.
+"""
+
+import hashlib
+import json
+import os
+import shutil
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+TEST_PROFILES_DIR = "/tmp/test_bedrock_guardrail_profiles"
+TEST_ACCOUNTS_FILE = "/tmp/test_bedrock_guardrail_accounts.json"
+
+ADMIN_KEY = "test_guardrails_admin_key"
+PARENT_KEY = "test_guardrails_parent_key"
+
+CONSENT_METADATA = {
+ "guardian_id": "guardian_test_001",
+ "relationship": "parent",
+ "consent_given": True,
+ "consent_method": "verified_test_form",
+ "privacy_policy_version": "test-v1",
+ "consented_at": "2025-01-01T00:00:00+00:00",
+}
+
+BEDROCK_STORY_TEXT = json.dumps({
+ "story": "The cat found a hat. A bat flew by and waved. They all smiled."
+})
+TEMPLATE_STORY_MARKER = "went on a big adventure"
+
+
+def _key_hash(raw_key):
+ return hashlib.sha256(raw_key.encode()).hexdigest()
+
+
+def auth(key=PARENT_KEY):
+ return {"Authorization": f"Bearer {key}"}
+
+
+def _mock_invoke_response(text):
+ """Build a mock Bedrock invoke_model response whose content is `text`."""
+ mock_response = MagicMock()
+ mock_response["body"].read.return_value = json.dumps({
+ "content": [{"text": text}]
+ }).encode()
+ return mock_response
+
+
+@pytest.fixture(autouse=True)
+def guardrails_env(monkeypatch):
+ """Fresh guardrail state and default limits for every test."""
+ from agent import bedrock_guardrails as guardrails
+
+ for name in (
+ "BEDROCK_RATE_LIMIT_PER_MINUTE",
+ "BEDROCK_DAILY_BUDGET",
+ "BEDROCK_MONTHLY_BUDGET",
+ ):
+ monkeypatch.delenv(name, raising=False)
+ guardrails.reset_state()
+ yield
+ guardrails.reset_state()
+
+
+@pytest.fixture(autouse=True)
+def patch_storage(monkeypatch):
+ os.makedirs(TEST_PROFILES_DIR, exist_ok=True)
+ monkeypatch.setattr("agent.profiler.PROFILES_DIR", TEST_PROFILES_DIR)
+
+ from agent import auth as auth_module
+
+ accounts = [
+ {
+ "account_id": "guardrail_admin",
+ "role": "admin",
+ "api_key_sha256": _key_hash(ADMIN_KEY),
+ "student_ids": [],
+ },
+ {
+ "account_id": "guardrail_parent",
+ "role": "parent",
+ "api_key_sha256": _key_hash(PARENT_KEY),
+ "student_ids": ["student_001", "student_002"],
+ },
+ ]
+ with open(TEST_ACCOUNTS_FILE, "w") as f:
+ json.dump(accounts, f)
+ monkeypatch.setattr(auth_module, "ACCOUNTS_FILE", TEST_ACCOUNTS_FILE)
+ auth_module.reset_registry()
+ yield
+ auth_module.reset_registry()
+ if os.path.exists(TEST_ACCOUNTS_FILE):
+ os.remove(TEST_ACCOUNTS_FILE)
+ shutil.rmtree(TEST_PROFILES_DIR, ignore_errors=True)
+
+
+@pytest.fixture
+def client():
+ from fastapi.testclient import TestClient
+
+ from main import app
+ return TestClient(app)
+
+
+@pytest.fixture
+def mock_bedrock_story():
+ """Successful Bedrock story generation; returns the invoke_model mock."""
+ with patch("agent.story_mode.boto3.client") as mock_client:
+ mock_client.return_value.invoke_model.return_value = _mock_invoke_response(
+ BEDROCK_STORY_TEXT
+ )
+ yield mock_client
+
+
+@pytest.fixture
+def mock_budget_logger(monkeypatch):
+ from agent import bedrock_guardrails as guardrails
+
+ logger = MagicMock()
+ monkeypatch.setattr(guardrails, "logger", logger)
+ return logger
+
+
+# ── Configuration ───────────────────────────────────────────────────────────
+
+class TestConfiguration:
+ def test_documented_defaults(self, monkeypatch):
+ from agent import bedrock_guardrails as guardrails
+
+ monkeypatch.delenv("BEDROCK_RATE_LIMIT_PER_MINUTE", raising=False)
+ monkeypatch.delenv("BEDROCK_DAILY_BUDGET", raising=False)
+ monkeypatch.delenv("BEDROCK_MONTHLY_BUDGET", raising=False)
+ assert guardrails.get_rate_limit_per_minute() == 10
+ assert guardrails.get_daily_budget() == 1000
+ assert guardrails.get_monthly_budget() == 20000
+
+ def test_limits_are_configurable_via_environment(self, monkeypatch):
+ from agent import bedrock_guardrails as guardrails
+
+ monkeypatch.setenv("BEDROCK_RATE_LIMIT_PER_MINUTE", "3")
+ monkeypatch.setenv("BEDROCK_DAILY_BUDGET", "50")
+ monkeypatch.setenv("BEDROCK_MONTHLY_BUDGET", "-1")
+ assert guardrails.get_rate_limit_per_minute() == 3
+ assert guardrails.get_daily_budget() == 50
+ assert guardrails.get_monthly_budget() == -1
+
+ def test_invalid_daily_or_monthly_budget_fails_loudly(self, monkeypatch):
+ from agent import bedrock_guardrails as guardrails
+
+ monkeypatch.setenv("BEDROCK_DAILY_BUDGET", "not-a-number")
+ with pytest.raises(ValueError):
+ guardrails.try_consume_budget("story")
+
+ def test_invalid_rate_limit_fails_loudly(self, monkeypatch):
+ from agent import bedrock_guardrails as guardrails
+
+ monkeypatch.setenv("BEDROCK_RATE_LIMIT_PER_MINUTE", "not-a-number")
+ with pytest.raises(ValueError):
+ guardrails.acquire_request_slot("principal_a")
+
+
+# ── Per-principal token bucket (unit level) ─────────────────────────────────
+
+class TestTokenBucket:
+ def test_burst_up_to_limit_then_blocked(self, monkeypatch):
+ from agent import bedrock_guardrails as guardrails
+
+ monkeypatch.setenv("BEDROCK_RATE_LIMIT_PER_MINUTE", "2")
+ guardrails.acquire_request_slot("principal_a")
+ guardrails.acquire_request_slot("principal_a")
+ with pytest.raises(guardrails.RateLimitExceededError) as excinfo:
+ guardrails.acquire_request_slot("principal_a")
+ assert excinfo.value.retry_after_seconds >= 1
+
+ def test_principals_are_isolated(self, monkeypatch):
+ from agent import bedrock_guardrails as guardrails
+
+ monkeypatch.setenv("BEDROCK_RATE_LIMIT_PER_MINUTE", "1")
+ guardrails.acquire_request_slot("principal_a")
+ # A different principal still has its own allowance.
+ guardrails.acquire_request_slot("principal_b")
+
+ def test_negative_limit_disables_enforcement(self, monkeypatch):
+ from agent import bedrock_guardrails as guardrails
+
+ monkeypatch.setenv("BEDROCK_RATE_LIMIT_PER_MINUTE", "-1")
+ for _ in range(50):
+ guardrails.acquire_request_slot("principal_a")
+
+ def test_zero_limit_blocks_everything(self, monkeypatch):
+ from agent import bedrock_guardrails as guardrails
+
+ monkeypatch.setenv("BEDROCK_RATE_LIMIT_PER_MINUTE", "0")
+ with pytest.raises(guardrails.RateLimitExceededError) as excinfo:
+ guardrails.acquire_request_slot("principal_a")
+ assert excinfo.value.retry_after_seconds <= 60
+
+ def test_tracked_buckets_are_bounded(self, monkeypatch):
+ """Memory safety valve: stuffing far more principals than the
+ tracking cap must evict stale buckets instead of growing forever."""
+ import time as time_module
+
+ from agent import bedrock_guardrails as guardrails
+
+ monkeypatch.setenv("BEDROCK_RATE_LIMIT_PER_MINUTE", "1")
+ now = time_module.monotonic()
+ for i in range(guardrails._MAX_TRACKED_PRINCIPALS + 500):
+ bucket = guardrails._TokenBucket(1.0, now - 1000)
+ guardrails._buckets[f"stale:{i}"] = bucket
+ guardrails.acquire_request_slot("fresh_principal")
+ assert len(guardrails._buckets) < guardrails._MAX_TRACKED_PRINCIPALS
+
+
+# ── Global budget (unit level) ──────────────────────────────────────────────
+
+class TestGlobalBudget:
+ def test_slots_are_consumed_until_exhausted_then_refused(self, monkeypatch):
+ from agent import bedrock_guardrails as guardrails
+
+ monkeypatch.setenv("BEDROCK_DAILY_BUDGET", "2")
+ assert guardrails.try_consume_budget("hint") is True
+ assert guardrails.try_consume_budget("story") is True
+ assert guardrails.try_consume_budget("hint") is False
+
+ def test_cutover_is_logged_once_per_window(
+ self, monkeypatch, mock_budget_logger
+ ):
+ from agent import bedrock_guardrails as guardrails
+
+ monkeypatch.setenv("BEDROCK_DAILY_BUDGET", "1")
+ guardrails.try_consume_budget("hint")
+ guardrails.try_consume_budget("hint")
+ guardrails.try_consume_budget("hint")
+ exhausted_calls = [
+ c for c in mock_budget_logger.warning.call_args_list
+ if c.kwargs.get("extra", {}).get("outcome") == "budget_exhausted"
+ ]
+ assert len(exhausted_calls) == 1
+
+ def test_usage_is_counted_even_when_limits_are_disabled(self, monkeypatch):
+ from agent import bedrock_guardrails as guardrails
+
+ monkeypatch.setenv("BEDROCK_DAILY_BUDGET", "-1")
+ monkeypatch.setenv("BEDROCK_MONTHLY_BUDGET", "-1")
+ guardrails.try_consume_budget("story")
+ snapshot = guardrails.usage_snapshot()
+ assert snapshot["daily"]["used"] == 1
+ assert snapshot["daily"]["limit"] is None
+ assert snapshot["budget_exhausted"] is False
+
+ def test_budget_recovers_when_the_daily_window_resets(self, monkeypatch):
+ from agent import bedrock_guardrails as guardrails
+
+ monkeypatch.setenv("BEDROCK_DAILY_BUDGET", "1")
+ assert guardrails.try_consume_budget("hint") is True
+ assert guardrails.try_consume_budget("hint") is False
+ monkeypatch.setattr(guardrails, "_utc_day", lambda: "2099-01-02")
+ assert guardrails.try_consume_budget("hint") is True
+
+ def test_budget_recovers_when_the_monthly_window_resets(self, monkeypatch):
+ from agent import bedrock_guardrails as guardrails
+
+ monkeypatch.setenv("BEDROCK_MONTHLY_BUDGET", "1")
+ assert guardrails.try_consume_budget("story") is True
+ assert guardrails.try_consume_budget("story") is False
+ monkeypatch.setattr(guardrails, "_utc_month", lambda: "2099-02")
+ assert guardrails.try_consume_budget("story") is True
+
+ def test_either_limit_being_exhausted_blocks_all_features(
+ self, monkeypatch
+ ):
+ from agent import bedrock_guardrails as guardrails
+
+ monkeypatch.setenv("BEDROCK_DAILY_BUDGET", "-1")
+ monkeypatch.setenv("BEDROCK_MONTHLY_BUDGET", "1")
+ assert guardrails.try_consume_budget("hint") is True
+ # Daily is unlimited, but the monthly cap blocks everything.
+ assert guardrails.try_consume_budget("story") is False
+
+
+# ── HTTP surface: per-principal 429s ────────────────────────────────────────
+
+class TestPerStudentRateLimitHTTP:
+ def test_story_returns_429_with_retry_after_after_limit(
+ self, client, mock_bedrock_story, monkeypatch
+ ):
+ monkeypatch.setenv("BEDROCK_RATE_LIMIT_PER_MINUTE", "2")
+ from agent.profiler import load_profile
+ load_profile("student_001", consent_metadata=CONSENT_METADATA)
+
+ story_body = {
+ "student_id": "student_001",
+ "words": ["cat", "hat"],
+ "use_bedrock": True,
+ }
+ assert client.post("/api/v1/story", json=story_body, headers=auth()).status_code == 200
+ assert client.post("/api/v1/story", json=story_body, headers=auth()).status_code == 200
+
+ limited = client.post("/api/v1/story", json=story_body, headers=auth())
+ assert limited.status_code == 429
+ assert int(limited.headers["Retry-After"]) >= 1
+
+ def test_rate_limit_is_per_student_not_global(
+ self, client, mock_bedrock_story, monkeypatch
+ ):
+ monkeypatch.setenv("BEDROCK_RATE_LIMIT_PER_MINUTE", "1")
+ from agent.profiler import load_profile
+ load_profile("student_001", consent_metadata=CONSENT_METADATA)
+ load_profile("student_002", consent_metadata=CONSENT_METADATA)
+
+ first = client.post("/api/v1/story", json={
+ "student_id": "student_001", "words": ["cat"], "use_bedrock": True,
+ }, headers=auth())
+ other = client.post("/api/v1/story", json={
+ "student_id": "student_002", "words": ["cat"], "use_bedrock": True,
+ }, headers=auth())
+ assert first.status_code == 200
+ assert other.status_code == 200
+
+ def test_template_only_requests_do_not_consume_allowance(
+ self, client, monkeypatch
+ ):
+ monkeypatch.setenv("BEDROCK_RATE_LIMIT_PER_MINUTE", "1")
+ from agent.profiler import load_profile
+ load_profile("student_001", consent_metadata=CONSENT_METADATA)
+
+ body = {
+ "student_id": "student_001",
+ "words": ["cat"],
+ "use_bedrock": False,
+ }
+ for _ in range(5):
+ response = client.post("/api/v1/story", json=body, headers=auth())
+ assert response.status_code == 200
+
+ def test_hint_returns_429_for_anonymous_caller_from_one_ip(
+ self, client, monkeypatch
+ ):
+ """The hint route is public until auth lands there, so anonymous
+ callers are bucketed by client IP."""
+ monkeypatch.setenv("BEDROCK_RATE_LIMIT_PER_MINUTE", "1")
+ hint_body = {"word": "cat", "theme": "animals", "attempt_number": 1}
+ with patch("agent.hint_generator.boto3.client") as mock_client:
+ mock_client.side_effect = Exception("no AWS in tests")
+ first = client.post("/api/v1/hint", json=hint_body)
+ second = client.post("/api/v1/hint", json=hint_body)
+ assert first.status_code == 200
+ assert second.status_code == 429
+ assert int(second.headers["Retry-After"]) >= 1
+
+ def test_hint_attempts_without_bedrock_do_not_consume_allowance(
+ self, client, monkeypatch
+ ):
+ """Only attempt-1 hints consult Bedrock; later attempts are
+ deterministic reveals and must not burn rate-limit allowance."""
+ monkeypatch.setenv("BEDROCK_RATE_LIMIT_PER_MINUTE", "1")
+ for attempt in (2, 3, 4):
+ response = client.post("/api/v1/hint", json={
+ "word": "cat", "theme": "animals", "attempt_number": attempt,
+ "use_bedrock": True,
+ })
+ assert response.status_code == 200
+
+
+# ── HTTP surface: global budget hard cutover ────────────────────────────────
+
+class TestGlobalBudgetCutoverHTTP:
+ def test_budget_exhaustion_forces_template_only_fallback(
+ self, client, mock_bedrock_story, monkeypatch, mock_budget_logger
+ ):
+ """Simulated budget exhaustion: the last allowed call still gets a
+ generated story; every later request silently degrades to the
+ template until the window resets."""
+ monkeypatch.setenv("BEDROCK_DAILY_BUDGET", "1")
+ from agent.profiler import load_profile
+ load_profile("student_001", consent_metadata=CONSENT_METADATA)
+
+ body = {
+ "student_id": "student_001",
+ "words": ["cat", "hat"],
+ "use_bedrock": True,
+ }
+ within_budget = client.post("/api/v1/story", json=body, headers=auth())
+ assert within_budget.status_code == 200
+ assert json.loads(BEDROCK_STORY_TEXT)["story"][:10] in within_budget.json()["story"]
+
+ exhausted = client.post("/api/v1/story", json=body, headers=auth())
+ assert exhausted.status_code == 200
+ assert TEMPLATE_STORY_MARKER in exhausted.json()["story"]
+ assert exhausted.json()["story"] != json.loads(BEDROCK_STORY_TEXT)["story"]
+
+ # The provider must not have been invoked again.
+ assert mock_bedrock_story.return_value.invoke_model.call_count == 1
+
+ def test_cutover_applies_to_hints_and_logs_once(
+ self, client, monkeypatch, mock_budget_logger
+ ):
+ monkeypatch.setenv("BEDROCK_DAILY_BUDGET", "1")
+ with patch("agent.hint_generator.boto3.client") as mock_client:
+ mock_client.side_effect = AssertionError("Bedrock must not be called after cutover")
+ first = client.post("/api/v1/hint", json={
+ "word": "cat", "theme": "animals", "attempt_number": 1,
+ })
+ second = client.post("/api/v1/hint", json={
+ "word": "hat", "theme": "objects", "attempt_number": 1,
+ })
+ third = client.post("/api/v1/hint", json={
+ "word": "bat", "theme": "animals", "attempt_number": 1,
+ })
+ assert first.status_code == 200
+ assert second.status_code == 200
+ assert third.status_code == 200
+ exhausted_logs = [
+ c for c in mock_budget_logger.warning.call_args_list
+ if c.kwargs.get("extra", {}).get("outcome") == "budget_exhausted"
+ ]
+ assert len(exhausted_logs) == 1
+
+ def test_budget_recovery_restores_generation(
+ self, client, mock_bedrock_story, monkeypatch
+ ):
+ from agent import bedrock_guardrails as guardrails
+ from agent.profiler import load_profile
+
+ monkeypatch.setenv("BEDROCK_DAILY_BUDGET", "1")
+ load_profile("student_001", consent_metadata=CONSENT_METADATA)
+ body = {
+ "student_id": "student_001",
+ "words": ["cat", "hat"],
+ "use_bedrock": True,
+ }
+ assert TEMPLATE_STORY_MARKER not in client.post(
+ "/api/v1/story", json=body, headers=auth(),
+ ).json()["story"]
+ assert TEMPLATE_STORY_MARKER in client.post(
+ "/api/v1/story", json=body, headers=auth(),
+ ).json()["story"]
+
+ # Window resets (next UTC day): generated stories resume.
+ monkeypatch.setattr(guardrails, "_utc_day", lambda: "2099-01-02")
+ assert TEMPLATE_STORY_MARKER not in client.post(
+ "/api/v1/story", json=body, headers=auth(),
+ ).json()["story"]
+
+
+# ── Admin observability endpoint ────────────────────────────────────────────
+
+class TestAdminBedrockUsageEndpoint:
+ def test_requires_authentication(self, client):
+ assert client.get("/api/v1/admin/bedrock-usage").status_code == 401
+
+ def test_rejects_non_admin_roles(self, client):
+ assert client.get(
+ "/api/v1/admin/bedrock-usage", headers=auth(PARENT_KEY),
+ ).status_code == 403
+
+ def test_reports_counters_limits_and_flags(
+ self, client, mock_bedrock_story, monkeypatch
+ ):
+ monkeypatch.setenv("BEDROCK_RATE_LIMIT_PER_MINUTE", "7")
+ monkeypatch.setenv("BEDROCK_DAILY_BUDGET", "100")
+ monkeypatch.setenv("BEDROCK_MONTHLY_BUDGET", "5000")
+ from agent.profiler import load_profile
+ load_profile("student_001", consent_metadata=CONSENT_METADATA)
+ body = {
+ "student_id": "student_001",
+ "words": ["cat", "hat"],
+ "use_bedrock": True,
+ }
+ for _ in range(3):
+ client.post("/api/v1/story", json=body, headers=auth())
+
+ response = client.get("/api/v1/admin/bedrock-usage", headers=auth(ADMIN_KEY))
+ assert response.status_code == 200
+ payload = response.json()
+ assert payload["daily"]["used"] == 3
+ assert payload["daily"]["limit"] == 100
+ assert payload["monthly"]["used"] == 3
+ assert payload["monthly"]["limit"] == 5000
+ assert payload["rate_limit"] == {"per_minute": 7}
+ assert payload["budget_exhausted"] is False
+ assert payload["tracked_principals"] >= 1
+
+ def test_reports_exhausted_budget_and_unlimited_as_null(
+ self, client, mock_bedrock_story, monkeypatch
+ ):
+ monkeypatch.setenv("BEDROCK_DAILY_BUDGET", "1")
+ monkeypatch.setenv("BEDROCK_MONTHLY_BUDGET", "-1")
+ from agent.profiler import load_profile
+ load_profile("student_001", consent_metadata=CONSENT_METADATA)
+ body = {
+ "student_id": "student_001",
+ "words": ["cat"],
+ "use_bedrock": True,
+ }
+ client.post("/api/v1/story", json=body, headers=auth())
+ client.post("/api/v1/story", json=body, headers=auth())
+
+ payload = client.get(
+ "/api/v1/admin/bedrock-usage", headers=auth(ADMIN_KEY),
+ ).json()
+ assert payload["budget_exhausted"] is True
+ assert payload["daily"]["used"] == 1
+ assert payload["monthly"]["limit"] is None
+
+ def test_snapshot_never_contains_identifiers(
+ self, client, mock_bedrock_story
+ ):
+ """Counters are aggregates: neither student IDs nor raw principals
+ may appear anywhere in the admin payload."""
+ from agent.profiler import load_profile
+ load_profile("student_001", consent_metadata=CONSENT_METADATA)
+ client.post("/api/v1/story", json={
+ "student_id": "student_001", "words": ["cat"], "use_bedrock": True,
+ }, headers=auth())
+
+ raw = client.get(
+ "/api/v1/admin/bedrock-usage", headers=auth(ADMIN_KEY),
+ ).text
+ assert "student_001" not in raw
+ assert "guardrail_parent" not in raw
+ assert PARENT_KEY not in raw