litellm/tests
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
..
agent_tests [Release Fix] (#22411) 2026-02-28 09:46:35 -08:00
audio_tests [Feat] New provider TTS - Add AWS polly API for TTS (#18326) 2025-12-22 18:19:34 +05:30
basic_proxy_startup_tests Fix: remove merge markdown (#17586) 2025-12-06 05:38:16 -08:00
batches_tests Merge pull request #22625 from BerriAI/litellm_azure_ai_finetune 2026-03-03 19:42:17 +05:30
code_coverage_tests Merge pull request #22752 from BerriAI/litellm_search_api_add 2026-03-04 18:29:10 +05:30
documentation_tests fix failing tests 2026-02-21 15:48:26 -08:00
enterprise fix: support list of modes in Mode.default for tag-based guardrails 2026-03-04 01:50:28 +05:30
guardrails_tests fix: bump litellm-proxy-extras to 0.4.50 and fix 3 failing tests (#22417) 2026-02-28 10:20:03 -08:00
image_gen_tests [Release - 02/10/2026] v1.81.10-nightly 2026-02-10 16:26:30 -08:00
litellm Merge pull request #21233 from Chesars/feat/per-request-json-schema-validation 2026-03-03 15:29:54 -03:00
litellm-proxy-extras fix: resolve 'does not exist' migration errors as applied in setup_database (#19281) 2026-01-26 22:11:36 -08:00
litellm_core_utils Fix: empty assistant message for converse API 2026-02-04 09:50:58 +05:30
litellm_utils_tests Fix test_perform_health_check_filters_by_model_id 2026-02-26 10:43:05 +05:30
llm_responses_api_testing Fix anthropic responses 2026-02-20 17:30:42 -08:00
llm_translation fix(types): normalize null top_logprobs to [] in ChatCompletionTokenLogprob (#22245) 2026-02-28 15:59:29 -08:00
load_tests Add memory leak detection tests with CI integration (#18881) 2026-01-09 17:36:10 -08:00
local_testing Add support for wildcards models for files api 2026-03-04 11:48:42 +05:30
logging_callback_tests Merge pull request #22180 from BerriAI/litellm_fix_vllm_test 2026-02-26 18:43:52 +05:30
mcp_tests Fix flaky MCP streaming test by properly mocking inner aresponses call 2026-03-04 11:09:24 -03:00
multi_instance_e2e_tests
ocr_tests OCR test fixes 2026-01-11 08:00:31 -08:00
old_proxy_tests/tests remove prompt caching headers as the support has been removed 2026-01-02 11:08:35 +05:30
openai_endpoints_tests test(responses): add end-to-end test for responses API WebSocket mode 2026-03-02 17:24:39 +05:30
otel_tests test fix 2026-01-31 19:08:07 -08:00
pass_through_tests Litellm stability fix v2 (#22452) 2026-02-28 15:29:45 -08:00
pass_through_unit_tests [Release Fix] (#22411) 2026-02-28 09:46:35 -08:00
proxy_admin_ui_tests security: fix critical/high CVEs in OS-level libs and NPM transitive 2026-02-24 19:40:09 +05:30
proxy_e2e_anthropic_messages_tests Fix: litellm/tests/llm_responses_api_testing/test_anthropic_responses_api.py 2026-02-20 17:30:53 -08:00
proxy_e2e_azure_batches_tests Add tenacity to e2e Azure batch CI and revert importorskip 2026-03-04 11:45:14 -03:00
proxy_security_tests
proxy_unit_tests Merge pull request #22372 from BerriAI/litellm_jwt_vkey_map 2026-03-05 06:24:49 +05:30
router_unit_tests [Release Fix] (#22411) 2026-02-28 09:46:35 -08:00
scim_tests
search_tests Add doc and tests for google search api 2026-03-04 13:55:43 +05:30
spend_tracking_tests increasing time for spend tracking 2026-01-20 16:31:25 -08:00
store_model_in_db_tests fix: test mock 2026-01-02 17:38:52 +09:00
test_litellm feat(mcp): BYOK MCP servers with OAuth 2.1 PKCE authorization flow (#22850) 2026-03-04 21:19:25 -08:00
unified_google_tests Revert "[Infra] Changing Google Tests to use Gemini 3 Flash Preview" 2026-01-20 17:32:10 -08:00
vector_store_tests Add support for vertex ai for rag/ingest 2026-02-13 20:11:10 +05:30
windows_tests
__init__.py
gettysburg.wav
large_text.py
openai_batch_completions.jsonl
README.MD
test_budget_management.py
test_callbacks_on_proxy.py Fix CI: Revert security scan changes and add GitGuardian ignore rules (#18358) 2025-12-22 17:03:53 -08:00
test_config.py
test_debug_warning.py
test_default_encoding_non_root.py fix: resolve Read-only file system error in non-root images (#19449) 2026-01-20 19:00:52 -08:00
test_end_users.py
test_entrypoint.py
test_fallbacks.py Revert "fix: prevent error when max_fallbacks exceeds available models (#20071)" 2026-02-03 15:15:30 +05:30
test_gpt5_azure_temperature_support.py
test_health.py
test_keys.py Adding retries to flaky tests 2026-01-22 15:21:44 -08:00
test_litellm_proxy_responses_config.py
test_logging.conf
test_models.py
test_openai_endpoints.py
test_organizations.py Adding retries to flaky tests 2026-01-22 15:21:44 -08:00
test_otel_thread_leak.py Fix thread leak in OpenTelemetry dynamic header path (#19946) 2026-01-28 10:35:37 -08:00
test_passthrough_endpoints.py
test_presidio_latency.py fix(presidio): reuse HTTP connections to prevent OOMs (#19964) 2026-01-28 16:08:53 -08:00
test_proxy_server_non_root.py deactivating non root tests 2026-01-23 22:55:36 -08:00
test_ratelimit.py
test_resource_cleanup.py
test_service_logger_otel.py fix(langfuse_otel): prevent empty proxy request spans from being sent to Langfuse 2026-01-28 15:35:35 +01:00
test_spend_logs.py Fix CI: Revert security scan changes and add GitGuardian ignore rules (#18358) 2025-12-22 17:03:53 -08:00
test_team.py fix(test): skip 'projects' field in team update assertion (#21777) 2026-02-21 10:24:53 -08:00
test_team_logging.py
test_team_members.py
test_users.py fix(tests): update deprecated Anthropic model in test_user_model_access (#21826) 2026-02-21 14:18:24 -08:00

In total litellm runs 1000+ tests

[02/20/2025] Update:

To make it easier to contribute and map what behavior is tested,

we've started mapping the litellm directory in tests/test_litellm

This folder can only run mock tests.