mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Merge pull request #40440 from BerriAI/litellm_mcp_upstream_error_log_detail
fix(mcp): log upstream request method, body and response on tool-list and OAuth2 token failures
This commit is contained in:
joshua-berri
GitHub
commit
6882f057b4
9 changed files with 680 additions and 8 deletions
|
|
@ -10,6 +10,7 @@ from contextlib import AbstractAsyncContextManager
|
|||
from datetime import timedelta
|
||||
from functools import partial
|
||||
from importlib import metadata
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Protocol, TypeAlias, TypeVar
|
||||
|
||||
import httpx
|
||||
|
|
@ -77,6 +78,7 @@ from litellm._logging import verbose_logger
|
|||
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR, MCP_TOOL_LISTING_TIMEOUT
|
||||
from litellm.experimental_mcp_client.tools import list_tools_with_pagination
|
||||
from litellm.llms.custom_httpx.http_handler import get_ssl_configuration
|
||||
from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response
|
||||
from litellm.types.llms.custom_http import VerifyTypes
|
||||
from litellm.types.mcp import (
|
||||
MCPAuth,
|
||||
|
|
@ -631,7 +633,9 @@ class MCPClient:
|
|||
auth=effective_auth,
|
||||
verify=ssl_config,
|
||||
follow_redirects=True,
|
||||
event_hooks={"request": [guard]} if guard else {},
|
||||
event_hooks=MappingProxyType(
|
||||
{"response": [capture_upstream_error_response], "request": [guard] if guard else []}
|
||||
), # mutable-ok: httpx types require lists of hooks
|
||||
)
|
||||
|
||||
return factory
|
||||
|
|
|
|||
|
|
@ -100,14 +100,24 @@ Usage with curl::
|
|||
http://localhost:4000/mcp/atlassian_mcp
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
from collections.abc import Callable, Mapping
|
||||
import re
|
||||
from collections.abc import AsyncIterator, Callable, Mapping
|
||||
from http.cookies import CookieError, SimpleCookie
|
||||
from itertools import islice
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qsl, quote, quote_plus, unquote_plus, urlencode
|
||||
|
||||
import httpx
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
from starlette.requests import HTTPConnection
|
||||
from starlette.types import Message, Send
|
||||
|
||||
from litellm.litellm_core_utils.secret_redaction import REDACTED, redact_string
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution
|
||||
|
||||
|
|
@ -221,6 +231,10 @@ class MCPDebug:
|
|||
@staticmethod
|
||||
def _mask(value: str | None) -> str:
|
||||
"""Mask a single value for safe display in headers."""
|
||||
return MCPDebug.mask_secret(value)
|
||||
|
||||
@staticmethod
|
||||
def mask_secret(value: str | None) -> str:
|
||||
if not value:
|
||||
return "(none)"
|
||||
return MCPDebug._masker._mask_value(value)
|
||||
|
|
@ -378,3 +392,230 @@ class MCPDebug:
|
|||
server_url=server_url,
|
||||
server_auth_type=server_auth_type,
|
||||
)
|
||||
|
||||
|
||||
_BODY_PREVIEW_CHARS: Final = 512
|
||||
_BODY_CAPTURE_BYTES: Final = 16384
|
||||
_CAPTURE_TIMEOUT_SECONDS: Final = 1.0
|
||||
_CAPTURE_EXTENSION: Final = "litellm_mcp_error_preview"
|
||||
_SAFE_HEADER_NAMES: Final = frozenset({"content-type", "content-length", "accept"})
|
||||
_PUBLIC_HEADER_NAMES: Final = _SAFE_HEADER_NAMES | frozenset(("host", "user-agent", "accept-encoding", "connection"))
|
||||
_JSON_BODY: Final = TypeAdapter(JsonValue)
|
||||
_LOG_MASKER: Final = SensitiveDataMasker(visible_prefix=0, visible_suffix=0)
|
||||
|
||||
|
||||
def _safe_text(value: str, limit: int = _BODY_PREVIEW_CHARS) -> str:
|
||||
escaped: Final = "".join(json.dumps(char)[1:-1] if ord(char) < 32 or ord(char) == 127 else char for char in value)
|
||||
return escaped if len(escaped) <= limit else f"{escaped[:limit]}...(truncated)"
|
||||
|
||||
|
||||
def safe_upstream_url(url: httpx.URL) -> str:
|
||||
return _safe_text(str(url.copy_with(username="", password="", path="/", query=None, fragment=None)))
|
||||
|
||||
|
||||
def _sensitive_field(key: str) -> bool:
|
||||
normalized: Final = re.sub(r"[^a-z0-9]", "", key.casefold())
|
||||
return normalized in ("code", "clientassertion") or any(
|
||||
pattern in normalized for pattern in _LOG_MASKER.sensitive_patterns
|
||||
)
|
||||
|
||||
|
||||
def _redact_object(
|
||||
fields: Mapping[str, JsonValue],
|
||||
) -> dict[str, JsonValue]: # mutable-ok: the standard JSON encoder requires dict objects
|
||||
return { # mutable-ok: construct the JSON object once for the standard parser and encoder
|
||||
key: REDACTED if _sensitive_field(key) else value for key, value in fields.items()
|
||||
}
|
||||
|
||||
|
||||
def _header_secret_values(name: str, value: str) -> tuple[str, ...]:
|
||||
if name == "cookie":
|
||||
cookie: Final = SimpleCookie[str]()
|
||||
try:
|
||||
cookie.load(value)
|
||||
except CookieError:
|
||||
return (value,)
|
||||
return (value, *(item.value for item in cookie.values()))
|
||||
if name not in ("authorization", "proxy-authorization"):
|
||||
return (value,)
|
||||
scheme, _, credential = value.partition(" ")
|
||||
if scheme.lower() != "basic":
|
||||
return (value, credential)
|
||||
try:
|
||||
decoded: Final = base64.b64decode(credential, validate=True).decode("utf-8")
|
||||
except ValueError:
|
||||
return (value, credential)
|
||||
password: Final = decoded.partition(":")[2]
|
||||
return (value, credential, decoded, password, unquote_plus(password))
|
||||
|
||||
|
||||
def _body_secret_values(request: httpx.Request) -> tuple[str, ...] | None:
|
||||
try:
|
||||
raw: Final = request.content
|
||||
except httpx.RequestNotRead:
|
||||
return None
|
||||
if not raw:
|
||||
return ()
|
||||
if len(raw) > _BODY_CAPTURE_BYTES:
|
||||
return None
|
||||
if request.headers.get("content-type", "").split(";", 1)[0].strip().lower() == "application/x-www-form-urlencoded":
|
||||
return tuple(value for key, value in parse_qsl(raw.decode("utf-8", errors="replace")) if _sensitive_field(key))
|
||||
try:
|
||||
body: Final = _JSON_BODY.validate_json(raw)
|
||||
except ValueError:
|
||||
return None
|
||||
from litellm.proxy._experimental.mcp_server.utils import ( # noqa: PLC0415 # MCP utils imports clients; inspect bodies only after initialization
|
||||
json_string_leaves,
|
||||
)
|
||||
|
||||
leaves: Final = json_string_leaves(body)
|
||||
if leaves is None:
|
||||
return None
|
||||
return tuple(
|
||||
value
|
||||
for path, value in leaves
|
||||
if not path or any(isinstance(part, str) and _sensitive_field(part) for part in path)
|
||||
)
|
||||
|
||||
|
||||
def _request_secret_values(request: httpx.Request) -> tuple[str, ...] | None:
|
||||
body_values: Final = _body_secret_values(request)
|
||||
if body_values is None:
|
||||
return None
|
||||
values: Final = (
|
||||
*body_values,
|
||||
request.url.password,
|
||||
*(value for _, value in request.url.params.multi_items()),
|
||||
*(
|
||||
secret
|
||||
for name, value in request.headers.items()
|
||||
if name not in _PUBLIC_HEADER_NAMES
|
||||
for secret in _header_secret_values(name, value)
|
||||
),
|
||||
)
|
||||
return tuple(sorted(frozenset(value for value in values if value), key=len, reverse=True))
|
||||
|
||||
|
||||
def _mask_known_values(value: str, secrets: tuple[str, ...]) -> str:
|
||||
variants: Final = tuple(
|
||||
sorted(
|
||||
frozenset(
|
||||
variant
|
||||
for secret in secrets
|
||||
for variant in (secret, json.dumps(secret)[1:-1], quote(secret, safe=""), quote_plus(secret))
|
||||
),
|
||||
key=len,
|
||||
reverse=True,
|
||||
)
|
||||
)
|
||||
return re.sub("|".join(re.escape(secret) for secret in variants), REDACTED, value) if variants else value
|
||||
|
||||
|
||||
def _preview(raw: bytes, content_type: str = "", secrets: tuple[str, ...] = ()) -> str:
|
||||
if not raw:
|
||||
return "(empty)"
|
||||
if len(raw) > _BODY_CAPTURE_BYTES:
|
||||
return "(omitted: body exceeds capture limit)"
|
||||
try:
|
||||
parsed: Final = _JSON_BODY.validate_python(json.loads(raw, object_hook=_redact_object))
|
||||
except (ValueError, RecursionError):
|
||||
text: Final = raw.decode("utf-8", errors="replace")
|
||||
if (
|
||||
content_type.split(";", 1)[0].strip().lower() != "application/x-www-form-urlencoded"
|
||||
or "=" not in text
|
||||
or any(char in text for char in "<>\n\r")
|
||||
):
|
||||
return "(omitted: unstructured body)"
|
||||
fields: Final = parse_qsl(text, keep_blank_values=True)
|
||||
return _safe_text(
|
||||
_mask_known_values(
|
||||
urlencode(tuple((key, REDACTED if _sensitive_field(key) else value) for key, value in fields)), secrets
|
||||
)
|
||||
)
|
||||
if not isinstance(parsed, (dict, list)):
|
||||
return "(omitted: unstructured body)"
|
||||
return _safe_text(redact_string(_mask_known_values(json.dumps(parsed, separators=(",", ":")), secrets)))
|
||||
|
||||
|
||||
def _masked_headers(headers: httpx.Headers) -> str:
|
||||
return _safe_text(", ".join(f"{name}={value}" for name, value in headers.items() if name in _SAFE_HEADER_NAMES))
|
||||
|
||||
|
||||
def _request_body_preview(request: httpx.Request, secrets: tuple[str, ...] | None) -> str:
|
||||
try:
|
||||
return _preview(request.content, request.headers.get("content-type", ""), secrets or ())
|
||||
except httpx.RequestNotRead:
|
||||
return "(streamed, not captured)"
|
||||
|
||||
|
||||
def _response_body_preview(response: httpx.Response, secrets: tuple[str, ...] | None) -> str:
|
||||
if secrets is None:
|
||||
return "(omitted: request credentials unavailable)"
|
||||
captured: Final = response.extensions.get(_CAPTURE_EXTENSION)
|
||||
if isinstance(captured, str):
|
||||
return captured
|
||||
try:
|
||||
return _preview(response.content, response.headers.get("content-type", ""), secrets)
|
||||
except httpx.ResponseNotRead:
|
||||
return "(not read)"
|
||||
|
||||
|
||||
async def _read_error_prefix(chunks: AsyncIterator[bytes], limit: int) -> bytes:
|
||||
buffer: Final = io.BytesIO()
|
||||
async for chunk in chunks:
|
||||
buffer.write(chunk[: limit - buffer.tell()])
|
||||
if buffer.tell() >= limit:
|
||||
break
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
async def capture_upstream_error_response(response: httpx.Response) -> None:
|
||||
if not response.is_error:
|
||||
return
|
||||
try:
|
||||
prefix: Final = await asyncio.wait_for(
|
||||
_read_error_prefix(response.aiter_bytes(chunk_size=4096), _BODY_CAPTURE_BYTES + 1),
|
||||
timeout=_CAPTURE_TIMEOUT_SECONDS,
|
||||
)
|
||||
response._content = prefix # pyright: ignore[reportPrivateUsage] # rebind-ok: httpx has no public setter to retain consumed bytes for auth retries
|
||||
secrets: Final = _request_secret_values(response.request)
|
||||
preview: Final = (
|
||||
_preview(prefix, response.headers.get("content-type", ""), secrets)
|
||||
if secrets is not None
|
||||
else "(omitted: request credentials unavailable)"
|
||||
)
|
||||
except (asyncio.TimeoutError, httpx.HTTPError, httpx.StreamError):
|
||||
response._content = b"" # pyright: ignore[reportPrivateUsage] # rebind-ok: httpx auth retries must survive diagnostic read failures
|
||||
response.extensions[_CAPTURE_EXTENSION] = (
|
||||
"(unavailable: error body read failed)" # rebind-ok: httpx response hooks communicate through extensions
|
||||
)
|
||||
return
|
||||
response.extensions[_CAPTURE_EXTENSION] = preview # rebind-ok: httpx response hooks communicate through extensions
|
||||
|
||||
|
||||
def describe_upstream_response(response: httpx.Response) -> str:
|
||||
try:
|
||||
request: Final = response.request
|
||||
except RuntimeError:
|
||||
return f"HTTP {response.status_code} | request unavailable"
|
||||
secrets: Final = _request_secret_values(request)
|
||||
return (
|
||||
f"{_safe_text(request.method)} {safe_upstream_url(request.url)} -> HTTP {response.status_code}"
|
||||
f" | request headers: {_masked_headers(request.headers)}"
|
||||
f" | request body: {_request_body_preview(request, secrets)}"
|
||||
f" | response body: {_response_body_preview(response, secrets)}"
|
||||
)
|
||||
|
||||
|
||||
def describe_upstream_http_failure(exc: BaseException) -> str | None:
|
||||
from litellm.proxy._experimental.mcp_server.faults.traversal import ( # noqa: PLC0415 # fault package initialization imports the credential resolver
|
||||
iter_exception_tree,
|
||||
)
|
||||
|
||||
lines: Final = tuple(
|
||||
describe_upstream_response(response)
|
||||
for current in islice(iter_exception_tree(exc), 16)
|
||||
for response in (getattr(current, "response", None),)
|
||||
if isinstance(response, httpx.Response)
|
||||
)
|
||||
return " | ".join(lines) or None
|
||||
|
|
|
|||
|
|
@ -81,7 +81,7 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
|
|||
raise_classified_list_failure,
|
||||
upstream_auth_challenge,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.mcp_debug import record_auth_resolution
|
||||
from litellm.proxy._experimental.mcp_server.mcp_debug import describe_upstream_http_failure, record_auth_resolution
|
||||
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import (
|
||||
MCPPerUserTokenCache,
|
||||
mcp_per_user_token_cache,
|
||||
|
|
@ -1405,6 +1405,11 @@ def _extract_upstream_auth_failure(
|
|||
return upstream_auth_challenge(exc)
|
||||
|
||||
|
||||
def _upstream_failure_suffix(exc: BaseException) -> str:
|
||||
detail: Final = describe_upstream_http_failure(exc)
|
||||
return f"\n upstream exchange: {detail}" if detail else ""
|
||||
|
||||
|
||||
def _obo_retry_applies(server: MCPServer, subject_token: str | None) -> bool:
|
||||
"""Whether an upstream 401/403 should invalidate the minted credential and retry once.
|
||||
|
||||
|
|
@ -4365,7 +4370,9 @@ class MCPServerManager:
|
|||
except MCPServerListError:
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_logger.warning("Failed to get tools from server %s: %s", server.name, e)
|
||||
verbose_logger.warning(
|
||||
"Failed to get tools from server %s: %s%s", server.name, type(e).__name__, _upstream_failure_suffix(e)
|
||||
)
|
||||
raise_classified_list_failure(e, server.name, suppress_challenge=server.is_dcr_bridge)
|
||||
|
||||
async def get_prompts_from_server(
|
||||
|
|
@ -5110,7 +5117,9 @@ class MCPServerManager:
|
|||
verbose_logger.warning("Connection error while listing tools from %s: %s", server_name, e)
|
||||
raise MCPServerListError(ServerListFault(tag="unreachable"), server_name) from e
|
||||
except Exception as e:
|
||||
verbose_logger.warning("Error listing tools from %s: %s", server_name, e)
|
||||
verbose_logger.warning(
|
||||
"Error listing tools from %s: %s%s", server_name, type(e).__name__, _upstream_failure_suffix(e)
|
||||
)
|
||||
raise_classified_list_failure(e, server_name)
|
||||
|
||||
_SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ import httpx
|
|||
from pydantic import BaseModel, ConfigDict, Field, SecretStr, TypeAdapter, ValidationError
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
|
||||
InMemoryTokenCacheBackend,
|
||||
OAuthToken,
|
||||
|
|
@ -101,6 +102,11 @@ async def post_client_credentials_grant(
|
|||
from litellm.llms.custom_httpx.http_handler import ( # noqa: PLC0415 # defer heavy handler import to call time
|
||||
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # handler factory params are coarsely typed
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.mcp_debug import ( # noqa: PLC0415 # diagnostics import credential enums through this package
|
||||
describe_upstream_http_failure,
|
||||
describe_upstream_response,
|
||||
safe_upstream_url,
|
||||
)
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider # noqa: PLC0415 # deferred with the handler import
|
||||
|
||||
try:
|
||||
|
|
@ -110,15 +116,28 @@ async def post_client_credentials_grant(
|
|||
)
|
||||
except httpx.HTTPStatusError as status_err:
|
||||
status_code: Final = status_err.response.status_code
|
||||
verbose_logger.warning(
|
||||
"OAuth2 client_credentials token request denied:\n upstream exchange: %s",
|
||||
describe_upstream_http_failure(status_err),
|
||||
)
|
||||
return TokenEndpointDenied(status_code=status_code, detail=f"token endpoint returned HTTP {status_code}")
|
||||
except Exception as exc: # noqa: BLE001 # any transport failure is the same outcome: unreachable
|
||||
return TokenEndpointUnreachable(detail=str(exc))
|
||||
verbose_logger.warning(
|
||||
"OAuth2 client_credentials POST %s failed: %s", safe_upstream_url(httpx.URL(url)), type(exc).__name__
|
||||
)
|
||||
return TokenEndpointUnreachable(detail=type(exc).__name__)
|
||||
try:
|
||||
body: Final = _TOKEN_BODY_ADAPTER.validate_json(response.content)
|
||||
except ValidationError:
|
||||
verbose_logger.warning("OAuth2 client_credentials invalid response: %s", describe_upstream_response(response))
|
||||
return TokenEndpointDenied(
|
||||
status_code=response.status_code, detail="token endpoint returned a non-JSON-object body"
|
||||
)
|
||||
access_token: Final = body.get("access_token")
|
||||
if not isinstance(access_token, str) or not access_token:
|
||||
verbose_logger.warning(
|
||||
"OAuth2 client_credentials response has no access token | %s", describe_upstream_response(response)
|
||||
)
|
||||
return TokenEndpointSuccess(body=body)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -965,7 +965,9 @@ if MCP_AVAILABLE:
|
|||
apply_tool_filters=apply_tool_filters,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception("Error getting tools from %s: %s", server.name, e)
|
||||
verbose_logger.warning(
|
||||
"Error getting tools from %s: %s", server.name, classify_list_exception(e).tag
|
||||
)
|
||||
return (), classify_list_exception(e)
|
||||
return tools_result, ServerListOk(tool_count=len(tools_result))
|
||||
|
||||
|
|
|
|||
|
|
@ -1895,3 +1895,15 @@ async def test_optional_discovery_preserves_cancellation(method: str) -> None:
|
|||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=3)
|
||||
|
||||
|
||||
|
||||
def test_client_import_before_proxy_credentials_succeeds_in_fresh_process():
|
||||
import subprocess
|
||||
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", "import litellm.experimental_mcp_client.client; from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager; print(MCPServerManager.__name__)"],
|
||||
capture_output=True, text=True, timeout=60, check=False,
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
assert result.stdout.strip() == "MCPServerManager"
|
||||
|
|
|
|||
|
|
@ -478,3 +478,47 @@ async def test_bearer_auth_advertises_the_header_it_will_occupy():
|
|||
assert ClientCredentialsBearerAuth("t", refetch, ClientCredentialsConfig()).header_name == "Authorization"
|
||||
default_carrier = ClientCredentialsConfig(header_name="esb-oauth")
|
||||
assert ClientCredentialsBearerAuth("t", refetch, default_carrier).header_name == "esb-oauth"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mode", ["denied", "invalid", "missing", "success", "timeout", "connect", "cancel"])
|
||||
async def test_token_exchange_failure_diagnostics(mode, monkeypatch, caplog):
|
||||
import asyncio
|
||||
import logging
|
||||
from litellm.llms.custom_httpx import http_handler
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import post_client_credentials_grant
|
||||
|
||||
class Poster:
|
||||
async def post(self, url, headers, data):
|
||||
request = httpx.Request("POST", url, headers=headers, data=data)
|
||||
if mode == "timeout":
|
||||
raise httpx.ReadTimeout("private-transport-message", request=request)
|
||||
if mode == "connect":
|
||||
raise httpx.ConnectError("private-transport-message", request=request)
|
||||
if mode == "cancel":
|
||||
raise asyncio.CancelledError
|
||||
response = httpx.Response(401 if mode == "denied" else 200, request=request,
|
||||
content=b"not-json-private" if mode == "invalid" else None,
|
||||
json=None if mode == "invalid" else {"error": "invalid_client", "client_secret":"first second", **({"access_token":"private-token"} if mode == "success" else {})})
|
||||
response.raise_for_status()
|
||||
return response
|
||||
|
||||
monkeypatch.setattr(http_handler, "get_async_httpx_client", lambda **kwargs: Poster())
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
if mode == "cancel":
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await post_client_credentials_grant("https://idp/token", {}, {})
|
||||
assert not caplog.text
|
||||
return
|
||||
result = await post_client_credentials_grant("https://idp/token?key=query-secret", {"client_secret":"first second"}, {"X-Custom":"header-secret"})
|
||||
for secret in ("first", "second", "query-secret", "header-secret", "private-token", "not-json-private", "private-transport-message"):
|
||||
assert secret not in caplog.text
|
||||
if mode == "success":
|
||||
assert isinstance(result, TokenEndpointSuccess) and result.body["access_token"] == "private-token"
|
||||
assert not caplog.text
|
||||
elif mode in {"timeout", "connect"}:
|
||||
assert isinstance(result, TokenEndpointUnreachable)
|
||||
assert "POST https://idp/ failed" in caplog.text
|
||||
else:
|
||||
assert "POST https://idp/ -> HTTP" in caplog.text
|
||||
assert {"denied":"denied", "invalid":"invalid response", "missing":"no access token"}[mode] in caplog.text
|
||||
|
|
|
|||
|
|
@ -10,9 +10,13 @@ from starlette.types import Message
|
|||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.mcp_debug import (
|
||||
MCP_DEBUG_REQUEST_HEADER,
|
||||
MCPDebug,
|
||||
describe_upstream_http_failure,
|
||||
|
||||
MCPAuthDiagnostics,
|
||||
)
|
||||
|
||||
|
|
@ -206,6 +210,206 @@ class TestWrapSendWithDebugHeaders:
|
|||
assert captured[0] == body_msg
|
||||
|
||||
|
||||
class TestDescribeUpstreamHttpFailure:
|
||||
@staticmethod
|
||||
def _status_error(*, body: bytes, response_body: bytes | None = None) -> httpx.HTTPStatusError:
|
||||
request = httpx.Request(
|
||||
"POST",
|
||||
"https://upstream.example/apis/mcp",
|
||||
headers={"Authorization": "Bearer secret-token-abcdef0123456789", "Content-Type": "application/json" if body.startswith(b"{") else "application/x-www-form-urlencoded"},
|
||||
content=body,
|
||||
)
|
||||
response = (
|
||||
httpx.Response(500, request=request, content=response_body)
|
||||
if response_body is not None
|
||||
else httpx.Response(500, request=request, stream=httpx.ByteStream(b'{"error":"boom"}'))
|
||||
)
|
||||
return httpx.HTTPStatusError("500", request=request, response=response)
|
||||
|
||||
def test_includes_method_url_status_and_request_body(self):
|
||||
exc = self._status_error(
|
||||
body=b'{"method":"initialize","jsonrpc":"2.0","id":0}',
|
||||
response_body=b'{"error":"boom"}',
|
||||
)
|
||||
described = describe_upstream_http_failure(exc)
|
||||
assert described is not None
|
||||
assert "POST https://upstream.example/ -> HTTP 500" in described
|
||||
assert '{"method":"initialize"' in described
|
||||
assert 'response body: {"error":"boom"}' in described
|
||||
|
||||
def test_masks_authorization_header_and_secret_body_fields(self):
|
||||
exc = self._status_error(
|
||||
body=b"grant_type=client_credentials&client_id=abc&client_secret=super-secret-value-1234",
|
||||
response_body=b"{}",
|
||||
)
|
||||
described = describe_upstream_http_failure(exc)
|
||||
assert described is not None
|
||||
assert "secret-token-abcdef0123456789" not in described
|
||||
assert "super-secret-value-1234" not in described
|
||||
assert "client_id=abc" in described
|
||||
assert "client_secret=" in described
|
||||
|
||||
def test_reports_unread_streamed_response_body(self):
|
||||
described = describe_upstream_http_failure(self._status_error(body=b"{}"))
|
||||
assert described is not None
|
||||
assert "response body: (not read)" in described
|
||||
|
||||
def test_finds_response_behind_cause_chain(self):
|
||||
wrapper = RuntimeError("token minting failed")
|
||||
wrapper.__cause__ = self._status_error(body=b"{}", response_body=b'{"error":"invalid_client"}')
|
||||
described = describe_upstream_http_failure(wrapper)
|
||||
assert described is not None
|
||||
assert "invalid_client" in described
|
||||
|
||||
def test_returns_none_without_http_response(self):
|
||||
assert describe_upstream_http_failure(ConnectionError("refused")) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("body", [
|
||||
b'{"password":"first second","token":"demo-secret"}',
|
||||
b'{"nested":[{"access_token":"first,second"}]}',
|
||||
b'client%5Fsecret=first+second&token=demo-secret',
|
||||
])
|
||||
def test_failure_log_fully_redacts_structured_secrets(body):
|
||||
request = httpx.Request("POST", "https://upstream/mcp?credential=query-secret",
|
||||
headers={"X-Custom-Credential": "custom-secret"}, content=body)
|
||||
response = httpx.Response(500, request=request, content=body)
|
||||
detail = describe_upstream_http_failure(httpx.HTTPStatusError("failure", request=request, response=response))
|
||||
assert detail is not None
|
||||
for secret in ("first", "second", "demo-secret", "custom-secret", "query-secret"):
|
||||
assert secret not in detail
|
||||
|
||||
|
||||
def test_failure_log_omits_unstructured_body():
|
||||
request = httpx.Request("POST", "https://upstream/mcp", content=b"arbitrary-secret")
|
||||
response = httpx.Response(500, request=request, content=b"<html>arbitrary-secret</html>")
|
||||
detail = describe_upstream_http_failure(httpx.HTTPStatusError("failure", request=request, response=response))
|
||||
assert detail is not None
|
||||
assert "arbitrary-secret" not in detail
|
||||
assert "omitted" in detail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mode", ["error", "empty", "large", "timeout", "read_failure", "closed", "success", "cancel"])
|
||||
async def test_error_capture_is_bounded_and_preserves_success_and_cancellation(mode):
|
||||
from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response
|
||||
|
||||
class Stream(httpx.AsyncByteStream):
|
||||
def __init__(self):
|
||||
self.reads = 0
|
||||
|
||||
async def __aiter__(self):
|
||||
self.reads += 1
|
||||
if mode == "timeout":
|
||||
await asyncio.sleep(10)
|
||||
if mode == "closed":
|
||||
raise httpx.StreamClosed()
|
||||
if mode == "read_failure":
|
||||
raise httpx.ReadError("private-read-error")
|
||||
if mode == "cancel":
|
||||
raise asyncio.CancelledError
|
||||
yield b"" if mode == "empty" else b'{"error":"missing_scope","password":"first second"}' if mode != "large" else b"x" * 20000
|
||||
|
||||
stream = Stream()
|
||||
request = httpx.Request("POST", "https://upstream/mcp")
|
||||
response = httpx.Response(200 if mode == "success" else 500, request=request, stream=stream)
|
||||
if mode == "cancel":
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await capture_upstream_error_response(response)
|
||||
return
|
||||
await capture_upstream_error_response(response)
|
||||
if mode == "success":
|
||||
assert stream.reads == 0
|
||||
assert await response.aread() == b'{"error":"missing_scope","password":"first second"}'
|
||||
return
|
||||
detail = describe_upstream_http_failure(httpx.HTTPStatusError("failure", request=request, response=response))
|
||||
assert detail is not None
|
||||
assert "first" not in detail and "second" not in detail and "private-read-error" not in detail
|
||||
expected = {"empty": "(empty)", "error": "missing_scope", "large": "capture limit", "timeout": "read failed", "read_failure": "read failed", "closed":"read failed"}
|
||||
assert expected[mode] in detail
|
||||
if mode == "error":
|
||||
assert await response.aread() == b'{"error":"missing_scope","password":"first second"}'
|
||||
|
||||
|
||||
@pytest.mark.parametrize("body", [b"", b'"scalar"', b'{"hint":"line1\\nline2"}', b'{"hint":"' + b'x' * 600 + b'"}'])
|
||||
def test_failure_preview_handles_empty_scalar_control_and_long_bodies(body):
|
||||
request = httpx.Request("POST", "https://user:secret@upstream/mcp?key=private#private", content=body)
|
||||
response = httpx.Response(500, request=request, content=body)
|
||||
detail = describe_upstream_http_failure(httpx.HTTPStatusError("failure", request=request, response=response))
|
||||
assert detail is not None
|
||||
assert "private" not in detail and "user:secret" not in detail and "\n" not in detail
|
||||
if not body:
|
||||
assert "(empty)" in detail
|
||||
elif body.startswith(b'"'):
|
||||
assert "omitted" in detail
|
||||
elif len(body) > 512:
|
||||
assert "truncated" in detail and len(detail) < 1300
|
||||
else:
|
||||
assert "line1\\nline2" in detail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("slow_error", [False, True])
|
||||
async def test_error_capture_preserves_httpx_auth_retry(slow_error):
|
||||
from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response
|
||||
|
||||
class RetryAuth(httpx.Auth):
|
||||
def auth_flow(self, request):
|
||||
response = yield request
|
||||
if response.status_code == 401:
|
||||
request.headers["Authorization"] = "Bearer refreshed"
|
||||
yield request
|
||||
|
||||
class SlowStream(httpx.AsyncByteStream):
|
||||
async def __aiter__(self):
|
||||
await asyncio.sleep(10)
|
||||
yield b'{"error":"expired_token"}'
|
||||
|
||||
def upstream(request):
|
||||
if request.headers.get("Authorization"):
|
||||
return httpx.Response(200, json={"ok": True})
|
||||
return httpx.Response(401, stream=SlowStream()) if slow_error else httpx.Response(401, json={"error":"expired_token"})
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(upstream), auth=RetryAuth(),
|
||||
event_hooks={"response":[capture_upstream_error_response]}) as client:
|
||||
response = await client.get("https://upstream/mcp")
|
||||
assert response.status_code == 200 and response.json() == {"ok":True}
|
||||
if slow_error:
|
||||
assert response.history[0].content == b""
|
||||
else:
|
||||
assert response.history[0].json() == {"error":"expired_token"}
|
||||
|
||||
|
||||
def test_failure_diagnostics_without_request_and_with_streamed_request():
|
||||
response = httpx.Response(503)
|
||||
exc = httpx.HTTPStatusError("failed", request=httpx.Request("GET", "https://upstream"), response=response)
|
||||
assert describe_upstream_http_failure(exc) == "HTTP 503 | request unavailable"
|
||||
request = httpx.Request("POST", "https://upstream", content=iter((b"private-body",)))
|
||||
response = httpx.Response(503, request=request)
|
||||
described = describe_upstream_http_failure(httpx.HTTPStatusError("failed", request=request, response=response))
|
||||
assert described is not None and "streamed, not captured" in described and "private-body" not in described
|
||||
|
||||
|
||||
|
||||
def test_deep_error_body_is_bounded_without_exposing_nested_values():
|
||||
body = b'{"nested":' * 18 + b'{"password":"hidden-value"}' + b'}' * 18
|
||||
request = httpx.Request("POST", "https://upstream/mcp", content=body)
|
||||
response = httpx.Response(500, request=request, content=body)
|
||||
detail = describe_upstream_http_failure(httpx.HTTPStatusError("failure", request=request, response=response))
|
||||
assert detail is not None and "hidden-value" not in detail
|
||||
assert "nested" in detail and "REDACTED" in detail
|
||||
|
||||
|
||||
@pytest.mark.parametrize("body", [b'client%5Fsecret=first+second&client_id=visible', b'client_secret=first%26second&client_id=visible'])
|
||||
def test_encoded_form_credentials_are_decoded_before_redaction(body):
|
||||
request = httpx.Request("POST", "https://upstream/token", content=body,
|
||||
headers={"Content-Type":"application/x-www-form-urlencoded"})
|
||||
response = httpx.Response(400, request=request, content=body,
|
||||
headers={"Content-Type":"application/x-www-form-urlencoded"})
|
||||
detail = describe_upstream_http_failure(httpx.HTTPStatusError("failure", request=request, response=response))
|
||||
assert detail is not None and "client_id=visible" in detail
|
||||
assert "first" not in detail and "second" not in detail
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("source", tuple(AuthResolution))
|
||||
@pytest.mark.parametrize("method", ("GET", "DELETE", "POST"))
|
||||
|
|
@ -291,3 +495,99 @@ async def test_concurrent_mcp_messages_record_on_their_own_http_scope() -> None:
|
|||
await asyncio.gather(record(first, AuthResolution.stored_user_token), record(second, AuthResolution.per_request_header))
|
||||
assert first.resolution() == "stored-user-token"
|
||||
assert second.resolution() == "per-request-header"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("source", ["header", "bearer", "basic", "cookie", "query", "form", "json"])
|
||||
def test_reflected_credentials_are_removed_from_normal_response_fields(source):
|
||||
import base64
|
||||
|
||||
secret = "generic-credential-123"
|
||||
headers = {"X-Custom":secret} if source == "header" else {"Authorization":"Bearer " + secret} if source == "bearer" else {"Authorization":"Basic " + base64.b64encode(("client:" + secret).encode()).decode()} if source == "basic" else {"Cookie":"session=" + secret} if source == "cookie" else {}
|
||||
request = httpx.Request("POST", "https://upstream/token" + ("?credential=" + secret if source == "query" else ""),
|
||||
headers=headers, data={"client_secret":secret} if source == "form" else None,
|
||||
json={"nested":{"client_secret":secret}} if source == "json" else None)
|
||||
response = httpx.Response(401, request=request, json={"error":"invalid_client", "error_description":"Rejected " + secret})
|
||||
detail = describe_upstream_http_failure(httpx.HTTPStatusError("failed", request=request, response=response))
|
||||
assert detail is not None and "invalid_client" in detail
|
||||
assert secret not in detail and "REDACTED" in detail
|
||||
|
||||
|
||||
|
||||
@pytest.mark.parametrize("secret", ['value"with\ncharacters€', "R"])
|
||||
def test_reflected_values_are_redacted_before_truncation_without_expanding_replacements(secret):
|
||||
request = httpx.Request("POST", "https://upstream/token", json={"client_secret":secret})
|
||||
response = httpx.Response(401, request=request, json={"error":"invalid_client", "detail":"x" * 460 + secret})
|
||||
detail = describe_upstream_http_failure(httpx.HTTPStatusError("failed", request=request, response=response))
|
||||
assert detail is not None and "invalid_client" in detail
|
||||
assert "value" not in detail and "characters" not in detail and len(detail) < 1400
|
||||
|
||||
|
||||
@pytest.mark.parametrize("headers", [{"Authorization":"Basic !!!"}, {"Cookie":"bad@key=opaque"}])
|
||||
def test_malformed_auth_headers_do_not_break_failure_diagnostics(headers):
|
||||
request = httpx.Request("POST", "https://upstream/token", headers=headers)
|
||||
response = httpx.Response(401, request=request, json={"error":"invalid_client"})
|
||||
detail = describe_upstream_http_failure(httpx.HTTPStatusError("failed", request=request, response=response))
|
||||
assert detail is not None and "invalid_client" in detail
|
||||
assert "!!!" not in detail and "opaque" not in detail
|
||||
|
||||
|
||||
def test_oversized_request_omits_potentially_reflected_response_credentials():
|
||||
request = httpx.Request("POST", "https://upstream/token", content=b"x" * 17000)
|
||||
response = httpx.Response(401, request=request, json={"error_description":"unknown-secret"})
|
||||
detail = describe_upstream_http_failure(httpx.HTTPStatusError("failed", request=request, response=response))
|
||||
assert detail is not None and "capture limit" in detail and "credentials unavailable" in detail
|
||||
assert "unknown-secret" not in detail
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streamed_error_redacts_reflected_credentials_before_capture():
|
||||
import json
|
||||
from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response
|
||||
|
||||
secret = "generic-credential-123"
|
||||
request = httpx.Request("POST", "https://upstream/token", data={"client_secret":secret})
|
||||
raw = json.dumps({"error":"invalid_client", "error_description":"Rejected " + secret}).encode()
|
||||
response = httpx.Response(401, request=request, stream=httpx.ByteStream(raw))
|
||||
await capture_upstream_error_response(response)
|
||||
detail = describe_upstream_http_failure(httpx.HTTPStatusError("failed", request=request, response=response))
|
||||
assert detail is not None and "invalid_client" in detail and "Rejected" in detail
|
||||
assert secret not in detail and "REDACTED" in detail
|
||||
assert await response.aread() == raw
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", ["/credential-path-value/mcp", "/oauth/credential-path-value/token"])
|
||||
def test_failure_diagnostics_omit_credential_bearing_url_paths(path):
|
||||
request = httpx.Request("POST", "https://upstream.example" + path)
|
||||
response = httpx.Response(401, request=request, json={"error": "access_denied"})
|
||||
error = httpx.HTTPStatusError("denied", request=request, response=response)
|
||||
diagnostic = describe_upstream_http_failure(error)
|
||||
assert diagnostic is not None
|
||||
assert "credential-path-value" not in diagnostic
|
||||
assert "POST https://upstream.example/ -> HTTP 401" in diagnostic
|
||||
assert "access_denied" in diagnostic
|
||||
|
||||
|
||||
def test_deep_request_omits_response_when_credentials_cannot_be_inspected():
|
||||
from litellm.proxy._experimental.mcp_server.utils import MAX_STRUCTURED_CONTENT_SCAN_DEPTH
|
||||
|
||||
raw = "[" * (MAX_STRUCTURED_CONTENT_SCAN_DEPTH + 1) + '{"client_secret":"nested-credential"}' + "]" * (MAX_STRUCTURED_CONTENT_SCAN_DEPTH + 1)
|
||||
request = httpx.Request("POST", "https://upstream/token", content=raw, headers={"Content-Type": "application/json"})
|
||||
response = httpx.Response(401, request=request, json={"error_description": "Rejected nested-credential"})
|
||||
detail = describe_upstream_http_failure(httpx.HTTPStatusError("failed", request=request, response=response))
|
||||
assert detail is not None and "HTTP 401" in detail
|
||||
assert "response body: (omitted: request credentials unavailable)" in detail
|
||||
assert "nested-credential" not in detail
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["accessToken", "refreshToken", "clientSecret", "apikey", "CLIENTASSERTION", "cost_token"])
|
||||
@pytest.mark.parametrize("encoding", ["json", "form"])
|
||||
def test_compact_credential_fields_and_reflected_values_are_redacted(field, encoding):
|
||||
secret = "generic-private-value"
|
||||
fields = {field: secret}
|
||||
request = httpx.Request("POST", "https://upstream/token", json=fields if encoding == "json" else None,
|
||||
data=fields if encoding == "form" else None)
|
||||
response = httpx.Response(401, request=request, json={field: secret, "error": "invalid_client", "detail": "Rejected " + secret})
|
||||
detail = describe_upstream_http_failure(httpx.HTTPStatusError("failed", request=request, response=response))
|
||||
assert detail is not None and "invalid_client" in detail
|
||||
assert "REDACTED" in detail and secret not in detail
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
"""Unit tests for MCP OAuth passthrough tool-fetch behavior."""
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
|
|
@ -11,7 +12,7 @@ if sys.version_info < (3, 11):
|
|||
from exceptiongroup import ExceptionGroup
|
||||
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError
|
||||
from litellm.proxy._experimental.mcp_server.exceptions import MCPServerListError, MCPUpstreamAuthError
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
_extract_upstream_auth_failure,
|
||||
|
|
@ -434,3 +435,43 @@ async def test_aggregate_with_single_accessible_server_still_absorbs():
|
|||
|
||||
assert listing.tools == []
|
||||
assert listing.outcomes["delegate_docs"].tag == "auth_required"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_tools_logs_upstream_request_details_on_500(caplog):
|
||||
manager = MCPServerManager()
|
||||
request = httpx.Request(
|
||||
"POST",
|
||||
"https://upstream/apis/mcp",
|
||||
headers={"Authorization": "Bearer upstream-token-0123456789"},
|
||||
content=b'{"method":"initialize","jsonrpc":"2.0","id":0}',
|
||||
)
|
||||
response = httpx.Response(500, request=request)
|
||||
mock_client = MagicMock()
|
||||
mock_client.list_tools = AsyncMock(
|
||||
side_effect=httpx.HTTPStatusError("500", request=request, response=response)
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
with pytest.raises(MCPServerListError):
|
||||
await manager._fetch_tools_with_timeout(mock_client, "sample_docs")
|
||||
|
||||
assert "POST https://upstream/ -> HTTP 500" in caplog.text
|
||||
assert '"method":"initialize"' in caplog.text
|
||||
assert "upstream-token-0123456789" not in caplog.text
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_creation_failure_logs_sanitized_exchange(monkeypatch, caplog):
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(server_id="sample", name="sample", url="https://upstream/mcp", transport=MCPTransport.http, auth_type=MCPAuth.none)
|
||||
request = httpx.Request("POST", "https://upstream/mcp?credential=query-secret")
|
||||
response = httpx.Response(500, request=request, json={"error":"missing_scope"})
|
||||
error = httpx.HTTPStatusError("query-secret", request=request, response=response)
|
||||
monkeypatch.setattr(manager, "_create_mcp_client", AsyncMock(side_effect=error))
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
with pytest.raises(MCPServerListError):
|
||||
await manager._get_tools_from_server(server)
|
||||
assert "POST https://upstream/ -> HTTP 500" in caplog.text
|
||||
assert "missing_scope" in caplog.text and "query-secret" not in caplog.text
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue