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
This commit is contained in:
Ishaan Jaffer 2026-03-04 21:00:21 -08:00
parent 5f3c4ed573
commit 8f902f925f
3 changed files with 21 additions and 3 deletions

View file

@ -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

View file

@ -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:

View file

@ -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(