From 8f902f925f92771ef505c46949f559f9b6dc6670 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 4 Mar 2026 21:00:21 -0800 Subject: [PATCH] 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 --- .../mcp_server/openapi_to_mcp_generator.py | 7 +++++-- litellm/proxy/_experimental/mcp_server/server.py | 13 ++++++++++++- .../mcp_management_endpoints.py | 4 ++++ 3 files changed, 21 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 3ac014b0439..bcbf91e5c56 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -219,11 +219,14 @@ def create_tool_function( The function safely handles parameter names that aren't valid Python identifiers by using **kwargs instead of named parameters. """ - # Allow per-request auth override (e.g. BYOK credential set via ContextVar) + # Allow per-request auth override (e.g. BYOK credential set via ContextVar). + # The ContextVar holds the full Authorization header value, including the + # correct prefix (Bearer / ApiKey / Basic) formatted by the caller in + # server.py based on the server's configured auth_type. effective_headers = dict(headers) override_auth = _request_auth_header.get() if override_auth: - effective_headers["Authorization"] = f"Bearer {override_auth}" + effective_headers["Authorization"] = override_auth # Build URL from base_url and path url = base_url + path diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 6e9e24b3d46..d97193c06d7 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1743,7 +1743,18 @@ if MCP_AVAILABLE: verbose_logger.debug(f"Executing local registry tool: {name}") # For BYOK servers the credential must be injected via a ContextVar # because the tool function has headers baked into its closure. - _auth_token = _request_auth_header.set(mcp_auth_header) + # Pre-format the full Authorization header value using the server's + # configured auth_type so the generator doesn't need to know the prefix. + auth_header_value: Optional[str] = None + if mcp_auth_header: + server_auth_type = getattr(mcp_server, "auth_type", None) if mcp_server else None + if server_auth_type == MCPAuth.api_key: + auth_header_value = f"ApiKey {mcp_auth_header}" + elif server_auth_type == MCPAuth.basic: + auth_header_value = f"Basic {mcp_auth_header}" + else: + auth_header_value = f"Bearer {mcp_auth_header}" + _auth_token = _request_auth_header.set(auth_header_value) try: local_content = await _handle_local_mcp_tool(name, arguments) finally: diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 912c161418d..b48db72a536 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1128,6 +1128,10 @@ if MCP_AVAILABLE: await delete_user_credential(prisma_client, user_id, server_id) except Exception: pass # Already deleted or didn't exist + from litellm.proxy._experimental.mcp_server.server import ( + _invalidate_byok_cred_cache, + ) + _invalidate_byok_cred_cache(user_id, server_id) return MCPUserCredentialResponse(server_id=server_id, has_credential=False) @router.put(