diff --git a/app/middleware/auth.py b/app/middleware/auth.py index 7b88ee2..1299af5 100644 --- a/app/middleware/auth.py +++ b/app/middleware/auth.py @@ -70,5 +70,8 @@ async def dispatch(self, request, call_next): "client_id": f"api_key:{user['api_key_id']}" if api_key else None, "permissions": user.get("permissions", []), "token_type": "api_key" if api_key else "user", + # M13: only ever True for an api_key identity (test mode is a + # property of the omni_sk_ key itself, not a session). + "test_mode": bool(user.get("test_mode", False)) if api_key else False, } return await call_next(request) diff --git a/app/routes/v1.py b/app/routes/v1.py index 84ac040..0b6a83e 100644 --- a/app/routes/v1.py +++ b/app/routes/v1.py @@ -16,7 +16,9 @@ - an org-level quota check maintained by omnibioai-billing (402), - optional Idempotency-Key replay, so a retried request is never run or billed twice, -- exactly one billable usage event per successful billable call, +- exactly one billable usage event per successful billable call -- except + for an omni_sk_test_ key, which gets a canned response instead, never + counted against quota or billed, - one error shape: {"error": {"type", "message", "request_id"}}. Request and response bodies of /v1/literature/answers are the frozen @@ -40,6 +42,8 @@ build_public_search, build_rag_query, build_rag_search_query, + build_test_answer, + build_test_search, ) from app.services.v1_store import V1Store, request_fingerprint @@ -154,8 +158,8 @@ def _upstream_error(status: int, response, request_id: str, headers: dict): async def _handle_billable_literature_call( - request: Request, *, resource: str, build_rag_body, build_public_response, quota_exceeded_message: str, - max_concurrent_answers: int | None = None, + request: Request, *, resource: str, build_rag_body, build_public_response, build_test_response, + quota_exceeded_message: str, max_concurrent_answers: int | None = None, ): """Shared lifecycle for every billable /v1/literature/* call (currently /answers and /search): auth'd-org check, contract @@ -165,9 +169,19 @@ async def _handle_billable_literature_call( their request/response, and their quota-exceeded wording -- every other step (in particular the idempotency/quota/usage sequencing) must stay identical between them, so it lives here once rather than - as two copies that could silently drift apart.""" + as two copies that could silently drift apart. + + An omni_sk_test_ key (identity.test_mode) short-circuits to + build_test_response's canned answer -- real rate limiting still + applies (abuse protection the gateway itself needs regardless of + whether a call is "real"), but quota, the concurrency cap, the real + RAG call, and usage emission are all skipped: a test key must never + consume real org quota, real RAG capacity, or be billed. + """ request_id = getattr(request.state, "trace_id", "") subject, org_id, user_id = _caller(request) + identity = getattr(request.state, "identity", None) or {} + test_mode = bool(identity.get("test_mode")) if not org_id: return _error(403, "organization_required", "This API is billed to an organization; your account has none.", request_id) @@ -206,7 +220,12 @@ async def _handle_billable_literature_call( return _error(409, "idempotency_in_progress", "A request with this Idempotency-Key is still running.", request_id, headers) - if max_concurrent_answers is not None: + # A test key never touches real RAG capacity, so it never needs (or + # holds) a concurrency slot -- acquiring one here, only to release it + # a few lines down having done no real work, would just be unearned + # contention against real callers sharing the same key/org budget. + acquire_concurrency = max_concurrent_answers is not None and not test_mode + if acquire_concurrency: if not await _acquire_concurrency_slots(org_id, subject, max_concurrent_answers): if idempotency_key is not None: await store.idempotency_finish(subject, idempotency_key, fingerprint, 429, None) @@ -217,6 +236,15 @@ async def _handle_billable_literature_call( ) try: + if test_mode: + public_response = build_test_response(domain=body.get("domain"), request_id=request_id, latency_ms=0) + if idempotency_key is not None: + await store.idempotency_finish(subject, idempotency_key, fingerprint, 200, public_response) + # No reserve_quota, no RAG forward, no emit_usage: a test key + # consumes no real quota, calls no real upstream, and is never + # billed -- that is the entire point of test mode. + return JSONResponse(public_response, status_code=200, headers={**headers, "X-Request-Id": request_id}) + # Atomic reserve-before-work: decrements the quota counter now, not # after the upstream call succeeds, so concurrent requests can # never all observe "quota available" and all succeed (see @@ -245,7 +273,6 @@ async def _handle_billable_literature_call( if idempotency_key is not None: await store.idempotency_finish(subject, idempotency_key, fingerprint, status, public_response) - identity = getattr(request.state, "identity", None) or {} await store.emit_usage( org_id=org_id, user_id=user_id, @@ -265,7 +292,7 @@ async def _handle_billable_literature_call( # denial, an upstream error, or success) -- a slot held by a # request that's already finished answering must never count # against the next one. - if max_concurrent_answers is not None: + if acquire_concurrency: await _release_concurrency_slots(org_id, subject) @@ -276,6 +303,7 @@ async def literature_answers(request: Request): resource=ANSWER_RESOURCE, build_rag_body=build_rag_query, build_public_response=build_public_answer, + build_test_response=build_test_answer, quota_exceeded_message="Your organization has used its included answers. " "Add a payment method or upgrade the plan.", max_concurrent_answers=Config.V1_MAX_CONCURRENT_ANSWERS, @@ -293,6 +321,7 @@ async def literature_search(request: Request): resource=SEARCH_RESOURCE, build_rag_body=build_rag_search_query, build_public_response=build_public_search, + build_test_response=build_test_search, quota_exceeded_message="Your organization has used its included searches. " "Add a payment method or upgrade the plan.", ) diff --git a/app/services/iam_client.py b/app/services/iam_client.py index 4eb7dc7..93db726 100644 --- a/app/services/iam_client.py +++ b/app/services/iam_client.py @@ -333,6 +333,11 @@ async def validate_api_key(self, api_key: str) -> Optional[dict]: "token_type": "api_key", "api_key_id": data["api_key_id"], "access_token": data["access_token"], + # M13 (design audit gap #9): an omni_sk_test_ key -- + # AuthMiddleware copies this onto request.state.identity, + # and app/routes/v1.py reads it there to serve a canned, + # unbilled response instead of forwarding to RAG. + "test_mode": bool(data.get("test_mode", False)), } ttl = min(Config.API_KEY_CACHE_TTL, int(data.get("expires_in", 0)) - 30) await self._set_cached_api_key(key_hash, user, ttl) diff --git a/app/services/literature_contract.py b/app/services/literature_contract.py index 183c66d..564592f 100644 --- a/app/services/literature_contract.py +++ b/app/services/literature_contract.py @@ -154,3 +154,45 @@ def build_public_answer(rag_result: dict, *, domain, request_id: str, latency_ms "latency_ms": latency_ms, }, } + + +# --------------------------------------------------------------------------- +# M13 (design audit gap #9): omni_sk_test_ keys get a canned, deterministic +# response in the exact shape build_public_answer/build_public_search +# produce -- never a real RAG call, never a real citation, never counted +# against quota or billed. Empty citations/results rather than fabricated +# literature data: a test key's whole point is exercising a caller's own +# integration code (idempotency handling, response parsing, error paths), +# not pretending to answer a real biomedical question. +# --------------------------------------------------------------------------- + + +def build_test_answer(*, domain, request_id: str, latency_ms: int) -> dict: + return { + "id": f"ans_test_{request_id or uuid.uuid4().hex}", + "answer": "This is a canned test-mode response. No literature service was called, " + "and this request was not billed.", + "citations": [], + "model": "test", + "model_source": "test", + "domain": domain or "default", + "usage": { + "queries": 0, + "input_tokens": None, + "output_tokens": None, + "billed_by": "query", + "latency_ms": latency_ms, + }, + } + + +def build_test_search(*, domain, request_id: str, latency_ms: int) -> dict: + return { + "id": f"srch_test_{request_id or uuid.uuid4().hex}", + "results": [], + "domain": domain or "default", + "usage": { + "searches": 0, + "latency_ms": latency_ms, + }, + } diff --git a/tests/test_api_key_auth.py b/tests/test_api_key_auth.py index 4f212df..72cfc86 100644 --- a/tests/test_api_key_auth.py +++ b/tests/test_api_key_auth.py @@ -110,6 +110,25 @@ def test_exchange_success_builds_identity_and_caches_by_hash(iam, exchange_secre http.post.assert_not_called() +def test_exchange_propagates_test_mode_true(iam, exchange_secret): + """M13: an omni_sk_test_ key's exchange response carries test_mode -- + AuthMiddleware reads it off this dict to build request.state.identity.""" + client, redis, http = iam + http.post.return_value = _response(200, {**EXCHANGE, "test_mode": True}) + user = asyncio.run(client.validate_api_key(KEY)) + assert user["test_mode"] is True + + +def test_exchange_defaults_test_mode_false_when_absent(iam, exchange_secret): + """A pre-M13 omnibioai-auth deployment's exchange response has no + test_mode field at all -- must default closed (not test mode), not + raise a KeyError.""" + client, redis, http = iam + http.post.return_value = _response(200, EXCHANGE) + user = asyncio.run(client.validate_api_key(KEY)) + assert user["test_mode"] is False + + def test_cache_ttl_never_outlives_minted_token(iam, exchange_secret): client, redis, http = iam http.post.return_value = _response(200, {**EXCHANGE, "expires_in": 20}) diff --git a/tests/test_v1_literature.py b/tests/test_v1_literature.py index bff4451..bd69778 100644 --- a/tests/test_v1_literature.py +++ b/tests/test_v1_literature.py @@ -434,6 +434,81 @@ def test_literature_search_has_no_concurrency_cap(client, redis, upstream, monke assert resp.status_code == 200 +# --------------------------------------------------------------------------- +# omni_sk_test_ keys: canned, unbilled responses (design audit gap #9's +# remaining "test keys and canned, unbilled responses are absent" bullet) +# --------------------------------------------------------------------------- + +TEST_MODE_USER = {**USER, "test_mode": True} + + +def test_test_mode_answers_returns_a_canned_response_without_calling_upstream(client, redis, upstream): + with patch.object(_main_mod.iam, "validate_api_key", AsyncMock(return_value=TEST_MODE_USER)): + resp = _post(client) + assert resp.status_code == 200 + body = resp.json() + assert body["id"].startswith("ans_test_") + assert body["model"] == "test" and body["model_source"] == "test" + assert body["citations"] == [] + assert body["usage"]["queries"] == 0 + upstream.assert_not_called() + assert _usage(redis) == [] # never billed + + +def test_test_mode_search_returns_a_canned_response_without_calling_upstream(client, redis, upstream): + headers = {"Authorization": f"Bearer {KEY}"} + with patch.object(_main_mod.iam, "validate_api_key", AsyncMock(return_value=TEST_MODE_USER)): + resp = client.post("/v1/literature/search", json={"question": "q"}, headers=headers) + assert resp.status_code == 200 + body = resp.json() + assert body["id"].startswith("srch_test_") + assert body["results"] == [] + assert body["usage"]["searches"] == 0 + upstream.assert_not_called() + assert _usage(redis) == [] + + +def test_test_mode_key_does_not_consume_quota(client, redis, upstream): + """A test key answers successfully even when the organization's real + quota is already exhausted -- test mode never checks it at all.""" + redis.kv["gateway:v1:quota:42:literature.answer"] = 0 + with patch.object(_main_mod.iam, "validate_api_key", AsyncMock(return_value=TEST_MODE_USER)): + resp = _post(client) + assert resp.status_code == 200 + + +def test_test_mode_key_is_still_rate_limited(client, redis, upstream, monkeypatch): + """Test mode skips quota/billing/RAG, but not rate limiting -- the + gateway's own resources still need abuse protection regardless of + whether a call is "real".""" + monkeypatch.setattr(Config, "V1_RATE_LIMIT_PER_MINUTE", 1) + with patch.object(_main_mod.iam, "validate_api_key", AsyncMock(return_value=TEST_MODE_USER)): + assert _post(client).status_code == 200 + assert _post(client).status_code == 429 + + +def test_test_mode_key_does_not_acquire_a_concurrency_slot(client, redis, upstream, monkeypatch): + """A test key never calls RAG, so it must not compete for -- or even + touch -- the real concurrency budget shared with live callers.""" + monkeypatch.setattr(Config, "V1_MAX_CONCURRENT_ANSWERS", 1) + redis.kv["gateway:v1:conc:api_key:7"] = 1 # already at the cap for a real call + + with patch.object(_main_mod.iam, "validate_api_key", AsyncMock(return_value=TEST_MODE_USER)): + resp = _post(client, idem="test-conc-1") + assert resp.status_code == 200 + assert redis.kv["gateway:v1:conc:api_key:7"] == 1 # untouched + + +def test_test_mode_key_supports_idempotency_replay(client, redis, upstream): + with patch.object(_main_mod.iam, "validate_api_key", AsyncMock(return_value=TEST_MODE_USER)): + first = _post(client, idem="test-idem-1") + second = _post(client, idem="test-idem-1") + assert first.status_code == 200 and second.status_code == 200 + assert second.headers.get("Idempotent-Replayed") == "true" + assert first.json() == second.json() + upstream.assert_not_called() + + def test_quota_exhausted_returns_402_without_calling_upstream(client, redis, upstream): redis.kv["gateway:v1:quota:42:literature.answer"] = 0 resp = _post(client, idem="q-1")