feat: dual JWT/per-user API key auth, RLS per user_id on mcp_servers, per-user LLM, RAG isolation - #62
Conversation
… per-user LLM, RAG isolation
- Dual auth: JWT (Authorization: Bearer, Logto OIDC JWKS via JwtAdapter) OR
per-user API keys (X-API-Key, shared api_keys table with composable-agents).
McpApiKeyMiddleware accepts both. New env vars LOGTO_URL, JWT_AUDIENCE.
- RLS per user_id on mcp_servers: migrations 002 (user_id column) + 003
(ENABLE/FORCE RLS + policy). McpRegistryStore filters by current_user_id
+ sets GUC app.user_id on asyncpg connection. BUG-001 fix: ON CONFLICT
DO UPDATE scoped by user_id → cross-user create same name → 409 (no silent
overwrite).
- Per-user LLM credentials: reads shared user_llm_settings table (owned by
composable-agents) via AsyncpgUserLlmReader (Fernet decrypt via existing
FernetSecretCipher). Factories get_embedding_for_user/get_chat_llm_for_user/
get_vector_store_for_user resolve per request from current_user_id; fallback
to LLMConfig env when contextvar None. LlmNotConfiguredError (422).
OPEN_ROUTER_API_KEY deprecated for chat+embeddings (still used for Kreuzberg
VLM — documented limitation).
- RAG isolation per user: chunks tagged with user_id in langchain_metadata at
index time; query filters {user_id: current_user_id} via langchain-postgres
metadata filter. RLS on dynamic classical_rag_* tables NOT applied (PGEngine
has its own pool, GUC not propagated — documented, application-level filter
is the mechanism).
- Shared tables api_keys + user_llm_settings read with SET LOCAL
row_security = off (privileged auth/credential-resolution op).
Tests: 557 unit tests pass (+105 new). QA: 195 tests pass on the local Docker
stack (auto-seed shared with composable-agents). 2 bugs found in QA and fixed
(cross-user mcp_servers overwrite, delete file requiring LLM creds).
5f119c2 to
09ef68c
Compare
Kaiohz
left a comment
There was a problem hiding this comment.
Code Review — PR #62: dual JWT/per-user API key auth, RLS per user_id, per-user LLM, RAG isolation
Score: 9/10 — solid, well-documented, well-tested PR. The hexagonal architecture pays off: every auth path is a port, every LLM resolution goes through a domain service, the RLS story is layered (GUC + app-level user_id filter = defence in depth). CI green, 105 new unit tests, 2 QA bugs caught and fixed in-flight. Below: a few non-blocking observations, mostly nits and one thing worth thinking about before merge.
What this PR does (and why the design is right)
- Auth surface (
security.py,auth_service.py):AuthServiceis the only component that knows JWT-vs-API-key precedence.ComposableAgentsSecurity.verify_credentials(FastAPI dep) andMcpApiKeyMiddleware.on_call_tool/on_list_toolsboth delegate to it. Precedence is JWT-first (matches the test contract — invalid JWT does not fall through to API key), which is the safer default: a forged or expired token should never silently succeed because an unrelated X-API-Key happens to validate. - JWT path (
jwt_adapter.py):httpx.AsyncClient+ in-memoryTTLCache(maxsize=1, ttl=300)+asyncio.Lockto prevent thundering herd on cache miss.kid→ key resolution falls back to the first key whenkidisNone(single-key JWKS — common in Logto). ReturnsNoneon any failure (never raises), so the dual-auth path can fall through cleanly. PII is never logged. - API key path (
api_key_reader.py): SHA-256 hash before lookup, readsapi_keyswithSET LOCAL row_security = off(justified — FORCE RLS table, user unknown at auth time). This is the correct shape for a privileged auth read. - RLS layering (
rls_context.py+mcp_registry_store.py+ migration 003): three layers of defence, in order: (1)app.user_idGUC on every asyncpg connection, (2)user_idclause on every SQL query, (3) PostgreSQL FORCE RLS policy that filters viacurrent_setting('app.user_id', true). The RLS contextvar +system_rls_contextasync ctx manager for migrations/cron is a clean way to keep the privileged read path intentional rather than implicit. - Per-user LLM (
user_llm_reader.py,credential_resolver.py,dependencies.py):get_embedding_for_user/get_chat_llm_for_user/get_vector_store_for_userreadcurrent_user_idand build per-request clients when a user is set; the env-based singletons stay as the legacy fallback.LlmNotConfiguredError(422) is a much better error than a 500 on a missing key. - RAG isolation (
classical_query_use_case.py,tag_documents_with_user_id): chunks taggeduser_idin metadata at index time, filter applied at query time via themetadata_filterarg. Clean, single point of wiring.
What's good (worth highlighting)
- Bug-001 fix documented and tested —
ON CONFLICT (name) DO UPDATEis scoped withWHERE mcp_servers.user_id = EXCLUDED.user_idto prevent cross-user ownership transfer. The 0-row-affected detection that raisesMcpServerAlreadyExistsErroris exactly the right way to surface the cross-user conflict. The RLS pitfall comment in thesave()docstring is the kind of thing that should be in a Confluence / runbook so future contributors don't reintroduce it. - Bug-002 fix —
delete_file/delete_folderuse cases swallowLlmNotConfiguredErrorand passvector_store=Noneso MinIO delete still works without LLM creds. Best-effort cascade is the right call here. - Contextvars propagate within the same asyncio task, which is exactly what makes the MCP middleware + RLS scoping work. The comment in
McpApiKeyMiddleware._check_api_keyabout the FastMCP handler running in the same task is a good defensive note. - READ-ONLY ports (
ApiKeyRepositoryPort,UserLlmSettingsRepositoryPort) make the ownership boundary explicit: composable-agents owns writes, mcp-raganything only reads. This is the kind of contract that prevents an accidentalDELETE FROM api_keysat 3am. - Error handling — the
_maskfunction forRegisteredMcpServer, the_BEARER_PREFIXextraction, the_ACCEPTED_ALGORITHMS = ["RS256", "ES256", "ES384"]whitelist (nononealgorithm — correct) — all small but important details done right. - Test coverage is well-shaped:
test_jwt_adapter.pymocks the JWKS endpoint,test_mcp_registry_user_isolation.pyexercises the RLS scoping,test_rag_user_isolation.pychecks chunk filter,test_rls_context.pycovers the contextvar reset,test_security_dual_auth.pyuses an in-memory FastMCP transport. Realistic, focused.
Observations / suggestions (non-blocking, sort by impact)
1. ⚠️ Race in JwtAdapter._get_cached_jwks lock re-check — verify the double-checked-locking pattern
The double-check pattern is correct if TTLCache.__getitem__ is atomic w.r.t. the asyncio event loop. Since the cache lives in a single process and is only mutated inside the lock, it is safe. However, the self._jwks_http_client is None check at the top of decode_token is read outside the lock — if a close() is called concurrently (e.g. during shutdown), get could raise on a closed client. Not a security issue, but worth a small if self._jwks_http_client is None or self._jwks_http_client.is_closed: return None check, or a graceful try/except around the httpx.get call to surface "client closed" as a None instead of a 500.
2. current_credential contextvar holds the raw JWT — confirm PII logging policy holds downstream
current_credential is set to the raw JWT in verify_credentials and to the raw API key in the middleware path. The PII-sensitive field. Currently nothing logs it (good), but it's a footgun: a future "log all RLS context for debugging" helper could leak. Consider one of:
- Wrap it in a
RedactedCredentialobject whose__repr__returns"***". - Or move the raw credential to a request-scoped cache keyed by a
correlation_idand only store the id in the contextvar.
Not required, but worth flagging.
3. verify_credentials returns AuthContext | str — the union return is a slight smell
async def verify_credentials(self, request: Request) -> AuthContext | str returns the legacy empty string "" when master-key auth is bypassed. FastAPI doesn't care, but downstream code that does _ctx: AuthContext = Depends(security.verify_credentials) and calls _ctx.user_id on a "" will crash with a confusing AttributeError on str. Two cleaner options:
- Two separate dependencies:
verify_credentials(raises on failure, returnsAuthContext) andverify_api_key_legacy(returnsstrfor the few legacy paths still using it). - Or always return
AuthContextand have the legacy fallback construct a syntheticAuthContext(user_id="", method="legacy_master_key", raw_credential="").
Currently the FastAPI routers only consume verify_credentials for the AuthContext path (no tests rely on the legacy return on the dual-auth branch), so this is more about future-proofing.
4. mcp_registry_store.save parses result.split()[-1] == "0" — robust but fragile
asyncpg's command tag format is well-defined ("INSERT n n" / "UPDATE n"), but the check result.split()[-1] == "0" treats "UPDATE 0" and "INSERT 0 0" the same. That's correct for the cross-user conflict path, but a same-user no-op (e.g. calling save twice with the same entry) also reports "UPDATE 0" and would incorrectly raise McpServerAlreadyExistsError. Edge case is unlikely in practice (the create use case is the only caller, and it only saves once), but a comment explaining the semantics would help future readers.
5. OPEN_ROUTER_API_KEY deprecation in the body is clean, but LLMConfig.api_key still falls back to it
The PR body says: "OPEN_ROUTER_API_KEY deprecated for chat+embed (still for VLM)". But LLMConfig.api_key (used by the legacy fallback path) still returns OPEN_ROUTER_API_KEY or OPENROUTER_API_KEY. For new deployments, this means setting OPEN_ROUTER_API_KEY is silently required for legacy mode — which is fine, but the app_config.api_key docstring in LLMConfig doesn't say "deprecated for chat/embed, kept for VLM only". Either update the docstring, or split into chat_api_key / vision_api_key properties to make the deprecation explicit in the type system.
6. auth_service.authenticate invalid-JWT → return None (no fallback to API key) — is this documented for callers?
if user is not None:
return AuthContext(...)
# Invalid JWT → no fallback to API key (matches the test contract).
return NoneThe comment is there. Just double-checking: a client sending both Authorization: Bearer <expired> and X-API-Key: <valid> will get 401. That's the right behavior (the expired JWT proves the user is trying to use a token they have, falling through to the API key would mask a clock skew or revocation issue). Document this in the API docs / OpenAPI description so integrators don't file "I sent both and got 401" as a bug.
7. Minor: LOGTO_URL=/oidc/jwks concatenation — no trailing slash normalisation
jwks_url=f"{app_config.LOGTO_URL}/oidc/jwks" if app_config.LOGTO_URL else ""If LOGTO_URL is set to "https://logto.example.com/" (trailing slash), the URL becomes "https://logto.example.com//oidc/jwks". Most HTTP clients handle double slashes, but to be safe either rstrip("/") on LOGTO_URL or document the expected format in the env-var description.
8. Test ergonomics — test_security_dual_auth.py mocks AuthService entirely
The 11 tests in test_security_dual_auth.py mock the AuthService so they never exercise the actual precedence or error paths inside AuthService. The 16 tests in test_auth_service.py cover those. The split is fine, but a couple of integration tests with a real JwtAdapter + real (mocked) ApiKeyReader would be high-value: a regression in precedence (e.g. someone changing the order) would be caught by an integration test, not by the unit-level mocks.
9. UserLlmCredentialResolver has no caching — every RAG request decrypts the API key
Per-request Fernet decrypt is fine (cheap), but the same (base_url, api_key) is reconstructed on every RAG call. A small in-memory TTLCache (e.g. 60s, per-user) would reduce DB pressure and Fernet CPU under high QPS. Not urgent — the PR body says "the per-user LLM/embeddings factories build fresh per-user instances", implying the cost was considered — but worth a follow-up ticket.
10. New env vars documented in .env.example ✅
LOGTO_URL and JWT_AUDIENCE are present in .env.example with the right comments. The breaking-change note for API_KEY and OPEN_ROUTER_API_KEY is in the PR body but not yet in the README — make sure the README section on environment variables gets the same treatment before merge (a quick git log -p README.md | grep -A 5 -B 2 OPEN_ROUTER would confirm).
Nits (won't fail the PR)
src/dependencies.pyis now 379 lines. Consider splitting into_deps_auth.py,_deps_llm.py,_deps_mcp.py(theget_*functions already group naturally). Not blocking, but a future maintainability win.classical_bm25_adapter: BM25EnginePort | None = Nonethenif bm25_config.BM25_ENABLED: try: ...; except: classical_bm25_adapter = None— thetry/excepthere is the legacy module-import pattern. A smallBM25_ENABLED and await self._init()lazy factory would be cleaner and avoid the silent fallback at import time.verify_credentialsis wired onapp.include_router(file_router, ...)etc. via_auth_dep = [Depends(security.verify_credentials)]. If a future router is added without_auth_dep, it silently falls back to no auth. Adefault=Depends(security.verify_credentials)on the router include, or a globalapp.dependency_overridespattern, would make the contract unmissable.
Verdict
Approve with comments. The architecture is right, the RLS layering is correct, the auth precedence is the safer choice, the tests cover the security boundaries, and the two bugs found+fixed in QA are exactly the kind of issues this kind of suite should catch. The observations above are non-blocking; the most impactful to address before merge are #2 (PII surface in current_credential), #6 (document the no-fallback contract), and a final pass on the README to mirror the .env.example deprecations.
Nice work — this is a clean, well-reasoned PR. 🛡️
Suite of composable-agents PR #39: mcp-raganything adopts dual JWT/per-user API key auth (shared api_keys table), RLS per user_id on mcp_servers, per-user LLM credentials (shared user_llm_settings table), RAG chunk isolation via user_id metadata filter.
Key changes:
Quality: 557 unit tests (+105), 195 QA tests (2 bugs found+fixed in QA), code review 0 critical, sonar 0 new blocking, trivy 0 vuln, README updated.
New env vars: LOGTO_URL, JWT_AUDIENCE. Breaking: API_KEY master deprecated, OPEN_ROUTER_API_KEY deprecated for chat+embeddings.