feat: canned, unbilled responses for omni_sk_test_ keys

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)

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.",
     )

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)

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,
+        },
+    }

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})

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")