From 28c38db968c7e09ce47c87ef472d908e2e90c61a Mon Sep 17 00:00:00 2001 From: man4ish Date: Sun, 4 Oct 2026 01:28:33 -0500 Subject: [PATCH] feat: expose BYOK provider-key storage through the public /v1 API PUT/GET/DELETE /v1/provider-keys(/{provider}) proxy into omnibioai-auth's own /orgs/{org_id}/provider-keys(/{provider}) (M14's storage service), the same forwarded-bearer-token pattern GET /v1/usage already uses for omnibioai-billing. Storage only, never a billable call; still rate-limited like every other /v1 route. The real authorization decision (manage_org) is enforced live by the destination endpoint itself, not this gateway. Co-Authored-By: Claude Sonnet 5 --- app/core/router.py | 20 ++++++++ app/routes/v1.py | 89 +++++++++++++++++++++++++++++++++++ tests/test_v1_literature.py | 93 +++++++++++++++++++++++++++++++++++++ 3 files changed, 202 insertions(+) diff --git a/app/core/router.py b/app/core/router.py index 6bef649..b829bc3 100644 --- a/app/core/router.py +++ b/app/core/router.py @@ -23,6 +23,17 @@ # for the "rag" entry above. See SERVICE_PERMISSION_MAP below for the # permission this requires. "billing": "http://billing-service:8005", + # Public API v1 ("M15": PUT/GET/DELETE /v1/provider-keys/{provider}): + # proxies into omnibioai-auth's own PUT/GET/DELETE + # /orgs/{id}/provider-keys(/{provider}) (M14's BYOK provider-key + # storage), the exact same forwarded-bearer-token pattern the + # "billing" entry above already uses for GET /v1/usage. The real + # authorization decision is manage_org, checked live by that + # endpoint's own require_org_permission_or_platform_admin dependency + # -- not this gateway, and not the JWT's permissions claim, which + # never embeds org-scoped permissions at all (they're resolved fresh + # per request against the caller's live membership). + "auth": "http://omnibioai-auth:8000", } @@ -68,6 +79,14 @@ def resolve_service(service: str) -> str | None: # Foundation integration made them real. This is usage.read's first # real consumer. "billing": "usage.read", + # Public API v1 (PUT/GET/DELETE /v1/provider-keys/{provider}): + # omnibioai-auth's Permission Registry already has + # "provider_keys.manage" registered as "reserved -- not yet enforced + # by any route" -- the exact same state usage.read was in before + # this gateway sent it as context. This is provider_keys.manage's + # first real consumer; the real enforcement is still manage_org, + # live, at the destination (see the SERVICE_MAP entry's own comment). + "auth": "provider_keys.manage", } @@ -86,6 +105,7 @@ def resolve_required_permission(service: str) -> str | None: # this so /v1/literature/* is authorized exactly like the rag service. V1_SERVICE_MAP = { "literature": "rag", + "provider-keys": "auth", } diff --git a/app/routes/v1.py b/app/routes/v1.py index 0b6a83e..e4f787b 100644 --- a/app/routes/v1.py +++ b/app/routes/v1.py @@ -447,3 +447,92 @@ async def literature_models(request: Request): {"model": "default", "source": "omnibioai_gpu", "billed_by": "query", "price": None}, ] return JSONResponse({"models": models}, status_code=200, headers={**headers, "X-Request-Id": request_id}) + + +# --------------------------------------------------------------------------- +# M15 (design audit gap #4, first slice of BYOK): proxies into +# omnibioai-auth's own PUT/GET/DELETE /orgs/{org_id}/provider-keys(/ +# {provider}) (M14's storage service) -- never implemented here, just +# forwarded, the same pattern GET /v1/usage above already uses for +# omnibioai-billing. Free (not a billable /v1/literature/* call); still +# rate-limited like every other /v1 route. +# --------------------------------------------------------------------------- + + +def _provider_key_error(status: int, response, request_id: str, headers: dict): + if status >= 500: + return _error(502, "upstream_error", + "The identity service failed to process this request.", request_id, headers) + if status == 403: + return _error(403, "forbidden", + "You do not have permission to manage this organization's provider keys.", + request_id, headers) + if status == 404: + return _error(404, "not_found", "No key is configured for that provider.", request_id, headers) + return _error(status, "invalid_request", "The identity service rejected the request.", + request_id, headers, detail=response) + + +@router.put("/provider-keys/{provider}") +async def set_provider_key(request: Request, provider: str): + """Store this organization's own Claude/OpenAI key, encrypted at + rest by omnibioai-auth. Storage only -- nothing yet routes a real + /v1/literature/answers call through it (design audit gap #4's + remaining bullets).""" + request_id = getattr(request.state, "trace_id", "") + subject, org_id, _ = _caller(request) + if not org_id: + return _error(403, "organization_required", + "This API is billed to an organization; your account has none.", request_id) + try: + body = await request.json() + except Exception: + return _error(400, "invalid_request", "Request body must be JSON.", request_id) + + headers, limited = await _rate_limited(request, subject, org_id, request_id) + if limited: + return limited + + status, response = await _forward_to("auth", request, "PUT", f"orgs/{org_id}/provider-keys/{provider}", body) + if not 200 <= status < 300: + return _provider_key_error(status, response, request_id, headers) + return JSONResponse(response, status_code=status, headers={**headers, "X-Request-Id": request_id}) + + +@router.get("/provider-keys") +async def get_provider_key(request: Request): + """Whether (and which provider's) key is configured -- never the + key itself, the same write-only contract omnibioai-auth's own + endpoint already enforces.""" + request_id = getattr(request.state, "trace_id", "") + subject, org_id, _ = _caller(request) + if not org_id: + return _error(403, "organization_required", + "This API is billed to an organization; your account has none.", request_id) + + headers, limited = await _rate_limited(request, subject, org_id, request_id) + if limited: + return limited + + status, response = await _forward_to("auth", request, "GET", f"orgs/{org_id}/provider-keys") + if not 200 <= status < 300: + return _provider_key_error(status, response, request_id, headers) + return JSONResponse(response, status_code=status, headers={**headers, "X-Request-Id": request_id}) + + +@router.delete("/provider-keys/{provider}") +async def delete_provider_key(request: Request, provider: str): + request_id = getattr(request.state, "trace_id", "") + subject, org_id, _ = _caller(request) + if not org_id: + return _error(403, "organization_required", + "This API is billed to an organization; your account has none.", request_id) + + headers, limited = await _rate_limited(request, subject, org_id, request_id) + if limited: + return limited + + status, response = await _forward_to("auth", request, "DELETE", f"orgs/{org_id}/provider-keys/{provider}") + if not 200 <= status < 300: + return _provider_key_error(status, response, request_id, headers) + return JSONResponse(response, status_code=status, headers={**headers, "X-Request-Id": request_id}) diff --git a/tests/test_v1_literature.py b/tests/test_v1_literature.py index bd69778..93df7d2 100644 --- a/tests/test_v1_literature.py +++ b/tests/test_v1_literature.py @@ -741,6 +741,99 @@ def test_models_is_free_static_and_never_calls_upstream(client, redis, upstream) assert _usage(redis) == [] +# --------------------------------------------------------------------------- +# M15 (design audit gap #4's BYOK storage, public surface): PUT/GET/DELETE +# /v1/provider-keys(/{provider}) proxy into omnibioai-auth's own +# /orgs/{org_id}/provider-keys(/{provider}) -- storage only, never a +# billable call. +# --------------------------------------------------------------------------- + + +def test_set_provider_key_requires_an_organization(client, redis, upstream): + with patch.object(_main_mod.iam, "validate_api_key", AsyncMock(return_value={**USER, "org_id": None})): + resp = client.request( + "PUT", "/v1/provider-keys/claude", json={"api_key": "sk-x"}, headers={"Authorization": f"Bearer {KEY}"}, + ) + assert resp.status_code == 403 + assert resp.json()["error"]["type"] == "organization_required" + upstream.assert_not_called() + + +def test_set_provider_key_forwards_to_auth_service(client, redis, upstream): + upstream.return_value = (200, {"provider": "claude", "has_key": True, "updated_at": None, "updated_by_email": None}) + resp = client.request( + "PUT", "/v1/provider-keys/claude", json={"api_key": "sk-secret"}, headers={"Authorization": f"Bearer {KEY}"}, + ) + assert resp.status_code == 200 + assert resp.json() == {"provider": "claude", "has_key": True, "updated_at": None, "updated_by_email": None} + + kwargs = upstream.call_args.kwargs + assert kwargs["url"] == "http://omnibioai-auth:8000/orgs/42/provider-keys/claude" + assert kwargs["method"] == "PUT" and kwargs["body"] == {"api_key": "sk-secret"} + assert "sk-secret" not in kwargs["headers"].get("Authorization", "") + assert _usage(redis) == [] # never billed + + +def test_get_provider_key_forwards_to_auth_service(client, redis, upstream): + upstream.return_value = (200, {"provider": None, "has_key": False, "updated_at": None, "updated_by_email": None}) + resp = client.get("/v1/provider-keys", headers={"Authorization": f"Bearer {KEY}"}) + assert resp.status_code == 200 + assert resp.json()["has_key"] is False + kwargs = upstream.call_args.kwargs + assert kwargs["url"] == "http://omnibioai-auth:8000/orgs/42/provider-keys" + assert kwargs["method"] == "GET" + + +def test_delete_provider_key_forwards_to_auth_service(client, redis, upstream): + upstream.return_value = (200, {"provider": None, "has_key": False, "updated_at": None, "updated_by_email": None}) + resp = client.request("DELETE", "/v1/provider-keys/claude", headers={"Authorization": f"Bearer {KEY}"}) + assert resp.status_code == 200 + kwargs = upstream.call_args.kwargs + assert kwargs["url"] == "http://omnibioai-auth:8000/orgs/42/provider-keys/claude" + assert kwargs["method"] == "DELETE" + + +def test_set_provider_key_maps_403_to_forbidden(client, redis, upstream): + upstream.return_value = (403, {"detail": "Forbidden"}) + resp = client.request( + "PUT", "/v1/provider-keys/claude", json={"api_key": "sk-x"}, headers={"Authorization": f"Bearer {KEY}"}, + ) + assert resp.status_code == 403 + assert resp.json()["error"]["type"] == "forbidden" + + +def test_delete_provider_key_maps_404_to_not_found(client, redis, upstream): + upstream.return_value = (404, {"detail": "No claude key is configured for this organization."}) + resp = client.request("DELETE", "/v1/provider-keys/claude", headers={"Authorization": f"Bearer {KEY}"}) + assert resp.status_code == 404 + assert resp.json()["error"]["type"] == "not_found" + + +def test_set_provider_key_maps_upstream_5xx_to_502(client, redis, upstream): + upstream.return_value = (500, {"detail": "CONFIG_ENCRYPTION_KEY is not set"}) + resp = client.request( + "PUT", "/v1/provider-keys/claude", json={"api_key": "sk-x"}, headers={"Authorization": f"Bearer {KEY}"}, + ) + assert resp.status_code == 502 + assert resp.json()["error"]["type"] == "upstream_error" + + +def test_set_provider_key_rejects_non_json_body(client, redis, upstream): + resp = client.request( + "PUT", "/v1/provider-keys/claude", content=b"not json", + headers={"Authorization": f"Bearer {KEY}", "Content-Type": "application/json"}, + ) + assert resp.status_code == 400 + upstream.assert_not_called() + + +def test_provider_key_routes_are_rate_limited(client, redis, upstream, monkeypatch): + monkeypatch.setattr(Config, "V1_RATE_LIMIT_PER_MINUTE", 1) + upstream.return_value = (200, {"provider": None, "has_key": False, "updated_at": None, "updated_by_email": None}) + assert client.get("/v1/provider-keys", headers={"Authorization": f"Bearer {KEY}"}).status_code == 200 + assert client.get("/v1/provider-keys", headers={"Authorization": f"Bearer {KEY}"}).status_code == 429 + + def test_unauthenticated_v1_is_rejected(client): assert client.post("/v1/literature/answers", json=BODY).status_code == 401