diff --git a/artifacts/ai-testing-academy/src/lib/locales/en.ts b/artifacts/ai-testing-academy/src/lib/locales/en.ts index 30b885f..7136eee 100644 --- a/artifacts/ai-testing-academy/src/lib/locales/en.ts +++ b/artifacts/ai-testing-academy/src/lib/locales/en.ts @@ -236,7 +236,7 @@ export const en = { errBlockedOpenUrl: '', errApiPrefix: 'API error (', errProxyBusy: - 'The free AI allowance is used up for now. Wait a little, or connect your own provider key in Settings to keep going.', + 'The shared AI service is busy, or its free allowance is spent. Wait a moment and try again, or connect your own provider key in Settings to keep going.', errProxyUnavailable: 'The academy’s own AI key is unavailable right now. Connect your own provider key in Settings to keep going.', errNoJson: 'Could not parse JSON from the model response. Please try again.', diff --git a/artifacts/ai-testing-academy/src/lib/locales/he.ts b/artifacts/ai-testing-academy/src/lib/locales/he.ts index c5e1427..6859833 100644 --- a/artifacts/ai-testing-academy/src/lib/locales/he.ts +++ b/artifacts/ai-testing-academy/src/lib/locales/he.ts @@ -232,7 +232,7 @@ export const he: Locale = { errBlockedOpenUrl: '', errApiPrefix: 'שגיאת API (', errProxyBusy: - 'מכסת ה-AI החינמית נוצלה כרגע. המתינו מעט, או חברו מפתח ספק משלכם בהגדרות כדי להמשיך.', + 'שירות ה-AI המשותף עמוס כרגע, או שהמכסה החינמית נוצלה. המתינו רגע ונסו שוב, או חברו מפתח ספק משלכם בהגדרות כדי להמשיך.', errProxyUnavailable: 'מפתח ה-AI של האקדמיה אינו זמין כרגע. חברו מפתח ספק משלכם בהגדרות כדי להמשיך.', errNoJson: 'לא ניתן לנתח JSON מתגובת המודל. נסה שוב.', diff --git a/replit.md b/replit.md index c31319a..7c8b8f8 100644 --- a/replit.md +++ b/replit.md @@ -63,6 +63,7 @@ is absent fails closed with a controlled 4xx/5xx response. The package scripts d | `UPSTREAM_API_BASE_URL` | Secretless Replit API artifact relays `/api/*` to Fly | FastAPI handles routes locally | | `METRICS_TOKEN` | Protects production `/metrics` scrapes | Metrics are available only outside production | | `METRICS_ID_SALT` | HMAC-pseudonymizes user labels in metrics | Authenticated users are labeled `redacted` | +| `TRUSTED_PROXY_HOPS` | How many rightmost `X-Forwarded-For` entries the platform appends, so the caller can be read past them | Defaults to 1. Too high trusts an entry the caller forged; too low keys everyone behind one proxy to the same quota | | `RATE_LIMIT_SALT` | Production quotas are counted in Postgres | **In production every rate-limited route refuses every caller with 429** — sign-in, the AI proxy and the admin seed route. `METRICS_ID_SALT` is accepted instead. `/api/readyz` names it in a `rateLimiting` field | ### API structure @@ -89,6 +90,10 @@ Two consequences worth knowing before editing: - **Adding an AI provider** is a class in `ai_gateway.py` plus an entry in `PROVIDERS`. Dispatch, `/api/ai/config` and the routes all follow from the registry. +- **An ordinary request falls through to the next provider** when the first answers 429 or + 5xx — a provider saying "not me, not now". A 4xx below that is the request being wrong, so + it is not retried elsewhere. A grounded request has no fallback: search grounding is why it + chose that provider. The outcome names whoever answered, so metrics stay truthful. - **Tests substitute collaborators through `app.dependency_overrides`**, not by patching module globals. `server/tests/conftest.py` exposes `override_dependency` for this. diff --git a/server/app/ai_gateway.py b/server/app/ai_gateway.py index b35d0ce..61e4132 100644 --- a/server/app/ai_gateway.py +++ b/server/app/ai_gateway.py @@ -35,12 +35,21 @@ class Completion: truncated: bool +UNKNOWN = "unknown" + + @dataclass(frozen=True) class AiOutcome: - """A proxied result: the upstream status travels with the body deliberately.""" + """A proxied result: the upstream status travels with the body deliberately. + + `provider` and `model` name whoever actually answered, which after a + fallback is not the one the request started with. + """ status: int payload: dict + provider: str = UNKNOWN + model: str = UNKNOWN class AiProvider(Protocol): @@ -122,23 +131,29 @@ class GeminiProvider(_Provider): supports_grounding = True def build_request(self, model: str, key: str, body: GenerateBody) -> UpstreamRequest: + payload: dict = { + "system_instruction": {"parts": [{"text": body.system}]}, + "contents": [ + { + "role": "model" if message.role == "assistant" else "user", + "parts": [{"text": message.content}], + } + for message in body.messages + ], + "generationConfig": {"maxOutputTokens": body.maxTokens}, + } + # Search grounding is what the caller asked for, not what this provider + # is. It was unconditional while Gemini only ever served grounded + # requests; as a fallback it also serves ordinary ones, and those must + # not quietly become searches. + if body.grounded: + payload["tools"] = [{"google_search": {}}] return UpstreamRequest( url=( f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent" ), headers={"content-type": "application/json", "x-goog-api-key": key}, - payload={ - "system_instruction": {"parts": [{"text": body.system}]}, - "contents": [ - { - "role": "model" if message.role == "assistant" else "user", - "parts": [{"text": message.content}], - } - for message in body.messages - ], - "generationConfig": {"maxOutputTokens": body.maxTokens}, - "tools": [{"google_search": {}}], - }, + payload=payload, ) def read_completion(self, data: dict) -> Completion: @@ -179,12 +194,49 @@ def advertised_config(self) -> dict[str, dict[str, object]]: for provider in self._providers } + def candidates(self, body: GenerateBody) -> tuple[_Provider, ...]: + """The provider for this request, then any other that could also serve it. + + A grounded request has no alternative: search grounding is the reason it + chose that provider, and a provider without it would answer a different + question. An ordinary request can be served by any of them. + """ + primary = self.provider_for(body) + if body.grounded: + return (primary,) + return (primary, *(p for p in self._providers if p is not primary)) + + @staticmethod + def _worth_another_provider(status: int) -> bool: + """Whether the refusal was about this provider rather than the request. + + A rate limit or a fault is the provider saying "not me, not now", and + another one may well answer. A 4xx below that is the request itself + being wrong, and every provider will say the same thing — retrying it + elsewhere only spends a second call to be told so twice. + """ + return status == 429 or status >= 500 + async def generate(self, body: GenerateBody) -> AiOutcome: - provider = self.provider_for(body) + outcome: AiOutcome | None = None + for provider in self.candidates(body): + attempt = await self._attempt(provider, body) + if attempt.status == 200 or not self._worth_another_provider(attempt.status): + return attempt + outcome = outcome or attempt + logger.warning( + "%s could not serve the request (status %s); trying the next provider", + provider.name, + attempt.status, + ) + return outcome or AiOutcome(503, {"error": "No AI provider is configured."}) + + async def _attempt(self, provider: _Provider, body: GenerateBody) -> AiOutcome: + model = provider.resolve_model(body.model) key = provider.api_key() if not key: - return AiOutcome(503, {"error": provider.missing_key_message()}) - request = provider.build_request(provider.resolve_model(body.model), key, body) + return AiOutcome(503, {"error": provider.missing_key_message()}, provider.name, model) + request = provider.build_request(model, key, body) try: async with httpx.AsyncClient(timeout=UPSTREAM_TIMEOUT) as client: upstream = await client.post( @@ -192,23 +244,32 @@ async def generate(self, body: GenerateBody) -> AiOutcome: ) data = upstream.json() if not upstream.is_success: - return self._refused(upstream.status_code) + return self._refused(upstream.status_code, provider.name, model) completion = provider.read_completion(data) - return AiOutcome(200, {"text": completion.text, "truncated": completion.truncated}) + return AiOutcome( + 200, + {"text": completion.text, "truncated": completion.truncated}, + provider.name, + model, + ) except httpx.TimeoutException: - return AiOutcome(504, {"error": "Provider request timed out"}) + return AiOutcome(504, {"error": "Provider request timed out"}, provider.name, model) except Exception: logger.exception("AI provider request failed") - return AiOutcome(502, {"error": "Failed to reach provider"}) + return AiOutcome(502, {"error": "Failed to reach provider"}, provider.name, model) - def _refused(self, status: int) -> AiOutcome: + def _refused(self, status: int, provider: str, model: str) -> AiOutcome: """Upstream detail stays in the log; the caller gets a correlation id.""" request_id = secrets.token_urlsafe(8) - logger.warning("AI provider refused request id=%s status=%s", request_id, status) + logger.warning( + "AI provider %s refused request id=%s status=%s", provider, request_id, status + ) return AiOutcome( status, { "error": "The AI provider could not complete this request.", "requestId": request_id, }, + provider, + model, ) diff --git a/server/app/database.py b/server/app/database.py index 305594c..49532ab 100644 --- a/server/app/database.py +++ b/server/app/database.py @@ -127,6 +127,22 @@ def _hit_rate_limit( return hits <= limit, max(0, limit - hits) +async def release_rate_limit(bucket: str, key_hash: str) -> int: + """Give back one hit, and report what is left. Never goes below zero.""" + return await asyncio.to_thread(_release_rate_limit, bucket, key_hash) + + +def _release_rate_limit(bucket: str, key_hash: str) -> int: + with psycopg.connect(database_url()) as connection: + row = connection.execute( + """UPDATE api_rate_limits SET hits = GREATEST(0, hits - 1) + WHERE bucket = %s AND key_hash = %s + RETURNING hits""", + (bucket, key_hash), + ).fetchone() + return int(row[0]) if row else 0 + + async def find_course_access( subject: str | None, email: str | None, diff --git a/server/app/dependencies.py b/server/app/dependencies.py index 2bda431..a5cc922 100644 --- a/server/app/dependencies.py +++ b/server/app/dependencies.py @@ -27,7 +27,7 @@ from .google_auth import GoogleUser, verify_google_id_token from .rate_limit import SharedRateLimiter from .sessions import read_session -from .settings import BURST_LIMIT, BURST_WINDOW, DAILY_QUOTA +from .settings import BURST_LIMIT, BURST_WINDOW, DAILY_QUOTA, TRUSTED_PROXY_HOPS logger = logging.getLogger(__name__) @@ -51,12 +51,41 @@ def bearer_token(authorization: str | None) -> str: return token.strip() if separator and scheme.lower() == "bearer" else "" +def _forwarded_for(request: Request) -> str | None: + """The caller's address from X-Forwarded-For, counting from the right. + + Everything a client sends itself lands at the left of this header, so the + only entries worth reading are the ones our own infrastructure appended. + `TRUSTED_PROXY_HOPS` says how many of those there are; a list too short to + have that many falls back to its leftmost entry, which is the closest thing + to the caller that exists. + """ + entries = [ + value.strip() + for value in request.headers.get("x-forwarded-for", "").split(",") + if value.strip() + ] + if not entries: + return None + index = len(entries) - 1 - TRUSTED_PROXY_HOPS + return entries[index] if index >= 0 else entries[0] + + def client_ip(request: Request) -> str: - # Fly terminates TLS and supplies this header itself. It is accepted only in - # production; local/test callers cannot forge a different quota identity. + """Who to charge a quota to, when there is no signed-in identity to use. + + Read from a forwarded header only in production, where a proxy we control is + the one writing it; a local caller cannot forge a different quota identity. + The socket address is the last resort and behind a proxy it names the proxy, + not the caller — which is how anonymous callers ended up sharing, and + shuffling between, each other's quota buckets. + """ if os.getenv("NODE_ENV") == "production": + # Fly terminates TLS and supplies this header itself. if forwarded := request.headers.get("fly-client-ip", "").strip(): return forwarded + if forwarded_for := _forwarded_for(request): + return forwarded_for return request.client.host if request.client else "unknown" diff --git a/server/app/rate_limit.py b/server/app/rate_limit.py index 7c15cfc..996aad4 100644 --- a/server/app/rate_limit.py +++ b/server/app/rate_limit.py @@ -10,7 +10,7 @@ from typing import Literal from .config import database_url, env -from .database import hit_rate_limit +from .database import hit_rate_limit, release_rate_limit logger = logging.getLogger(__name__) @@ -50,6 +50,13 @@ async def hit(self, key: str) -> tuple[bool, int]: hits.append(now) return True, self.limit - len(hits) + async def release(self, key: str) -> int: + async with self._lock: + hits = self._hits[key] + if hits: + hits.pop() + return max(0, self.limit - len(hits)) + WhenUnavailable = Literal["refuse", "degrade"] @@ -88,6 +95,28 @@ def __init__( self.memory = MemoryRateLimiter(limit, window_seconds) self._warned: set[str] = set() + async def release(self, key: str) -> int | None: + """Give back a hit that bought the caller nothing, or None if it could not. + + A quota exists to ration what the caller receives. Charging for a request + the provider refused spends an allowance on nothing, and with ten a day + that is most of a visit. Retrying is still bounded by the burst limiter, + so giving one back cannot become a way around the quota. + """ + if os.getenv("NODE_ENV") != "production": + return await self.memory.release(key) + if shared_quota_problem(): + return None + salt = env("RATE_LIMIT_SALT") or env("METRICS_ID_SALT") + assert salt is not None + digest = hmac.new(salt.encode(), key.encode(), hashlib.sha256).hexdigest() + try: + hits = await release_rate_limit(self.bucket, digest) + except Exception: + logger.exception("Could not release a rate limit hit for bucket %r", self.bucket) + return None + return max(0, self.limit - hits) + async def hit(self, key: str) -> tuple[bool, int]: if os.getenv("NODE_ENV") != "production": return await self.memory.hit(key) diff --git a/server/app/routes/ai.py b/server/app/routes/ai.py index d1a553d..eb98ced 100644 --- a/server/app/routes/ai.py +++ b/server/app/routes/ai.py @@ -48,9 +48,21 @@ async def ai_generate(request: Request, ai: Ai, session: SessionUser): issues = validation_issues(exc.errors()) if isinstance(exc, ValidationError) else [] return error_response("Invalid request body", 400, issues=issues, headers=headers) - provider, model = ai.target(body) outcome = await ai.generate(body) - observe_ai(request, provider=provider, model=model, email=email, status=outcome.status) + if outcome.status != 200: + # The caller received nothing, so the allowance they just spent buys + # them nothing either. Retrying stays bounded by the burst limiter. + if (remaining := await daily_limiter.release(key)) is not None: + headers["X-AI-Quota-Remaining"] = str(remaining) + # Labelled with whoever actually answered, which after a fallback is not the + # provider the request started with. + observe_ai( + request, + provider=outcome.provider, + model=outcome.model, + email=email, + status=outcome.status, + ) return JSONResponse(outcome.payload, status_code=outcome.status, headers=headers) diff --git a/server/app/settings.py b/server/app/settings.py index b92c58b..93ee964 100644 --- a/server/app/settings.py +++ b/server/app/settings.py @@ -15,5 +15,14 @@ UPSTREAM_TIMEOUT = positive_int("AI_UPSTREAM_TIMEOUT_MS", 30_000) / 1000 MAX_REQUEST_BODY = positive_int("MAX_REQUEST_BODY_BYTES", 96 * 1024) +# How many rightmost X-Forwarded-For entries are written by infrastructure we +# control and must therefore be skipped to reach the caller. A client can put +# anything in that header, but only ever at the left, so counting from the right +# is what makes the value trustworthy. One matches a single load balancer that +# appends the caller and then itself, which is the common cloud shape; set it to +# match the deployment, because too high trusts a forged entry and too low keys +# every caller behind one proxy to the same bucket. +TRUSTED_PROXY_HOPS = positive_int("TRUSTED_PROXY_HOPS", 1) + COURSE_SKU = "ai-testing-bootcamp" TERMS_VERSION = "2026-08-20" diff --git a/server/tests/test_ai_gateway.py b/server/tests/test_ai_gateway.py index b39fc1a..91f0cd3 100644 --- a/server/tests/test_ai_gateway.py +++ b/server/tests/test_ai_gateway.py @@ -281,3 +281,129 @@ def _json_body(request: httpx.Request) -> dict: import json return json.loads(request.content) + + +@pytest.fixture +def upstreams(monkeypatch: pytest.MonkeyPatch): + """Answer each provider's host differently, and record who was called.""" + + def install(*, groq: httpx.Response | None, gemini: httpx.Response | None) -> list[str]: + called: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + name = "groq" if "groq" in request.url.host else "gemini" + called.append(name) + reply = groq if name == "groq" else gemini + if reply is None: + raise AssertionError(f"{name} should not have been called") + return httpx.Response(reply.status_code, content=reply.content, request=request) + + original = httpx.AsyncClient + monkeypatch.setattr( + "app.ai_gateway.httpx.AsyncClient", + lambda **kwargs: original(transport=httpx.MockTransport(handler), **kwargs), + ) + return called + + return install + + +@pytest.fixture +def both_keys(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("GROQ_API_KEY", "gsk_fixture") + monkeypatch.setenv("GEMINI_API_KEY", "AIza_fixture") + + +def gemini_ok() -> httpx.Response: + return httpx.Response( + 200, json={"candidates": [{"content": {"parts": [{"text": "from gemini"}]}}]} + ) + + +@pytest.mark.asyncio +async def test_a_rate_limited_provider_hands_the_request_to_the_other(both_keys, upstreams) -> None: + """The reported case: Groq refuses a large request on tokens-per-minute.""" + called = upstreams(groq=httpx.Response(429, json={"error": "rate"}), gemini=gemini_ok()) + + outcome = await AiGateway().generate(body()) + + assert outcome.status == 200 + assert outcome.payload["text"] == "from gemini" + assert called == ["groq", "gemini"] + + +@pytest.mark.asyncio +async def test_the_outcome_names_the_provider_that_actually_answered(both_keys, upstreams) -> None: + """Metrics label the real one, or a fallback would look like a Groq success.""" + upstreams(groq=httpx.Response(503, json={"error": "down"}), gemini=gemini_ok()) + + outcome = await AiGateway().generate(body()) + + assert outcome.provider == "gemini" + assert outcome.model in GeminiProvider.models + + +@pytest.mark.asyncio +async def test_a_bad_request_is_not_offered_to_a_second_provider(both_keys, upstreams) -> None: + """400 is the request being wrong; the other provider would only agree.""" + called = upstreams(groq=httpx.Response(400, json={"error": "malformed"}), gemini=None) + + outcome = await AiGateway().generate(body()) + + assert outcome.status == 400 + assert called == ["groq"] + + +@pytest.mark.asyncio +async def test_a_grounded_request_is_never_handed_to_a_provider_without_search( + both_keys, upstreams +) -> None: + """Answering it without grounding would answer a different question.""" + called = upstreams(groq=None, gemini=httpx.Response(429, json={"error": "rate"})) + + outcome = await AiGateway().generate(body(grounded=True)) + + assert outcome.status == 429 + assert called == ["gemini"] + + +@pytest.mark.asyncio +async def test_an_unconfigured_fallback_leaves_the_original_failure_showing( + monkeypatch: pytest.MonkeyPatch, upstreams +) -> None: + """Reporting the fallback's missing key would hide why the request failed.""" + monkeypatch.setenv("GROQ_API_KEY", "gsk_fixture") + monkeypatch.delenv("GEMINI_API_KEY", raising=False) + upstreams(groq=httpx.Response(429, json={"error": "rate"}), gemini=None) + + outcome = await AiGateway().generate(body()) + + assert outcome.status == 429 + + +@pytest.mark.asyncio +async def test_when_every_provider_refuses_the_first_refusal_is_reported( + both_keys, upstreams +) -> None: + called = upstreams( + groq=httpx.Response(429, json={"error": "rate"}), + gemini=httpx.Response(500, json={"error": "boom"}), + ) + + outcome = await AiGateway().generate(body()) + + assert outcome.status == 429 + assert called == ["groq", "gemini"] + + +@pytest.mark.asyncio +async def test_gemini_only_searches_when_the_request_asked_it_to( + both_keys, monkeypatch: pytest.MonkeyPatch, upstream +) -> None: + """Serving an ordinary request as a fallback must not turn it into a search.""" + monkeypatch.delenv("GROQ_API_KEY", raising=False) + sent = upstream(httpx.Response(200, json={"candidates": [{"content": {"parts": []}}]})) + + await AiGateway().generate(body()) + + assert "tools" not in _json_body(sent[0]) diff --git a/server/tests/test_client_identity.py b/server/tests/test_client_identity.py new file mode 100644 index 0000000..bc84d11 --- /dev/null +++ b/server/tests/test_client_identity.py @@ -0,0 +1,125 @@ +"""Who an anonymous request is charged to. + +Behind a proxy, `request.client.host` names the proxy rather than the caller, +and on a platform that fronts the app with several edge pods it names a +*different* proxy from one request to the next. That is not a small +inaccuracy: it means anonymous callers share whichever bucket their request +lands in, so one visitor meets a 429 caused by another's usage, and the daily +allowance is not enforced at all because the next request lands elsewhere. It +was visible from outside as `X-AI-Quota-Remaining` bouncing 8, 9, 9, 8, 9 +across five consecutive requests instead of counting down. + +The header can be set by anyone, so the direction it is read from is the whole +security argument: forged entries are always on the left. +""" + +from __future__ import annotations + +import pytest +from starlette.datastructures import Headers +from starlette.requests import Request + +from app.dependencies import client_ip + +CALLER = "203.0.113.10" +LOAD_BALANCER = "10.0.0.1" +SOCKET = "192.0.2.99" + + +def request_with(**headers: str) -> Request: + raw = [(name.encode(), value.encode()) for name, value in headers.items()] + scope = { + "type": "http", + "method": "POST", + "path": "/api/ai/generate", + "headers": raw, + "scheme": "https", + "server": ("api.example", 443), + "query_string": b"", + "client": (SOCKET, 51000), + } + request = Request(scope) + request._headers = Headers(scope=scope) + return request + + +@pytest.fixture(autouse=True) +def production(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("NODE_ENV", "production") + + +def test_a_caller_behind_one_load_balancer_is_identified(monkeypatch) -> None: + """`, ` is the common cloud shape, and the default.""" + request = request_with(**{"x-forwarded-for": f"{CALLER}, {LOAD_BALANCER}"}) + + assert client_ip(request) == CALLER + + +def test_a_forged_entry_on_the_left_is_ignored() -> None: + """The entry a client wrote itself must never become its quota identity.""" + forged = "1.1.1.1" + request = request_with(**{"x-forwarded-for": f"{forged}, {CALLER}, {LOAD_BALANCER}"}) + + assert client_ip(request) == CALLER + assert client_ip(request) != forged + + +def test_two_callers_behind_the_same_proxy_are_told_apart() -> None: + """The bug this fixes: both used to key to the proxy, so they shared a quota.""" + other = "198.51.100.7" + first = request_with(**{"x-forwarded-for": f"{CALLER}, {LOAD_BALANCER}"}) + second = request_with(**{"x-forwarded-for": f"{other}, {LOAD_BALANCER}"}) + + assert client_ip(first) != client_ip(second) + + +def test_the_same_caller_through_a_different_edge_is_still_the_same_caller() -> None: + """Edge pods vary between requests; the caller must not vary with them.""" + first = request_with(**{"x-forwarded-for": f"{CALLER}, 10.0.0.1"}) + second = request_with(**{"x-forwarded-for": f"{CALLER}, 10.0.0.2"}) + + assert client_ip(first) == client_ip(second) + + +def test_a_single_entry_is_used_as_it_stands() -> None: + """A proxy that appends the caller without appending itself leaves one entry.""" + request = request_with(**{"x-forwarded-for": CALLER}) + + assert client_ip(request) == CALLER + + +def test_whitespace_and_empty_entries_do_not_shift_the_count() -> None: + request = request_with(**{"x-forwarded-for": f" {CALLER} , , {LOAD_BALANCER} "}) + + assert client_ip(request) == CALLER + + +def test_flys_own_header_still_wins_where_it_is_present() -> None: + """Fly writes this itself and it is not part of the forwarded chain.""" + request = request_with( + **{"fly-client-ip": CALLER, "x-forwarded-for": f"1.1.1.1, {LOAD_BALANCER}"} + ) + + assert client_ip(request) == CALLER + + +def test_no_forwarding_header_falls_back_to_the_socket() -> None: + assert client_ip(request_with()) == SOCKET + + +def test_a_local_run_never_trusts_the_header(monkeypatch: pytest.MonkeyPatch) -> None: + """Outside production nothing sits in front, so the header is only a claim.""" + monkeypatch.setenv("NODE_ENV", "development") + request = request_with(**{"x-forwarded-for": f"1.1.1.1, {LOAD_BALANCER}"}) + + assert client_ip(request) == SOCKET + + +def test_the_hop_count_is_configurable_for_a_deeper_chain( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Two proxies append two entries, and the default of one would read the wrong one.""" + monkeypatch.setattr("app.dependencies.TRUSTED_PROXY_HOPS", 2) + request = request_with(**{"x-forwarded-for": f"{CALLER}, {LOAD_BALANCER}, 10.0.0.2"}) + + assert client_ip(request) == CALLER diff --git a/server/tests/test_rate_limit_config.py b/server/tests/test_rate_limit_config.py index 3980402..b1c1224 100644 --- a/server/tests/test_rate_limit_config.py +++ b/server/tests/test_rate_limit_config.py @@ -210,3 +210,94 @@ async def test_the_ai_proxy_still_refuses_when_the_quota_cannot_count( ) assert response.status_code == 429 + + +@pytest.mark.asyncio +async def test_a_hit_can_be_given_back(production) -> None: + """Outside a shared store this is the in-memory path, which tests exercise.""" + production.setenv("NODE_ENV", "development") + limiter = SharedRateLimiter("ai-daily", 10, 86_400) + + await limiter.hit("ip:198.51.100.4") + remaining = await limiter.release("ip:198.51.100.4") + + assert remaining == 10 + + +@pytest.mark.asyncio +async def test_giving_back_more_than_was_taken_cannot_go_negative(production) -> None: + production.setenv("NODE_ENV", "development") + limiter = SharedRateLimiter("ai-daily", 10, 86_400) + + await limiter.release("ip:198.51.100.4") + await limiter.release("ip:198.51.100.4") + + assert (await limiter.hit("ip:198.51.100.4"))[1] == 9 + + +@pytest.mark.asyncio +async def test_a_refused_provider_does_not_spend_the_daily_allowance( + api_client, override_dependency +) -> None: + """The reported case: Groq refuses a large request, and ten a day is not many. + + A caller who receives nothing must not be charged for it. Retrying is still + bounded by the burst limiter, so handing the allowance back cannot become a + way around the quota. + """ + from app.ai_gateway import AiOutcome + from app.dependencies import get_ai_gateway + + class RefusingGateway: + def target(self, body): + return "groq", "openai/gpt-oss-120b" + + def advertised_config(self): + return {} + + async def generate(self, body) -> AiOutcome: + return AiOutcome(429, {"error": "The AI provider could not complete this request."}) + + override_dependency(get_ai_gateway, RefusingGateway) + + first = await api_client.post( + "/api/ai/generate", json={"messages": [{"role": "user", "content": "hi"}]} + ) + second = await api_client.post( + "/api/ai/generate", json={"messages": [{"role": "user", "content": "hi"}]} + ) + + assert first.status_code == 429 + assert first.headers["X-AI-Quota-Remaining"] == second.headers["X-AI-Quota-Remaining"], ( + "two refused requests in a row must leave the allowance where it started" + ) + + +@pytest.mark.asyncio +async def test_an_answered_request_does_spend_the_allowance( + api_client, override_dependency +) -> None: + """The other half: a request that returns text is charged for, as it should be.""" + from app.ai_gateway import AiOutcome + from app.dependencies import get_ai_gateway + + class AnsweringGateway: + def target(self, body): + return "groq", "openai/gpt-oss-120b" + + def advertised_config(self): + return {} + + async def generate(self, body) -> AiOutcome: + return AiOutcome(200, {"text": "answer", "truncated": False}) + + override_dependency(get_ai_gateway, AnsweringGateway) + + first = await api_client.post( + "/api/ai/generate", json={"messages": [{"role": "user", "content": "hi"}]} + ) + second = await api_client.post( + "/api/ai/generate", json={"messages": [{"role": "user", "content": "hi"}]} + ) + + assert int(second.headers["X-AI-Quota-Remaining"]) < int(first.headers["X-AI-Quota-Remaining"])