Skip to content

feat: dual JWT/per-user API key auth, RLS per user_id on mcp_servers, per-user LLM, RAG isolation - #62

Merged
Kaiohz merged 1 commit into
mainfrom
feat/dual-auth-rls-llm-peruser
Jul 27, 2026
Merged

feat: dual JWT/per-user API key auth, RLS per user_id on mcp_servers, per-user LLM, RAG isolation#62
Kaiohz merged 1 commit into
mainfrom
feat/dual-auth-rls-llm-peruser

Conversation

@Kaiohz

@Kaiohz Kaiohz commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

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:

  • Dual auth: JwtAdapter (Logto JWKS) + AsyncpgApiKeyReader (shared api_keys, row_security=off bypass) + McpApiKeyMiddleware extended (Bearer JWT + X-API-Key).
  • RLS on mcp_servers: migrations 002/003, McpRegistryStore filters by current_user_id. BUG-001 fix: ON CONFLICT scoped by user_id → cross-user create same name → 409 (no silent overwrite).
  • Per-user LLM: AsyncpgUserLlmReader (shared user_llm_settings, Fernet), factories get_embedding/chat/vector_store_for_user. LlmNotConfiguredError 422. OPEN_ROUTER_API_KEY deprecated for chat+embed (still for VLM). BUG-002 fix: delete file/folder no longer requires LLM creds.
  • RAG isolation: chunks tagged user_id in langchain_metadata; query filters {user_id}. RLS on dynamic classical_rag_* skipped (PGEngine own pool).

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.

@Kaiohz
Kaiohz marked this pull request as ready for review July 27, 2026 09:07
… 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).
@Kaiohz
Kaiohz force-pushed the feat/dual-auth-rls-llm-peruser branch from 5f119c2 to 09ef68c Compare July 27, 2026 09:11

@Kaiohz Kaiohz left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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): AuthService is the only component that knows JWT-vs-API-key precedence. ComposableAgentsSecurity.verify_credentials (FastAPI dep) and McpApiKeyMiddleware.on_call_tool / on_list_tools both 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-memory TTLCache(maxsize=1, ttl=300) + asyncio.Lock to prevent thundering herd on cache miss. kid → key resolution falls back to the first key when kid is None (single-key JWKS — common in Logto). Returns None on 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, reads api_keys with SET 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_id GUC on every asyncpg connection, (2) user_id clause on every SQL query, (3) PostgreSQL FORCE RLS policy that filters via current_setting('app.user_id', true). The RLS contextvar + system_rls_context async 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_user read current_user_id and 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 tagged user_id in metadata at index time, filter applied at query time via the metadata_filter arg. Clean, single point of wiring.

What's good (worth highlighting)

  1. Bug-001 fix documented and testedON CONFLICT (name) DO UPDATE is scoped with WHERE mcp_servers.user_id = EXCLUDED.user_id to prevent cross-user ownership transfer. The 0-row-affected detection that raises McpServerAlreadyExistsError is exactly the right way to surface the cross-user conflict. The RLS pitfall comment in the save() docstring is the kind of thing that should be in a Confluence / runbook so future contributors don't reintroduce it.
  2. Bug-002 fixdelete_file / delete_folder use cases swallow LlmNotConfiguredError and pass vector_store=None so MinIO delete still works without LLM creds. Best-effort cascade is the right call here.
  3. Contextvars propagate within the same asyncio task, which is exactly what makes the MCP middleware + RLS scoping work. The comment in McpApiKeyMiddleware._check_api_key about the FastMCP handler running in the same task is a good defensive note.
  4. 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 accidental DELETE FROM api_keys at 3am.
  5. Error handling — the _mask function for RegisteredMcpServer, the _BEARER_PREFIX extraction, the _ACCEPTED_ALGORITHMS = ["RS256", "ES256", "ES384"] whitelist (no none algorithm — correct) — all small but important details done right.
  6. Test coverage is well-shaped: test_jwt_adapter.py mocks the JWKS endpoint, test_mcp_registry_user_isolation.py exercises the RLS scoping, test_rag_user_isolation.py checks chunk filter, test_rls_context.py covers the contextvar reset, test_security_dual_auth.py uses 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 RedactedCredential object whose __repr__ returns "***".
  • Or move the raw credential to a request-scoped cache keyed by a correlation_id and 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, returns AuthContext) and verify_api_key_legacy (returns str for the few legacy paths still using it).
  • Or always return AuthContext and have the legacy fallback construct a synthetic AuthContext(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 None

The 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.py is now 379 lines. Consider splitting into _deps_auth.py, _deps_llm.py, _deps_mcp.py (the get_* functions already group naturally). Not blocking, but a future maintainability win.
  • classical_bm25_adapter: BM25EnginePort | None = None then if bm25_config.BM25_ENABLED: try: ...; except: classical_bm25_adapter = None — the try/except here is the legacy module-import pattern. A small BM25_ENABLED and await self._init() lazy factory would be cleaner and avoid the silent fallback at import time.
  • verify_credentials is wired on app.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. A default=Depends(security.verify_credentials) on the router include, or a global app.dependency_overrides pattern, 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. 🛡️

@Kaiohz
Kaiohz merged commit b948245 into main Jul 27, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant