diff --git a/litellm/proxy/_experimental/mcp_server/mcp_debug.py b/litellm/proxy/_experimental/mcp_server/mcp_debug.py index 07e77856894..ccce0aaf615 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_debug.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_debug.py @@ -101,21 +101,24 @@ Usage with curl:: """ import asyncio +import base64 +import io import json +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, urlencode +from urllib.parse import parse_qsl, quote, quote_plus, unquote_plus, urlencode import httpx -from pydantic import JsonValue, TypeAdapter, ValidationError +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.faults.traversal import iter_exception_tree from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution # Header the client sends to opt into debug mode @@ -396,6 +399,7 @@ _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) @@ -413,33 +417,105 @@ def _sensitive_field(key: str) -> bool: return key.lower() in ("code", "cookie", "client_assertion") or _LOG_MASKER.is_sensitive_key(key) -def _redact_json(value: JsonValue, depth: int = 0) -> str: - if depth >= 16: - return json.dumps("(depth limit)") - if isinstance(value, dict): - return ( - "{" - + ",".join( - json.dumps(key) - + ":" - + (json.dumps(REDACTED) if _sensitive_field(key) else _redact_json(item, depth + 1)) - for key, item in value.items() - ) - + "}" +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, ) - if isinstance(value, list): - return "[" + ",".join(_redact_json(item, depth + 1) for item in value) + "]" - return json.dumps(redact_string(value) if isinstance(value, str) else value) + ) + return re.sub("|".join(re.escape(secret) for secret in variants), REDACTED, value) if variants else value -def _preview(raw: bytes, content_type: str = "") -> str: +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_json(raw) - except ValidationError: + 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" @@ -449,41 +525,45 @@ def _preview(raw: bytes, content_type: str = "") -> str: return "(omitted: unstructured body)" fields: Final = parse_qsl(text, keep_blank_values=True) return _safe_text( - urlencode( - tuple((key, REDACTED if _sensitive_field(key) else redact_string(value)) for key, value in fields) + _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_json(parsed)) + 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) -> str: +def _request_body_preview(request: httpx.Request, secrets: tuple[str, ...] | None) -> str: try: - return _preview(request.content, request.headers.get("content-type", "")) + 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) -> str: +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", "")) + return _preview(response.content, response.headers.get("content-type", ""), secrets) except httpx.ResponseNotRead: return "(not read)" -async def _read_error_prefix(chunks: AsyncIterator[bytes], remaining: int) -> bytes: - chunk: Final = await anext(chunks, b"") - if not chunk or len(chunk) >= remaining: - return chunk[:remaining] - return chunk + await _read_error_prefix(chunks, remaining - len(chunk)) +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: @@ -495,8 +575,13 @@ async def capture_upstream_error_response(response: httpx.Response) -> None: timeout=_CAPTURE_TIMEOUT_SECONDS, ) response._content = prefix # pyright: ignore[reportPrivateUsage] # rebind-ok: httpx has no public setter to retain consumed bytes for auth retries - preview: Final = _preview(prefix, response.headers.get("content-type", "")) - except (TimeoutError, httpx.HTTPError): + 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 (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 @@ -510,15 +595,20 @@ def describe_upstream_response(response: httpx.Response) -> str: 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)}" - f" | response body: {_response_body_preview(response)}" + 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) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py index 513fc2f2036..43d97abe4db 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py @@ -38,11 +38,6 @@ from pydantic import BaseModel, ConfigDict, Field, SecretStr, TypeAdapter, Valid from typing_extensions import assert_never from litellm._logging import verbose_logger -from litellm.proxy._experimental.mcp_server.mcp_debug import ( - describe_upstream_http_failure, - describe_upstream_response, - safe_upstream_url, -) from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( InMemoryTokenCacheBackend, OAuthToken, @@ -107,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: diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 713330a8280..3283af26cf7 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -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" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py index 0ed78fcfc9d..dced871e146 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py @@ -290,7 +290,7 @@ def test_failure_log_omits_unstructured_body(): @pytest.mark.asyncio -@pytest.mark.parametrize("mode", ["error", "large", "timeout", "read_failure", "success", "cancel"]) +@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 @@ -302,11 +302,13 @@ async def test_error_capture_is_bounded_and_preserves_success_and_cancellation(m 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'{"error":"missing_scope","password":"first second"}' if mode != "large" else b"x" * 20000 + 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") @@ -323,7 +325,7 @@ async def test_error_capture_is_bounded_and_preserves_success_and_cancellation(m 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 = {"error": "missing_scope", "large": "capture limit", "timeout": "read failed", "read_failure": "read failed"} + 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"}' @@ -381,13 +383,12 @@ def test_failure_diagnostics_without_request_and_with_streamed_request(): def test_deep_error_body_is_bounded_without_exposing_nested_values(): - import json 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 "depth limit" in detail and "hidden-value" not in detail - assert json.loads(detail.split("response body: ")[1])["nested"] + 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']) @@ -485,3 +486,62 @@ 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