- Add get_cached_byok_credential() public helper to server.py to expose cache
reads without tight coupling to private internals (_byok_cred_cache, TTL const)
- Update status endpoint to use get_cached_byok_credential() instead of
importing private cache internals directly
- Move prisma_client import in callback to after all early-exit paths so it is
only imported when actually needed for token exchange
- Add test_callback_url_encoded_success_response: verifies form-encoded token
response (GitHub default) is stored as plain string
- Add test_callback_url_encoded_error_response: verifies form-encoded error body
returns 502 with error key in HTML response
- Replace fragile "?" in url check with urlparse/urlunparse for safe query
string merging; handles fragments, trailing separators, existing params
- Cache None after status DB query when connected=False: reduces DB load from
~15 queries/user during 30s OAuth consent screen to just 1 (evicted on callback)
- Add move_to_end() on cache reads in _get_byok_credential and _check_byok_credential
for true LRU semantics: frequently-accessed entries stay alive in cache
- Move response.raise_for_status() inside httpx async context manager to avoid
ResponseClosed errors on streaming responses after context exit
- Require explicit user_id in /connect and /status: remove api_key fallback to
prevent credential sharing across users and orphaning on key rotation
- Switch _byok_cred_cache from plain dict to OrderedDict with move_to_end() for
true LRU eviction — frequently-accessed entries are no longer evicted early
- Consume state token on provider error redirect (user denied) so orphaned
entries don't fill the 1000-entry store and trigger 503 for all users
- Remove overly permissive "or b'provider' fallback in callback test assertion;
the primary check (b'invalid_grant' in body) is sufficient and specific
- Surface 401/403 auth errors in polling loop immediately instead of silently
retrying for 10 minutes; user sees "Session expired" message right away
- Remove "" sentinel write from status endpoint: only real token fetches
(via _get_byok_credential / _check_byok_credential) write to cache.
Status reads cache when real token present, falls through to DB otherwise.
Eliminates dual-interpretation of same cache key across status/auth paths.
- Revert _get_byok_credential and _check_byok_credential: no longer need
special "" handling since "" can no longer appear in cache
- Fix _purge_expired_states ordering in callback: check state first so
expired tokens get "State expired" error, not misleading "Invalid state"
- Fix test_connect_missing_client_secret_raises_400: replace dead-code
master_key mock with prisma_client mock to make validation-order-independent
- Remove handleError() from getMcpOAuth2ConnectUrl: caller (OAuth2ConnectButton)
already calls setError(), so double notification is now eliminated
- Fix _check_byok_credential to treat "" sentinel as cache miss: consistent with
_get_byok_credential, preventing a race where deleted credentials pass auth check
but return None from token lookup, causing silent 401 to backend
- Remove vestigial master_key import/check from openapi_oauth2_connect: state
tokens are now pure random (secrets.token_urlsafe), prisma_client check already
rejects unconfigured deployments
- Add token_endpoint_auth_method field to MCPServer ("post" default, "basic" for Okta/Auth0)
- Implement HTTP Basic auth branch in token exchange (client_secret_basic per RFC 6749 §2.3.1)
- Use patch.object(litellm, "require_byok_credential_store", True) in 503 tests to avoid global mutation
- Add callback_url to all three callback test state entries (exercises primary path, not fallback)
- Add test asserting Basic auth creds go in Authorization header, not POST body
- _write_byok_cred_cache: only evict when the key is new (not already present),
so updating an existing entry doesn't unnecessarily displace another user's
credential when the cache is at capacity.
- getMcpOAuth2Status: remove the wrapping try-catch that just re-throws
without any side effect (adds noise with no handling benefit).
- Improve test_execute_mcp_tool_uses_user_api_key_dict_as_fallback to also
verify the non-fallback path (explicit user_api_key_auth takes precedence).
The status endpoint writes '' to _byok_cred_cache for connected=True (no
actual token needed for status checks). Without this fix, _get_byok_credential
would return '' as the credential, causing an empty Bearer header and silent
401s on tool calls after the user just authorized.
Treat '' as a cache miss and fall through to the DB to fetch the real token.
- Status endpoint now writes to _byok_cred_cache after DB query so subsequent
2s polls are served from memory (cache read was correct but write was missing,
making every pre-auth poll hit the DB).
- Convert test_spec_path_server_uses_tool_registry to async to fix deprecated
asyncio.get_event_loop() usage that fails on Python 3.12+.
- Improve test_execute_mcp_tool_uses_user_api_key_dict_as_fallback to actually
exercise the rest_endpoints.py fallback expression and assert the correct value.
- Store callback_url in _pending_oauth2_states during /connect so /callback
reuses the exact same redirect_uri, preventing mismatches when the two
requests are routed differently through reverse proxies.
- getMcpOAuth2ConnectUrl: use .catch(() => ({})) for error JSON parse to
avoid SyntaxError on non-JSON error bodies (e.g. HTML 502 pages).
- Add regression test: callback_url stored in state dict matches retrieved value.
- Add regression test: user_api_key_auth fallback from user_api_key_dict.
- Add _get_callback_base_url() helper that prefers LITELLM_PROXY_BASE_URL
env var over X-Forwarded-Host-derived URL, avoiding header spoofing in
the security-sensitive OAuth2 redirect_uri construction.
- OAuth2ConnectButton.tsx: close popup on component unmount so orphaned
popups don't persist when the user navigates away.
- getMcpOAuth2Status: remove handleError() call before throwing so the
polling loop's catch block is the single error surface and no unexpected
UI notifications fire on transient 2-second poll failures.
- Error HTML page now auto-closes the popup after 4s (same as success
page), stopping the 2s polling loop and clearing the loading spinner
so users aren't left in a confused state after a provider error.
- Replace trivial `assert manager is not None` with a meaningful assertion:
patch _create_mcp_client and verify it is NOT called for spec_path servers,
directly testing the short-circuit in _get_tools_from_server.
- /connect now fails immediately with 503 if prisma_client is None, so
misconfigured deployments don't waste users through the full provider
consent flow before failing at /callback.
- Fix "LRU-style" → "FIFO" in _write_byok_cred_cache docstring since
eviction uses insertion order (first-inserted), not recency.
- Fix malformed authorization URL when base URL already has query params
(e.g. Google: ?access_type=offline). Use & instead of ? in that case.
- Simplify _make_state_token() to take no parameters since none were used.
Update call site and tests accordingly.
- Status endpoint now checks _byok_cred_cache before querying the DB,
avoiding a raw Prisma query on every 2-second poll from the UI.
The callback's _invalidate_byok_cred_cache call ensures the first poll
after successful auth always falls through to DB and returns connected=True.
- require_byok_credential_store now defaults to False so existing deployments
without a database are not broken on upgrade. Set True to opt into the strict
503-on-no-DB behavior.
- Add comment clarifying that get_user_credential() is the centralized DB helper
for LiteLLM_MCPUserCredentials, mitigated by 60s in-memory cache.
- Update tests to explicitly set the flag to True when asserting 503 behavior.
Add a feature flag `litellm.require_byok_credential_store` (default True)
that controls whether _get_byok_credential and _check_byok_credential raise
HTTP 503 or fall through to legacy behavior when prisma_client is None.
Default True preserves the new secure behavior (503 = infra problem, distinct
from 401 = missing credential). Set to False to restore legacy silent-bypass
for stateless/no-database deployments.
- _make_state_token now uses secrets.token_urlsafe(32) instead of an
HMAC digest that was never re-verified by the callback (dict lookup
only), making the intent explicit and removing misleading crypto.
- Fix misleading comment on managed MCP tool call: original_tool_name
is the unprefixed name, not the "full (potentially prefixed)" name.
- Document the known stale-token window trade-off in _BYOK_CRED_CACHE_TTL
so future contributors understand the failure mode.
- Update state-token tests to reflect the new implementation.
- _get_byok_credential now raises HTTP 503 when prisma_client is None,
matching _check_byok_credential behavior. Previously it returned None
silently, causing callers to surface a misleading 401 instead of the
correct 503 infrastructure error.
- execute_mcp_tool prefixed-name fallback now checks name doesn't already
start with the server prefix before calling add_server_prefix_to_name,
preventing double-prefixed names like "github-github-get_user".
- Add MCP_TOOL_PREFIX_SEPARATOR to server.py utils imports.
- Add regression tests for both fixes.
- _check_byok_credential now raises HTTP 503 when prisma_client is None
instead of silently returning; previously this bypassed BYOK identity
enforcement entirely when the database was not configured
- Add regression tests:
- test_check_byok_credential_raises_503_when_no_db: verifies 503 is
raised (not silent bypass) when credential store is unavailable
- test_spec_path_server_uses_tool_registry: documents spec_path
short-circuit invariant (OpenAPI servers bypass MCP client creation)
- test_mcp_server_byok_fields_propagated/default: cover the YAML
config byok-field propagation bug fix
_check_byok_credential wrote the raw credential from DB directly to
_byok_cred_cache without calling _extract_access_token. When the OAuth2
callback stores a JSON blob {"access_token": "...", "refresh_token": "..."},
a subsequent call that went through _check_byok_credential first (rather
than _get_byok_credential) would populate the cache with the blob. The
next _get_byok_credential call would return the blob as the Bearer token,
causing silent HTTP 401s on every upstream API call.
Fix: extract the access_token before writing to cache, consistent with
how _get_byok_credential already handles it.
- Replace _byok_cred_cache.clear() with single-entry eviction in
_write_byok_cred_cache: at capacity, evict only the oldest key
instead of wiping all 4096 entries (thundering-herd prevention)
- Add regression tests:
- Cache LRU eviction: verifies exactly 1 entry evicted at capacity
- MCPServer byok fields: is_byok/byok_description/byok_api_key_help_url
propagated correctly (regression for the load_servers_from_config
missing-fields bug that caused BYOK servers to lose their fields)
- Cache capacity: tests all paths in _write_byok_cred_cache
1. Separate try-blocks for store_user_credential and
_invalidate_byok_cred_cache: a cache-flush failure no longer
returns a 'Storage error' page when the write succeeded
2. Popup-close race fix in OAuth2ConnectButton: do one final status
check when popup is detected closed so fast OAuth flows (popup
auto-closes before the first 2s poll fires) are not silently missed
3. Fix test patches: `master_key` and `prisma_client` are imported
inline in the function body via `from proxy_server import X`;
patch `litellm.proxy.proxy_server.*` not the endpoint module
- Extract and store refresh_token alongside access_token when the
provider returns one; stored as JSON blob {"access_token": ...,
"refresh_token": ...} so users aren't forced back through OAuth2
consent when the access token expires
- Add _extract_access_token() helper in server.py that transparently
handles both the new JSON blob format and legacy plain-string
credentials (backward compatible)
- Add tests: refresh_token stored as JSON, plain token when no
refresh_token, _extract_access_token with all input shapes
- Add PKCE comment acknowledging RFC 9700 gap; confidential client
with client_secret is lower risk, full PKCE tracked as follow-up
When a server has is_byok=true and auth_type=oauth2, the Credentials
column in the MCP Servers table shows an OAuth2ConnectButton instead
of the static key entry modal.
- OAuth2ConnectButton: calls /v1/mcp/server/{id}/oauth2/connect,
opens the returned authorization_url in a popup, polls
/v1/mcp/server/{id}/oauth2/status every 2 s until connected=true,
then shows a Connected badge and calls onConnected() to refresh the table
- mcp_server_columns: branches on auth_type===oauth2 to render the
new button, passing accessToken and refreshServers
- mcp_servers: passes accessToken and refetch down to the columns factory
- networking: adds getMcpOAuth2ConnectUrl and getMcpOAuth2Status helpers
Three bugs prevented stored OAuth2 tokens from being injected as
Bearer headers when OpenAPI MCP tools were called:
1. _get_tools_from_server called _create_mcp_client before checking
spec_path, so for BYOK OAuth2 servers resolve_mcp_auth tried a
client_credentials token exchange (GitHub rejects this with a non-JSON
body), causing a JSON parse error that made tools/list return [].
Fix: check spec_path first and read from tool registry directly,
skipping MCP client creation entirely for OpenAPI servers.
2. REST tool call passed user_api_key_auth=data.get("user_api_key_auth")
which is None unless set in request metadata, so _get_byok_credential
couldn't look up the stored token by user_id.
Fix: fall back to user_api_key_dict from the route dependency.
3. OpenAPI tools are registered in the local registry under the prefixed
name (e.g. github_user-get_authenticated_user) but callers may use
the bare name (get_authenticated_user). get_tool(bare_name) returned
None and execution fell through to the managed MCP client path.
Fix: when local_tool is None and we know the server has a spec_path,
retry get_tool with the prefixed name.
Also fixes load_servers_from_config to propagate is_byok, byok_description,
and byok_api_key_help_url from YAML config into the MCPServer object.
Adds three new endpoints so users can authorize their own accounts
through a provider's OAuth2 consent screen (GitHub, Spotify, Linear, etc.)
instead of pasting static API keys for BYOK OpenAPI MCP servers.
- GET /v1/mcp/server/{server_id}/oauth2/connect — initiates the flow,
returns an authorization_url the UI opens as a popup
- GET /v1/mcp/oauth2/callback — receives code+state from provider,
exchanges for access token, stores in LiteLLM_MCPUserCredentials
- GET /v1/mcp/server/{server_id}/oauth2/status — returns connected:true/false
State tokens are HMAC-SHA256 signed with master_key and expire after 10 min.
The callback shows a success HTML page (with auto-close for popups) instead
of redirecting, avoiding 404s in environments without a full UI deploy.
Claude Code v2.1.69+ sends `custom: {defer_loading: true}` on tool
definitions. Anthropic's API accepts this field, but Bedrock rejects it
with "Extra inputs are not permitted", causing ~90% of requests to fail.
Strip the `custom` field from each tool in the request body before
sending to Bedrock, in both the Messages API and Chat API invoke paths.
Fixes#22847
Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com>
* fix(proxy): readiness check returns 200 when database is unreachable
_db_health_readiness_check() catches health_check() exceptions but
never updates db_health_cache to "disconnected" and never re-raises.
The caller health_readiness() always returns 200 with "db": "connected"
hardcoded, regardless of actual DB state.
In Kubernetes, this means pods with dead database connections stay in
the Service endpoints and continue receiving traffic they cannot serve.
Changes:
- Set db_health_cache to "disconnected" and re-raise the exception on
health_check failure so health_readiness() returns 503
- Use actual db_health_status["status"] in the response instead of
hardcoding "db": "connected"
- Reduce cache TTL from 2 minutes to 15 seconds. The 2-minute window
is too wide for readiness probes (typically 10-15s intervals) and
means a pod can report healthy for up to 2 minutes after the DB dies
- Only serve cached results when status is "connected". The previous
condition (status != "unknown") would also cache "disconnected" for
2 minutes, delaying recovery detection after a DB comes back
* fix(proxy): add DB connection self-healing to readiness check
When the Prisma query engine's internal TCP connection pool holds dead
connections (caused by network blips, Cloud SQL proxy restarts, or
node-level issues), health_check() fails with httpx.ConnectError.
The engine never recovers on its own because nothing triggers a
disconnect/connect cycle to restart the subprocess with fresh
connections.
This leaves pods permanently failing readiness checks until they are
manually restarted, even after the underlying DB becomes reachable
again.
Add a reconnect attempt to _db_health_readiness_check() when
health_check() fails:
1. disconnect() - kills the query engine subprocess and closes all
connections (has built-in backoff retry: 3 tries, 10s max)
2. connect() - starts a new engine with fresh TCP connections (has
built-in backoff retry: 3 tries, 10s max)
3. health_check() - verifies the new connection works (has built-in
backoff retry: 3 tries, 10s max)
If reconnect succeeds, the pod immediately returns to service (200).
If it fails, the original exception is re-raised (503). Reconnect
attempts are rate-limited by probe frequency (~10-15s), so a
permanently unreachable DB gets one attempt per cycle with no retry
loops.
This uses the same disconnect/connect mechanism that
PrismaWrapper.recreate_prisma_client() uses for IAM token refresh,
and aligns with the community-documented pattern for Prisma connection
recovery in long-running processes (prisma/prisma#24718, #27024).
* Add poetry lock and modify test_health_endpoints
* Address allow_requests_on_db_unavailable regression
* Address comments
* resolve greptile issue
* Restore accidentally deleted UI HTML files
These were removed in an earlier commit but still exist on main.
Restoring to keep the PR diff clean.
* Guard reconnect with is_database_transport_error
Only attempt disconnect/connect/health_check cycle for transport-level
failures (unreachable DB, dropped connection). Data-layer errors like
UniqueViolationError indicate the DB is reachable, so reconnecting
would be pointless churn.
* Address greptile's comments
* Fix module alias after rebase and add adversarial test coverage
- Unify module alias to _health_endpoints_module after rebase conflict
- Add test for non-transport error with flag on (exercises is_database_transport_error guard)
- Add test for disconnect() failure during reconnect cycle
- Split non-transport error test into flag-off (re-raises) and flag-on (skips reconnect) variants
* Remove stale UI HTML files reintroduced during rebase
* fix: don't close HTTP/SDK clients on LLMClientCache eviction
Removing the _remove_key override that eagerly called aclose()/close()
on evicted clients. Evicted clients may still be held by in-flight
streaming requests; closing them causes:
RuntimeError: Cannot send a request, as the client has been closed.
This is a regression from commit fb72979432. Clients that are no longer
referenced will be garbage-collected naturally. Explicit shutdown cleanup
happens via close_litellm_async_clients().
Fixes production crashes after the 1-hour cache TTL expires.
* test: update LLMClientCache unit tests for no-close-on-eviction behavior
Flip the assertions: evicted clients must NOT be closed. Replace
test_remove_key_closes_async_client → test_remove_key_does_not_close_async_client
and equivalents for sync/eviction paths.
Add test_remove_key_removes_plain_values for non-client cache entries.
Remove test_background_tasks_cleaned_up_after_completion (no more _background_tasks).
Remove test_remove_key_no_event_loop variant that depended on old behavior.
* test: add e2e tests for OpenAI SDK client surviving cache eviction
Add two new e2e tests using real AsyncOpenAI clients:
- test_evicted_openai_sdk_client_stays_usable: verifies size-based eviction
doesn't close the client
- test_ttl_expired_openai_sdk_client_stays_usable: verifies TTL expiry
eviction doesn't close the client
Both tests sleep after eviction so any create_task()-based close would
have time to run, making the regression detectable.
Also expand the module docstring to explain why the sleep is required.
* docs(AGENTS.md): add rule — never close HTTP/SDK clients on cache eviction
* docs(CLAUDE.md): add HTTP client cache safety guideline
* Include user_email in new user creation within get_user_object
Enhance the get_user_object function to include user_email in the parameters when creating a new user. This change is accompanied by a new test to verify that user_email is correctly included during the upsert process.
* Improve error handling in test_get_user_object by logging exceptions
Updated the test_get_user_object_upsert_includes_user_email function to log exceptions when they occur, enhancing the visibility of potential issues during testing. This change helps in diagnosing failures related to the mock LiteLLM_UserTable.
