From 29b2bc0e6fcd919fdc6afb9bccfe01cc4faa941f Mon Sep 17 00:00:00 2001 From: Amiel Peled Date: Thu, 20 Aug 2026 22:23:07 +0300 Subject: [PATCH 1/3] fix(server): identify the caller behind the proxy, not the proxy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Anonymous quotas were being charged to `request.client.host`, which behind a proxy is the proxy. On a platform that fronts the app with several edge pods it is a *different* proxy from one request to the next, so the effects compounded: visitors shared whichever bucket their request landed in, and the daily allowance was not enforced at all because the next request landed elsewhere. Visible from outside — five consecutive anonymous requests to the deployment returned X-AI-Quota-Remaining of 8, 9, 9, 8, 9 rather than counting down. It is also what produced the reported symptom of a 429 on the first click and success on the second: the first landed in a bucket another visitor had filled, the retry landed somewhere else. The caller now comes from X-Forwarded-For, read from the right. That direction is the whole security argument: anything a client writes itself lands at the left, so only the entries our own infrastructure appended are worth reading. TRUSTED_PROXY_HOPS says how many those are — one by default, matching a single load balancer that appends the caller and then itself. Fly's own header still wins where it is present, and outside production nothing is trusted at all. Ten tests cover it, including the two that state the bug: two callers behind one proxy must be told apart, and one caller through two different edges must not. 157 pytest, 58 api, 72 contract pass. Co-Authored-By: Claude Opus 5 (1M context) --- replit.md | 1 + server/app/dependencies.py | 35 +++++++- server/app/settings.py | 9 ++ server/tests/test_client_identity.py | 125 +++++++++++++++++++++++++++ 4 files changed, 167 insertions(+), 3 deletions(-) create mode 100644 server/tests/test_client_identity.py diff --git a/replit.md b/replit.md index c31319a..30ed807 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 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/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_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 From 3c38abf9fe4473eae77c7f0880a5e29b21b9c159 Mon Sep 17 00:00:00 2001 From: Amiel Peled Date: Thu, 20 Aug 2026 22:37:24 +0300 Subject: [PATCH 2/3] fix(server): do not charge a quota for a request the provider refused MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reproduced against the deployment: a small request returns 200, and a large one — the shape the résumé rewrite sends — returns 429 carrying "The AI provider could not complete this request". That is not our limiter. It is the shared Groq key meeting Groq's own tokens-per-minute ceiling, passed through with the upstream status. The daily allowance was spent before the call and kept whatever came back, so a visitor with ten requests a day could burn the lot on refusals and receive nothing. SharedRateLimiter can now give a hit back, and the AI route does so whenever the outcome is not a 200, correcting X-AI-Quota-Remaining with it. Retrying is still bounded by the burst limiter, so handing one back cannot become a way around the quota. Checked by mutation — removing the release from the route fails the test that two refusals in a row leave the allowance where it started, and the paired test keeps an answered request being charged for. The client wording is widened to match. It said the free allowance was spent, which is only one of the two causes and the wrong one here; it now names both and gives the step that works either way. 161 pytest, 246 unit, 58 api, 72 contract pass. Co-Authored-By: Claude Opus 5 (1M context) --- .../ai-testing-academy/src/lib/locales/en.ts | 2 +- .../ai-testing-academy/src/lib/locales/he.ts | 2 +- server/app/database.py | 16 ++++ server/app/rate_limit.py | 31 ++++++- server/app/routes/ai.py | 5 + server/tests/test_rate_limit_config.py | 91 +++++++++++++++++++ 6 files changed, 144 insertions(+), 3 deletions(-) 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/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/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..44f46ea 100644 --- a/server/app/routes/ai.py +++ b/server/app/routes/ai.py @@ -50,6 +50,11 @@ async def ai_generate(request: Request, ai: Ai, session: SessionUser): provider, model = ai.target(body) outcome = await ai.generate(body) + 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) observe_ai(request, provider=provider, model=model, email=email, status=outcome.status) return JSONResponse(outcome.payload, status_code=outcome.status, headers=headers) 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"]) From 4e8f82b80384ae49ee46ccef27e7687c770e08d8 Mon Sep 17 00:00:00 2001 From: Amiel Peled Date: Thu, 20 Aug 2026 22:44:01 +0300 Subject: [PATCH 3/3] feat(server): fall through to the next AI provider when one refuses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Groq's free tier refuses a large request on tokens-per-minute, which is what made the résumé rewrite fail while short prompts succeeded. Gemini is configured on the same deployment and can answer it, so an ordinary request now tries it rather than giving up. The boundaries are the interesting part, and each has a test: - A grounded request has no fallback. Search grounding is the reason it chose that provider, and answering without it answers a different question. - A 4xx below 429 is not retried. That is the request being wrong, and the second provider would only agree, one wasted call later. - A fallback with no key configured leaves the original failure showing rather than replacing it with its own 503, which would hide why the request failed. - The outcome now carries the provider and model that actually answered, and the route labels metrics with those. Otherwise a fallback would be recorded as a Groq success. One thing had to change for this to be honest: Gemini's request builder added the google_search tool unconditionally, which was correct while it only ever served grounded requests. Serving an ordinary one as a fallback would have turned it into a search. Grounding now follows the request, not the provider. That last bug was real and I nearly shipped it — a scripted edit failed to match and the change silently did not apply, because I had asserted on the batch rather than on that replacement. The test caught it. 168 pytest, 58 api, 72 contract pass. Co-Authored-By: Claude Opus 5 (1M context) --- replit.md | 4 + server/app/ai_gateway.py | 105 ++++++++++++++++++++------ server/app/routes/ai.py | 11 ++- server/tests/test_ai_gateway.py | 126 ++++++++++++++++++++++++++++++++ 4 files changed, 222 insertions(+), 24 deletions(-) diff --git a/replit.md b/replit.md index 30ed807..7c8b8f8 100644 --- a/replit.md +++ b/replit.md @@ -90,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/routes/ai.py b/server/app/routes/ai.py index 44f46ea..eb98ced 100644 --- a/server/app/routes/ai.py +++ b/server/app/routes/ai.py @@ -48,14 +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) 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) - observe_ai(request, provider=provider, model=model, email=email, status=outcome.status) + # 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/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])