litellm/tests/test_litellm/proxy
Ishaan Jaff 1bb713bc7b
feat(mcp): BYOK MCP servers with OAuth 2.1 PKCE authorization flow (#22850)
* feat(mcp): BYOK (Bring Your Own Key) for OpenAPI MCP servers with OAuth 2.1 flow

Adds per-user credential storage for BYOK MCP servers so external clients
can authenticate via standard OAuth 2.1 PKCE without needing a full identity
provider.

Backend:
- New DB table LiteLLM_MCPUserCredentials (user_id, server_id, credential_b64)
- is_byok, byok_description, byok_api_key_help_url fields on MCPServerTable
- OAuth 2.1 authorization server endpoints (/.well-known/oauth-authorization-server,
  /.well-known/oauth-protected-resource, /v1/mcp/oauth/authorize, /v1/mcp/oauth/token)
- 401 challenge with WWW-Authenticate header when BYOK server has no credential
- CRUD endpoints: POST/DELETE /v1/mcp/server/{id}/user-credential
- has_user_credential annotated on GET /v1/mcp/server response

UI:
- ByokCredentialModal: 2-step Connect flow (access description + API key entry)
- BYOK toggle + description fields on admin MCP server create form
- Connect/Connected state in MCP server table
- BYOK Demo page (/tools/byok-demo) showing full OAuth 2.1 PKCE flow

* feat(mcp/byok): redesign OAuth authorize page to match 2-step Connect mockup

- Step 1: L→S logos, requested access checklist, How it works box, Continue button
- Step 2: API key input, Save toggle, Duration pills (1h/24h/7d/30d/until_revoked), security note
- Matches screenshots: white modal on dark bg, progress dots, dark CTA buttons
- Authorize handler now fetches byok_description and byok_api_key_help_url from server registry
- CLAUDE.md: replace SQL snippet with proper DB migration troubleshooting guidance

* fix: address greptile review feedback (greploop iteration 1)

- XSS: escape all user-supplied values in _build_authorize_html() with html.escape()
- Open redirect: validate redirect_uri scheme and URL-encode code/state in redirect
- N+1 query: batch BYOK credential lookup into single find_many() call
- Critical path DB: add 60s TTL in-memory cache to _check_byok_credential()
- Encrypt BYOK credentials at rest using encrypt_value_helper/decrypt_value_helper

* fix(byok): update OAuth popup with LiteLLM logo, MCP title suffix, remove emojis

* fix(byok-demo): fix token endpoint URL (/v1/mcp/oauth/token not /v1/mcp/token)

* feat(byok): inject stored BYOK credential as mcp_auth_header on tool execution

* feat(byok): use contextvars to inject per-user credential into OpenAPI tool closures; remove byok-demo from LiteLLM UI

OpenAPI tools have auth headers baked into their closures at registration time. BYOK servers have
no static auth token, so per-user credentials were never reaching the HTTP calls.

Fix: add _request_auth_header ContextVar in openapi_to_mcp_generator.py. create_tool_function now
reads this var at call time and overrides the Authorization header if set. execute_mcp_tool resolves
the MCP server and performs BYOK checks before the local-tool dispatch branch, then sets the
ContextVar around _handle_local_mcp_tool so the credential flows into the HTTP request.

Also remove the /tools/byok-demo page from the LiteLLM UI dashboard — the demo lives at
~/Downloads/litellm-byok-demo/index.html (served separately on port 8080).

* fix: address greptile review feedback (greploop iteration 2)

- Cache invalidation: add _invalidate_byok_cred_cache() and call it after
  store_user_credential() in both token endpoint and management endpoint
- Unbounded cache: add _BYOK_CRED_CACHE_MAX_SIZE=4096 with clear-on-overflow
- Unbounded auth codes: add _AUTH_CODES_MAX_SIZE=1000 with 503 on overflow
- Double DB query: merge _check_byok_credential + _get_byok_credential into
  single _get_byok_credential call; raise 401 inline if None returned
- Sidebar: remove byok-demo entry (page was deleted in prior commit)
- JWT comment: document why byok_session HS256 token can't be used as proxy auth

* fix: address greptile review feedback (greploop iteration 3)

- auth_type: pre-format Authorization header (Bearer/ApiKey/Basic) in server.py
  before setting ContextVar so openapi_to_mcp_generator respects server auth_type
- cache invalidation on delete: call _invalidate_byok_cred_cache after
  delete_user_credential so stale True entries don't persist for 60s
- ContextVar guard: only set _request_auth_header when mcp_auth_header is set,
  avoiding unnecessary ContextVar overhead on non-BYOK tool calls

* fix: address greptile review feedback (greploop iteration 4)

- Unified credential cache: store actual credential value (Optional[str])
  instead of just bool so _get_byok_credential also benefits from caching —
  eliminates the DB hit on every BYOK tool call within the 60s TTL window
- Extracted _write_byok_cred_cache() helper for consistent cache writes
- Replaced has_user_credential with get_user_credential in _check_byok_credential
  so one DB call satisfies both existence check and value retrieval
- Remove false 'encrypted at rest' claim from OAuth HTML and ByokCredentialModal

* Update tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* Update tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

---------

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
2026-03-04 21:19:25 -08:00
..
_experimental/mcp_server feat(mcp): BYOK MCP servers with OAuth 2.1 PKCE authorization flow (#22850) 2026-03-04 21:19:25 -08:00
agent_endpoints [Feat] Agent RBAC Permission Fix - Ensure Internal Users cannot create agents (#22329) 2026-02-27 18:32:39 -08:00
anthropic_endpoints [Fix] 404 Not Found on /api/event_logging/batch endpoint (#20504) 2026-02-05 10:58:08 -08:00
auth Merge pull request #22678 from Harshit28j/litellm_custom_auth_opt_in 2026-03-04 14:53:44 +05:30
client add a new feature fix to expose the team alias when authenticating th… (#17725) 2025-12-10 10:10:28 -08:00
common_utils fix req change 2026-02-28 16:34:23 +05:30
db [Feat] Add Tool Policies for AI Gateway (#22732) 2026-03-03 20:22:20 -08:00
discovery_endpoints fix(azure): preserve content_policy_violation error details from Azure OpenAI (#20883) 2026-02-10 22:47:03 -08:00
experimental/mcp_server Litellm fix GitHub action testing (#11163) 2025-05-26 14:41:42 -07:00
google_endpoints fix: Metadata / Trace ID Missing in S3 Streaming Callbacks 2026-02-25 14:16:42 +05:30
guardrails azure content enhancement... (#22581) 2026-03-04 10:22:30 -08:00
health_endpoints [Release Fix] (#22411) 2026-02-28 09:46:35 -08:00
hooks Litellm stability fix v2 (#22452) 2026-02-28 15:29:45 -08:00
image_endpoints fixing core proxy tests 2026-02-12 17:54:32 -08:00
management_endpoints Fix tool management tests using wrong field name call_policy 2026-03-04 11:20:17 -03:00
management_helpers [MCP Gateway] Litellm mcp fixes team control (#15304) 2025-10-07 16:48:00 -07:00
middleware feat: add in_flight_requests metric to /health/backlog + prometheus (#22319) 2026-02-27 18:00:50 -08:00
openai_files_endpoint fix(proxy): improve team expiry enforcement validation 2026-03-03 17:29:39 -08:00
pass_through_endpoints Managed batches fixes for Gemini/Vertex 2026-02-28 20:45:16 -05:00
policy_engine Guardrail Policy Versioning (#21862) 2026-02-21 20:14:31 -08:00
prompts fix(prompts): fix prompt info lookup and delete using correct IDs (#19358) 2026-01-20 12:28:34 -08:00
public_endpoints [Feature] Add /public/endpoints endpoint for provider endpoint support 2026-02-26 18:17:37 -08:00
rag_endpoints tests and route permissions (#21508) 2026-02-18 16:58:38 -08:00
response_api_endpoints Fix x-litellm-key-spend update 2025-12-12 11:44:51 +05:30
spend_tracking [Release Fix] (#22411) 2026-02-28 09:46:35 -08:00
test_configs Litellm fix GitHub action testing (#11163) 2025-05-26 14:41:42 -07:00
ui_crud_endpoints fix(ollama): thread api_base to get_model_info + graceful fallback (#21970) 2026-02-23 21:00:37 -08:00
vector_store_endpoints test_delete_vector_store_checks_access 2026-01-31 12:05:09 -08:00
__init__.py test fix 2025-10-17 10:46:42 -07:00
conftest.py Add health endpoint tests to CI with database and Redis support (#17877) 2025-12-12 07:35:50 -08:00
test_aiohttp_cleanup_closed.py fix(aiohttp): only set enable_cleanup_closed when required (#21897) 2026-02-23 21:06:29 -08:00
test_api_key_masking_in_errors.py fix: mask API keys in error responses for invalid/malformed keys (#20289) 2026-02-12 19:58:05 +05:30
test_audio_speech_prometheus_hooks.py fix req changes 2026-02-28 21:32:57 +05:30
test_batch_expiry.py fix(proxy): improve team expiry enforcement validation 2026-03-03 17:29:39 -08:00
test_batch_metadata_none_fix.py Fix issue #13995: Handle None metadata in batch requests (#13996) 2025-08-27 14:51:09 -07:00
test_caching_routes.py [Bug Fix] Ensure /redis/info works on GCP Redis (#11732) 2025-06-14 15:35:09 -07:00
test_chat_completion_metadata.py fix: propagate JWT auth metadata to OTEL spans (#19627) 2026-01-23 21:21:23 -08:00
test_common_request_processing.py feat(proxy): add key_alias, key_hash, requested_model DD APM span tags (#22710) 2026-03-03 20:22:59 -08:00
test_custom_proxy.py fix(ui/): fix routing for custom server root path (#15701) 2025-10-23 13:59:29 -07:00
test_empty_model_list.py [Release Fix] (#22411) 2026-02-28 09:46:35 -08:00
test_enforce_user_param.py Enforce support of enforce_user_param to openai post endpoints 2025-12-03 12:19:21 +05:30
test_fallback_management_endpoints.py Add fallback endpoints support 2026-01-16 10:51:33 +05:30
test_fastapi_offline_routes.py [Bug Fix] - Get Routes (#13466) 2025-08-09 12:52:23 -07:00
test_health_check_functions.py Fix_mapped tests part 2 2026-02-26 12:43:39 +05:30
test_health_check_max_tokens.py add docs and formatting 2026-02-28 14:08:09 +05:30
test_litellm_pre_call_utils.py Agent Tracing - support context_id based trace id propogation + nested llm calls (#22626) 2026-03-03 18:19:12 -08:00
test_model_dump_with_preserved_fields.py Fix_mapped tests part 2 2026-02-26 12:43:39 +05:30
test_model_id_header_propagation.py (fix) propagate x-litellm-model-id in responses (#16986) 2025-11-24 20:40:43 -08:00
test_prometheus_cleanup.py Add Prometheus child_exit cleanup for gunicorn workers 2026-02-27 16:11:15 -08:00
test_proxy_cli.py [Release Fix] (#22411) 2026-02-28 09:46:35 -08:00
test_proxy_server.py Fix inflight mypy 2026-03-02 19:41:32 +05:30
test_proxy_types.py fix: Add PROXY_ADMIN role to system user for key rotation (#21896) 2026-02-27 19:11:29 -08:00
test_proxy_utils.py Fix date overflow/division by zero in proxy utils (#19527) 2026-01-21 21:09:57 -08:00
test_pyroscope.py Fix CI/CD pyroscope test failure (#21219) 2026-02-14 12:07:20 -08:00
test_response_model_sanitization.py fix(proxy): prevent provider-prefixed model leaks (#19943) 2026-01-28 22:26:38 -08:00
test_route_a2a_models.py Fix test_route_a2a_model_bypasses_router 2026-02-05 09:47:05 +05:30
test_route_llm_request.py Override router settings 2026-01-31 16:04:52 -08:00
test_shared_health_check.py Fix_mapped tests part 2 2026-02-26 12:43:39 +05:30
test_spend_log_cleanup.py Fix spend log cleanup: lock tracking, integer retention, skip log level 2026-03-03 10:12:08 -08:00
test_swagger_chat_completions.py [Release Fix] (#22411) 2026-02-28 09:46:35 -08:00
test_team_member_update.py fix mapped tests (#12320) 2025-07-04 10:04:43 -07:00
test_tools_allowlist_enforcement.py [Feat] Add Tool Policies for AI Gateway (#22732) 2026-03-03 20:22:20 -08:00
test_update_llm_router_resilience.py fix(proxy): isolate get_config failures from model loading in sync loop 2026-02-26 17:49:44 -03:00