From b0a1d99747a0bafc75f883b908d61563a0067eab Mon Sep 17 00:00:00 2001 From: man4ish Date: Sun, 4 Oct 2026 00:32:42 -0500 Subject: [PATCH] feat: BYOK provider-key storage for organizations An organization can now store its own encrypted Claude/OpenAI key (PUT/GET/DELETE /orgs/{org_id}/provider-keys), gated on manage_org -- an org-wide infrastructure setting, not a self-service literature-API action. Built on the organization_config table, which has existed since the multi-tenant schema migration with no service or route ever using it. One provider/key slot per organization, mirroring GlobalConfig's own platform-wide shape; setting a new provider replaces the previous one. Credentials are write-only (has_key flags only, never echoed back) and refused with 500 rather than stored in plaintext when CONFIG_ENCRYPTION_KEY isn't configured, reusing the same crypto module config_service.py already established. Deliberately scoped to storage only: actually routing a /v1/literature/answers call through the stored provider (token accounting, llm.tokens.* events), and exposing this through the public gateway's own /v1/provider-keys/{provider} path, are follow-up work -- this is reachable today the same way Studio's own LLM settings page would be, through the org-admin endpoint directly. Co-Authored-By: Claude Sonnet 5 --- app/api/routes_organization_config.py | 82 +++++++ app/main.py | 2 + app/schemas/organization_config.py | 15 ++ app/services/audit_service.py | 5 + app/services/organization_config_service.py | 113 ++++++++++ tests/test_organization_provider_keys.py | 228 ++++++++++++++++++++ tests/test_route_authorization_coverage.py | 9 +- 7 files changed, 453 insertions(+), 1 deletion(-) create mode 100644 app/api/routes_organization_config.py create mode 100644 app/schemas/organization_config.py create mode 100644 app/services/organization_config_service.py create mode 100644 tests/test_organization_provider_keys.py diff --git a/app/api/routes_organization_config.py b/app/api/routes_organization_config.py new file mode 100644 index 0000000..d35c78b --- /dev/null +++ b/app/api/routes_organization_config.py @@ -0,0 +1,82 @@ +from fastapi import APIRouter, Depends, HTTPException +from sqlalchemy.orm import Session + +from app.db.models import OrganizationMembership, User +from app.db.session import get_db +from app.rbac import require_org_permission_or_platform_admin +from app.schemas.organization_config import ProviderKeyIn, ProviderKeyOut +from app.services import organization_config_service + +# BYOK provider-key storage (design audit gap #4). Gated on manage_org -- +# an org-wide, billing-adjacent infrastructure setting, the same +# administrative character as manage_billing/manage_teams, not a +# literature-API self-service action, so it is deliberately NOT one of +# the public literature:read/usage:read scopes a self-service key can +# be issued with (see routes_apikeys.py's PUBLIC_SCOPE_TO_PERMISSION). +router = APIRouter(prefix="/orgs/{org_id}/provider-keys", tags=["provider-keys"]) + +MANAGE_ORG = "manage_org" + + +def _to_out(db: Session, config) -> ProviderKeyOut: + updated_by_email = None + if config and config.updated_by_user_id: + user = db.query(User).filter(User.id == config.updated_by_user_id).first() + updated_by_email = user.email if user else None + + return ProviderKeyOut( + provider=config.llm_provider if config else None, + has_key=bool(config and config.llm_api_key_encrypted), + updated_at=config.updated_at.isoformat() if config and config.updated_at else None, + updated_by_email=updated_by_email, + ) + + +@router.get("", response_model=ProviderKeyOut) +def get_provider_key( + org_id: int, + db: Session = Depends(get_db), + membership: OrganizationMembership = Depends(require_org_permission_or_platform_admin(MANAGE_ORG)), +): + return _to_out(db, organization_config_service.get_organization_config(db, org_id)) + + +@router.put("/{provider}", response_model=ProviderKeyOut) +def put_provider_key( + org_id: int, + provider: str, + body: ProviderKeyIn, + db: Session = Depends(get_db), + membership: OrganizationMembership = Depends(require_org_permission_or_platform_admin(MANAGE_ORG)), +): + try: + config = organization_config_service.set_provider_key( + db, org_id, provider, body.api_key, updated_by_user_id=membership.user_id, + ) + except ValueError as e: + raise HTTPException(400, str(e)) + except RuntimeError as e: + # crypto.encrypt() raises when CONFIG_ENCRYPTION_KEY isn't set -- + # a loud, deliberate 500, never a silently-dropped or plaintext- + # stored credential (same contract routes_config.py's own + # update_config endpoint follows for the platform-wide key). + raise HTTPException(500, str(e)) + return _to_out(db, config) + + +@router.delete("/{provider}", response_model=ProviderKeyOut) +def delete_provider_key( + org_id: int, + provider: str, + db: Session = Depends(get_db), + membership: OrganizationMembership = Depends(require_org_permission_or_platform_admin(MANAGE_ORG)), +): + try: + config = organization_config_service.clear_provider_key( + db, org_id, provider, updated_by_user_id=membership.user_id, + ) + except ValueError as e: + raise HTTPException(400, str(e)) + if config is None: + raise HTTPException(404, f"No {provider} key is configured for this organization.") + return _to_out(db, config) diff --git a/app/main.py b/app/main.py index 87b8d34..a578c98 100644 --- a/app/main.py +++ b/app/main.py @@ -13,6 +13,7 @@ from app.api.routes_apikeys import router as apikeys_router from app.api.routes_apikeys import exchange_router as apikeys_exchange_router from app.api.routes_apikeys import me_router as apikeys_me_router +from app.api.routes_organization_config import router as provider_keys_router from app.api.routes_oauth_clients import router as oauth_clients_router from app.api.routes_oauth_token import router as oauth_token_router from app.api.routes_org_sso import router as org_sso_router @@ -104,6 +105,7 @@ app.include_router(apikeys_router) app.include_router(apikeys_exchange_router) app.include_router(apikeys_me_router) +app.include_router(provider_keys_router) app.include_router(oauth_clients_router) app.include_router(oauth_token_router) app.include_router(org_sso_router) diff --git a/app/schemas/organization_config.py b/app/schemas/organization_config.py new file mode 100644 index 0000000..832e0bc --- /dev/null +++ b/app/schemas/organization_config.py @@ -0,0 +1,15 @@ +from pydantic import BaseModel + + +class ProviderKeyIn(BaseModel): + api_key: str + + +class ProviderKeyOut(BaseModel): + """No credential field, ever, regardless of caller role -- has_key + only, the same write-only convention app/schemas/config.py's + GlobalConfigOut already established for the platform-wide LLM key.""" + provider: str | None + has_key: bool + updated_at: str | None + updated_by_email: str | None diff --git a/app/services/audit_service.py b/app/services/audit_service.py index 8516921..316a32d 100644 --- a/app/services/audit_service.py +++ b/app/services/audit_service.py @@ -43,6 +43,11 @@ class AuditEventType: # status, and the key material itself are unaffected and keep their own # existing event types. API_KEY_RENAMED = "api_key_renamed" + # M14 (BYOK provider-key storage, design audit gap #4): an + # organization's own Claude/OpenAI key. Never the key itself in + # metadata -- same rule API_KEY_CREATED's own entry follows. + PROVIDER_KEY_SET = "provider_key_set" + PROVIDER_KEY_CLEARED = "provider_key_cleared" OAUTH_CLIENT_CREATED = "oauth_client_created" OAUTH_CLIENT_REVOKED = "oauth_client_revoked" SSO_CONFIGURATION_CREATED = "sso_configuration_created" diff --git a/app/services/organization_config_service.py b/app/services/organization_config_service.py new file mode 100644 index 0000000..7d41d8a --- /dev/null +++ b/app/services/organization_config_service.py @@ -0,0 +1,113 @@ +"""M14 (design audit gap #4's "organisation-scoped BYOK provider-key +service"): an organization can store its own Claude/OpenAI key, used +instead of the platform's own when the public API answers that +organization's questions. The underlying organization_config table has +existed since the multi-tenant schema migration with nobody reading or +writing it -- this module is its first real consumer. + +Scoped deliberately narrowly, matching how far "storage" goes on its +own: this only stores and retrieves the encrypted key. Actually routing +a /v1/literature/answers call through the stored provider (Claude/ +OpenAI client calls, token accounting, llm.tokens.* usage events) is +the rest of gap #4 and is not built here -- storing a key an org can +never yet be routed through is a real, honestly-scoped gap, not a +hidden one (the same "mechanism before the business/product decision +that uses it" shape this project's M3/M6/M7 all already established, +here applied to an engineering dependency rather than a pricing one). + +Also not yet exposed through the public gateway's own /v1/provider-keys/ +{provider} path the design doc names -- only through this org-admin +endpoint (PUT/GET/DELETE /orgs/{org_id}/provider-keys), the same one +Studio's own LLM settings page would call. Routing this through the +gateway for direct external-developer use is follow-up work. +""" +from datetime import datetime + +from sqlalchemy.orm import Session + +from app.core import crypto +from app.db.models import OrganizationConfig +from app.services import audit_service +from app.services.audit_service import AuditEventType + +# One provider/key pair per organization at a time -- the same shape +# omnibioai-auth's own GlobalConfig already uses for the platform-wide +# LLM credential (config_service.py), not a new multi-provider-per-org +# design. Setting a key for a different provider replaces whichever one +# was previously configured, exactly like PUT /auth/config with a new +# llm_provider does there. +SUPPORTED_PROVIDERS = {"claude", "openai"} + + +def get_organization_config(db: Session, organization_id: int) -> OrganizationConfig | None: + return db.query(OrganizationConfig).filter_by(organization_id=organization_id).first() + + +def _get_or_create(db: Session, organization_id: int) -> OrganizationConfig: + config = get_organization_config(db, organization_id) + if config is None: + config = OrganizationConfig(organization_id=organization_id) + db.add(config) + db.flush() + return config + + +def set_provider_key( + db: Session, organization_id: int, provider: str, api_key: str, updated_by_user_id: int, +) -> OrganizationConfig: + """Raises ValueError for an unsupported provider name, and (via + crypto.encrypt) RuntimeError if CONFIG_ENCRYPTION_KEY isn't + configured -- never silently stores a credential in plaintext.""" + if provider not in SUPPORTED_PROVIDERS: + raise ValueError(f"Unsupported provider {provider!r}; supported: {sorted(SUPPORTED_PROVIDERS)}") + + config = _get_or_create(db, organization_id) + config.llm_provider = provider + config.llm_api_key_encrypted = crypto.encrypt(api_key) + config.updated_at = datetime.utcnow() + config.updated_by_user_id = updated_by_user_id + db.flush() + # Never the key itself or its ciphertext in audit metadata -- the + # same "name/scope only" rule apikey_service.create_api_key's own + # audit entry follows. + audit_service.log_event( + db, AuditEventType.PROVIDER_KEY_SET, actor_user_id=updated_by_user_id, + organization_id=organization_id, resource_type="organization_config", resource_id=config.id, + after_state={"provider": provider}, metadata={"provider": provider}, + commit=False, + ) + db.commit() + db.refresh(config) + return config + + +def clear_provider_key( + db: Session, organization_id: int, provider: str, updated_by_user_id: int, +) -> OrganizationConfig | None: + """Returns None (the route turns this into 404) if the organization + has no config row at all, or its currently-configured provider + doesn't match `provider` -- there is only one slot, so "clear + openai's key" when the org actually has a claude key configured + isn't a no-op clear of the wrong thing, it's simply nothing to do. + """ + if provider not in SUPPORTED_PROVIDERS: + raise ValueError(f"Unsupported provider {provider!r}; supported: {sorted(SUPPORTED_PROVIDERS)}") + + config = get_organization_config(db, organization_id) + if config is None or config.llm_provider != provider: + return None + + config.llm_provider = None + config.llm_api_key_encrypted = None + config.updated_at = datetime.utcnow() + config.updated_by_user_id = updated_by_user_id + db.flush() + audit_service.log_event( + db, AuditEventType.PROVIDER_KEY_CLEARED, actor_user_id=updated_by_user_id, + organization_id=organization_id, resource_type="organization_config", resource_id=config.id, + before_state={"provider": provider}, metadata={"provider": provider}, + commit=False, + ) + db.commit() + db.refresh(config) + return config diff --git a/tests/test_organization_provider_keys.py b/tests/test_organization_provider_keys.py new file mode 100644 index 0000000..0f284a8 --- /dev/null +++ b/tests/test_organization_provider_keys.py @@ -0,0 +1,228 @@ +"""/orgs/{org_id}/provider-keys: BYOK storage for an organization's own +Claude/OpenAI key (design audit gap #4). Gated on manage_org (org +administration, not a self-service literature-API action); credentials +are never returned, refused (500) rather than stored in plaintext +without an encryption key configured, and genuinely encrypted at rest +once one is -- same contract app/schemas/config.py's GlobalConfig +already established for the platform-wide LLM key, applied here per +organization. + +Developer: Manish Kumar +""" +import uuid + +import pytest +from cryptography.fernet import Fernet + +from app.services import organization_config_service + + +def _auth_header(token): + return {"Authorization": f"Bearer {token}"} + + +def _register_and_login(client, email=None): + email = email or f"provkey-{uuid.uuid4().hex[:8]}@omnibioai.test" + password = "TestPassword123!" + client.post("/auth/register", json={"email": email, "password": password}) + login = client.post("/auth/login", json={"email": email, "password": password}) + return {"email": email, "access_token": login.json()["access_token"]} + + +@pytest.fixture +def org(client): + """An organization owned by a freshly registered user (org_admin, + holding manage_org) -- the same shape test_apikeys.py's own `org` + fixture uses.""" + owner = _register_and_login(client) + headers = _auth_header(owner["access_token"]) + created = client.post( + "/orgs", + json={"name": "Provider Key Test Org", "slug": f"provkey-org-{uuid.uuid4().hex[:8]}"}, + headers=headers, + ).json() + return {"id": created["id"], "owner": owner, "owner_headers": headers} + + +@pytest.fixture +def configured_crypto(monkeypatch): + """Patches app.core.crypto's module-level singleton directly, since it's + computed once at import time from CONFIG_ENCRYPTION_KEY -- the same + pattern tests/test_config.py's own fixture of the same name uses.""" + import app.core.crypto as crypto + + key = Fernet.generate_key() + monkeypatch.setattr(crypto, "_fernet", Fernet(key)) + return crypto + + +# ── Permission gating ──────────────────────────────────────────────────────── + + +def test_get_requires_auth(client, org): + resp = client.get(f"/orgs/{org['id']}/provider-keys") + assert resp.status_code in (401, 403) + + +def test_get_returns_no_key_configured_by_default(client, org): + resp = client.get(f"/orgs/{org['id']}/provider-keys", headers=org["owner_headers"]) + assert resp.status_code == 200 + body = resp.json() + assert body["provider"] is None and body["has_key"] is False + + +def test_non_admin_member_cannot_set_a_key(client, org): + """A member without manage_org (no role assignment at all here) gets 403.""" + outsider = _register_and_login(client) + outsider_headers = _auth_header(outsider["access_token"]) + resp = client.put( + f"/orgs/{org['id']}/provider-keys/claude", + json={"api_key": "sk-nope"}, + headers=outsider_headers, + ) + # Not a member at all -- the dependency 404s rather than 403, same as + # every other /orgs/{org_id}/* route's non-member behavior. + assert resp.status_code == 404 + + +# ── Validation ─────────────────────────────────────────────────────────────── + + +def test_unsupported_provider_name_is_rejected(client, org, configured_crypto): + resp = client.put( + f"/orgs/{org['id']}/provider-keys/not-a-real-provider", + json={"api_key": "sk-whatever"}, + headers=org["owner_headers"], + ) + assert resp.status_code == 400 + + +def test_delete_unsupported_provider_name_is_rejected(client, org): + resp = client.delete(f"/orgs/{org['id']}/provider-keys/not-a-real-provider", headers=org["owner_headers"]) + assert resp.status_code == 400 + + +# ── No silent plaintext fallback ───────────────────────────────────────────── + + +def test_set_key_fails_loudly_when_encryption_key_unset(client, org): + """conftest.py never sets CONFIG_ENCRYPTION_KEY, so app.core.crypto's + module-level _fernet is None for the whole test session -- this + exercises the real "not configured" state, not a simulated one.""" + resp = client.put( + f"/orgs/{org['id']}/provider-keys/claude", + json={"api_key": "sk-should-never-be-stored-plaintext"}, + headers=org["owner_headers"], + ) + assert resp.status_code == 500 + + +# ── Real encryption round-trip (key patched in for just these tests) ─────── + + +def test_set_key_returns_has_key_true_and_never_the_key_itself(client, org, configured_crypto): + resp = client.put( + f"/orgs/{org['id']}/provider-keys/claude", + json={"api_key": "sk-real-secret-value-12345"}, + headers=org["owner_headers"], + ) + assert resp.status_code == 200 + body = resp.json() + assert body["provider"] == "claude" + assert body["has_key"] is True + assert "api_key" not in body + assert "sk-real-secret-value-12345" not in str(body) + + +def test_set_key_is_genuinely_encrypted_at_rest(client, org, configured_crypto): + client.put( + f"/orgs/{org['id']}/provider-keys/claude", + json={"api_key": "sk-real-secret-value-12345"}, + headers=org["owner_headers"], + ) + from app.db.session import SessionLocal + + db = SessionLocal() + try: + config = organization_config_service.get_organization_config(db, org["id"]) + assert config.llm_api_key_encrypted != "sk-real-secret-value-12345" + assert configured_crypto.decrypt(config.llm_api_key_encrypted) == "sk-real-secret-value-12345" + finally: + db.close() + + +def test_get_reflects_the_configured_provider(client, org, configured_crypto): + client.put( + f"/orgs/{org['id']}/provider-keys/openai", + json={"api_key": "sk-openai-secret"}, + headers=org["owner_headers"], + ) + resp = client.get(f"/orgs/{org['id']}/provider-keys", headers=org["owner_headers"]) + assert resp.json() == { + "provider": "openai", + "has_key": True, + "updated_at": resp.json()["updated_at"], + "updated_by_email": org["owner"]["email"], + } + + +def test_setting_a_new_provider_replaces_the_previous_one(client, org, configured_crypto): + """One slot per organization -- setting openai after claude was + configured replaces it, it does not add a second provider.""" + client.put(f"/orgs/{org['id']}/provider-keys/claude", json={"api_key": "sk-claude"}, headers=org["owner_headers"]) + client.put(f"/orgs/{org['id']}/provider-keys/openai", json={"api_key": "sk-openai"}, headers=org["owner_headers"]) + + resp = client.get(f"/orgs/{org['id']}/provider-keys", headers=org["owner_headers"]) + assert resp.json()["provider"] == "openai" + + +# ── Clearing ───────────────────────────────────────────────────────────────── + + +def test_delete_clears_a_configured_key(client, org, configured_crypto): + client.put(f"/orgs/{org['id']}/provider-keys/claude", json={"api_key": "sk-claude"}, headers=org["owner_headers"]) + + resp = client.delete(f"/orgs/{org['id']}/provider-keys/claude", headers=org["owner_headers"]) + assert resp.status_code == 200 + assert resp.json()["provider"] is None + assert resp.json()["has_key"] is False + + listed = client.get(f"/orgs/{org['id']}/provider-keys", headers=org["owner_headers"]) + assert listed.json()["has_key"] is False + + +def test_delete_for_a_provider_that_is_not_the_configured_one_404s(client, org, configured_crypto): + """Only one slot exists -- "delete openai's key" when claude is + actually configured is 404, not a silent no-op success.""" + client.put(f"/orgs/{org['id']}/provider-keys/claude", json={"api_key": "sk-claude"}, headers=org["owner_headers"]) + + resp = client.delete(f"/orgs/{org['id']}/provider-keys/openai", headers=org["owner_headers"]) + assert resp.status_code == 404 + + # The actually-configured claude key must be untouched by the failed attempt. + listed = client.get(f"/orgs/{org['id']}/provider-keys", headers=org["owner_headers"]) + assert listed.json() == {"provider": "claude", "has_key": True, + "updated_at": listed.json()["updated_at"], "updated_by_email": org["owner"]["email"]} + + +def test_delete_with_nothing_configured_404s(client, org): + resp = client.delete(f"/orgs/{org['id']}/provider-keys/claude", headers=org["owner_headers"]) + assert resp.status_code == 404 + + +# ── Cross-org isolation ────────────────────────────────────────────────────── + + +def test_non_member_cannot_read_or_set_another_orgs_key(client, org, configured_crypto): + outsider = _register_and_login(client) + outsider_headers = _auth_header(outsider["access_token"]) + + assert client.get(f"/orgs/{org['id']}/provider-keys", headers=outsider_headers).status_code == 404 + resp = client.put( + f"/orgs/{org['id']}/provider-keys/claude", json={"api_key": "sk-stolen"}, headers=outsider_headers, + ) + assert resp.status_code == 404 + + # Confirm the attempted write had no effect. + listed = client.get(f"/orgs/{org['id']}/provider-keys", headers=org["owner_headers"]) + assert listed.json()["has_key"] is False diff --git a/tests/test_route_authorization_coverage.py b/tests/test_route_authorization_coverage.py index 92d42d0..2cf9bc3 100644 --- a/tests/test_route_authorization_coverage.py +++ b/tests/test_route_authorization_coverage.py @@ -171,8 +171,15 @@ def test_org_scoped_route_inventory_matches_expected_count(): require_org_permission_or_platform_admin(MANAGE_API_KEYS) -- the same dependency the existing POST/GET/DELETE routes at that same path already use. + + 44 -> 47 as of M14 (BYOK provider-key storage, design audit gap #4): + GET/PUT/DELETE /orgs/{org_id}/provider-keys(/{provider}) were added, + using require_org_permission_or_platform_admin(MANAGE_ORG) -- an + org-wide infrastructure setting, the same administrative character + manage_org already gates elsewhere (e.g. GET/PUT /orgs/{org_id} + itself), not a new permission. """ - assert len(list(_org_scoped_routes())) == 44 + assert len(list(_org_scoped_routes())) == 47 def test_every_org_scoped_route_uses_the_platform_admin_aware_dependency():