diff --git a/litellm/proxy/mcp_server/protocol_error_logging.py b/litellm/proxy/mcp_server/protocol_error_logging.py new file mode 100644 index 00000000000..026626109bc --- /dev/null +++ b/litellm/proxy/mcp_server/protocol_error_logging.py @@ -0,0 +1,309 @@ +"""Spend-log records for JSON-RPC requests the MCP SDK rejects at the protocol layer. + +The MCP SDK writes protocol errors (unknown method, malformed params) straight to +the ASGI ``send``, bypassing ``call_mcp_tool``'s logging, so such rejections +produced no ``standard_logging_object``. See +https://github.com/BerriAI/litellm/issues/28929. +""" + +import json +import traceback +import types +import uuid +from collections.abc import Mapping +from datetime import datetime +from typing import Final + +from pydantic import TypeAdapter, ValidationError +from starlette.types import Message, Send +from typing_extensions import ReadOnly, TypedDict + +from litellm._logging import verbose_logger +from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + get_chain_id_from_headers, +) +from litellm.types.utils import CallTypes +from litellm.utils import Rules, function_setup + +# Upper bound on a response body the send wrapper will parse to detect a +# JSON-RPC rejection. Rejection envelopes are tiny; anything larger is a +# permitted tool result and must not be re-deserialized after it was sent. +_MCP_PROTOCOL_ERROR_PEEK_MAX_BYTES: Final = 64 * 1024 +EMPTY_MCP_PROTOCOL_PARAMS: Final[Mapping[str, object]] = types.MappingProxyType({}) + + +class _JSONRPCRequestForLogging(TypedDict, total=False): + jsonrpc: ReadOnly[object] + method: ReadOnly[object] + params: ReadOnly[object] + id: ReadOnly[object] + + +class _JSONRPCErrorForLogging(TypedDict, total=False): + code: ReadOnly[object] + message: ReadOnly[object] + + +class _JSONRPCErrorEnvelopeForLogging(TypedDict): + error: ReadOnly[_JSONRPCErrorForLogging] + + +_JSONRPC_REQUEST_FOR_LOGGING_ADAPTER: Final[TypeAdapter[_JSONRPCRequestForLogging | tuple[object, ...]]] = TypeAdapter( + _JSONRPCRequestForLogging | tuple[object, ...] +) +_JSONRPC_PARAMS_FOR_LOGGING_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_JSONRPC_ERROR_FOR_LOGGING_ADAPTER: Final[TypeAdapter[_JSONRPCErrorEnvelopeForLogging]] = TypeAdapter( + _JSONRPCErrorEnvelopeForLogging +) + + +def _jsonrpc_rejection_reason(error_code: int | None) -> str: + """Map a JSON-RPC error code to a stable, low-cardinality reason string + for spend-log metadata. Unknown codes degrade to ``protocol_error``.""" + match error_code: + case -32700: + return "parse_error" + case -32600: + return "invalid_request" + case -32601: + return "unknown_method" + case -32602: + return "malformed_params" + case -32603: + return "internal_error" + case _: + return "protocol_error" + + +def parse_jsonrpc_request_for_logging( + body: bytes, +) -> tuple[str | None, Mapping[str, object], object]: + """Parse a peeked JSON-RPC request body into ``(method, params, id)`` for + protocol-rejection logging (#28929). + + - A single request object yields its ``method`` (only when ``jsonrpc`` + is ``"2.0"`` and ``method`` is a string), ``params`` (only if a dict), + and ``id``. + - A top-level array (JSON-RPC batch) yields the synthetic method + ``"batch"``: the MCP SDK validates the whole array as one message and, + if it can't, emits a single error response (it does NOT process + elements individually), so we record the one rejection without + attributing a tool name. + - A truncated/oversized peek or any non-request shape yields + ``(None, {}, None)`` so the caller installs no logging seam (a wrong + tool-name attribution is worse than no record). + """ + try: + request: Final = _JSONRPC_REQUEST_FOR_LOGGING_ADAPTER.validate_json(body) + except ValidationError: + return None, EMPTY_MCP_PROTOCOL_PARAMS, None + if isinstance(request, tuple): + return "batch", EMPTY_MCP_PROTOCOL_PARAMS, None + method: Final = request.get("method") + if request.get("jsonrpc") != "2.0" or not isinstance(method, str): + return None, EMPTY_MCP_PROTOCOL_PARAMS, None + raw_params: Final = request.get("params") + try: + params: Final = _JSONRPC_PARAMS_FOR_LOGGING_ADAPTER.validate_python(raw_params) + except ValidationError: + return method, EMPTY_MCP_PROTOCOL_PARAMS, request.get("id") + return method, params, request.get("id") + + +async def _log_mcp_protocol_rejection( + request_method: str | None, + params: Mapping[str, object], + jsonrpc_error: Mapping[str, object], + request_id: object, + user_api_key_auth: UserAPIKeyAuth | None, + raw_headers: Mapping[str, str] | None, + start_time: datetime, +) -> None: + """Emit a real ``standard_logging_object`` for a JSON-RPC request + rejected at the protocol layer by the MCP SDK (e.g. unknown method, + malformed params). Without this, such rejections produce no spend-log + record at all — see https://github.com/BerriAI/litellm/issues/28929. + + This builds a genuine ``LiteLLMLoggingObj`` via ``function_setup`` and + calls its ``async_failure_handler`` — the same mechanism + ``list_mcp_tools`` / ``call_mcp_tool`` use — so the + ``standard_logging_object`` is actually synthesized and dispatched to + registered callbacks. (``proxy_logging_obj.post_call_failure_hook`` is + NOT used here: it only synthesizes a ``standard_logging_object`` for + proxy-only *LLM-API* errors and requires a pre-existing + ``litellm_logging_obj`` in ``request_data`` — neither holds for a + protocol-level MCP rejection, so it would deliver a raw dict and no + record. Driving ``async_failure_handler`` directly also means a + malformed *client* request never fires the ``llm_exceptions`` alert that + lives in ``post_call_failure_hook`` -- these are 4xx client errors, not + LLM API failures.) + + Purely additive and best-effort: any failure here is swallowed so a + logging/serialization error can never alter or drop the client response. + """ + try: + if user_api_key_auth is None: + return + + raw_tool_name: Final = params.get("name") or params.get("mcp_tool_name") + tool_name: Final = str(raw_tool_name) if raw_tool_name else None + raw_error_code: Final = jsonrpc_error.get("code") + error_code: Final[int | None] = raw_error_code if isinstance(raw_error_code, int) else None + raw_error_message: Final = jsonrpc_error.get("message") + error_message: Final = ( + raw_error_message if isinstance(raw_error_message, str) else "MCP JSON-RPC protocol error" + ) + + spend_logs_metadata: Final[dict[str, object]] = { # mutable-ok: callback metadata requires a concrete dict + "mcp_operation": request_method, + "rejection_reason": _jsonrpc_rejection_reason(error_code), + "jsonrpc_error_code": error_code, + } + if tool_name is not None: + spend_logs_metadata["mcp_tool_name"] = tool_name + + request_data: Final[dict[str, object]] = { # mutable-ok: auth enrichment mutates this legacy request payload + "model": "MCP: protocol_error", + "call_type": CallTypes.call_mcp_tool.value, + "litellm_call_id": str(uuid.uuid4()), + "litellm_trace_id": get_chain_id_from_headers(dict(raw_headers) if raw_headers is not None else None), + "metadata": { + "spend_logs_metadata": spend_logs_metadata, + }, + "input": [ + { + "role": "system", + # Keep message content a string: the downstream + # standard-logging builder may run token-counting / + # serialization over messages, and a non-string content + # can raise there (silently swallowed below -> the very + # record this fix emits would be dropped). The structured + # fields live in spend_logs_metadata above. + "content": json.dumps( + { + "mcp_operation": request_method, + "jsonrpc_id": request_id, + "jsonrpc_error_code": error_code, + } + ), + } + ], + } + + LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( # pyright: ignore[reportUnknownMemberType] # legacy helper is untyped + data=request_data, + user_api_key_dict=user_api_key_auth, + _metadata_variable_name="metadata", + ) + + # Build a genuine logging object (mirrors list_mcp_tools at the + # ``function_setup`` call site) and drive its failure handler, which + # synthesizes ``model_call_details["standard_logging_object"]`` and + # dispatches it to every registered async failure callback. A + # function_setup failure is caught by the outer best-effort handler. + litellm_logging_obj, _ = function_setup( + original_function="call_mcp_tool", + rules_obj=Rules(), + start_time=start_time, + **request_data, # pyright: ignore[reportAny, reportArgumentType] # function_setup erases its forwarded keyword types + ) + + if litellm_logging_obj is None: # pyright: ignore[reportUnnecessaryComparison] # keep best-effort guard for runtime substitutes + return + + litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value + litellm_logging_obj.model = "MCP: protocol_error" + + await litellm_logging_obj.async_failure_handler( + Exception(error_message), + traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG), + start_time, + datetime.now(), # noqa: DTZ005 # matches the logging pipeline start time + ) + except Exception as log_exc: # noqa: BLE001 # best-effort: logging must never break the MCP response + verbose_logger.debug("MCP protocol-error logging failed (continuing): %s", log_exc) + + +def _parse_jsonrpc_error_response_for_logging( + raw_body: bytes, +) -> _JSONRPCErrorEnvelopeForLogging | None: + stripped_body: Final = raw_body.lstrip() + data_lines: Final = tuple( + line.removeprefix(b"data:").lstrip() for line in raw_body.splitlines() if line.startswith(b"data:") + ) + if not stripped_body.startswith(b"{") and len(data_lines) != 1: + return None + json_body: Final = stripped_body if stripped_body.startswith(b"{") else data_lines[0] + try: + return _JSONRPC_ERROR_FOR_LOGGING_ADAPTER.validate_json(json_body) + except ValidationError: + return None + + +def wrap_send_for_protocol_error_logging( + inner_send: Send, + request_method: str | None, + params: Mapping[str, object], + request_id: object, + user_api_key_auth: UserAPIKeyAuth | None, + raw_headers: Mapping[str, str] | None, + start_time: datetime, +) -> Send: + """Wrap an ASGI ``send`` so a JSON-RPC error response emitted by the MCP + SDK (an ``error`` envelope with no ``result``) produces exactly one + ``standard_logging_object``. Only JSON-RPC *requests* (carrying a + ``method``) are loggable; everything else is forwarded untouched. + + The wrapper always forwards the original message FIRST and never mutates + it, so the client response is byte-for-byte unchanged on every path. + """ + if request_method is None: + return inner_send + + already_logged: Final = {"done": False} + + async def _logging_send(message: Message) -> None: + await inner_send(message) + if already_logged["done"]: + return + try: + if message.get("type") != "http.response.body": + return + raw_body: Final[object] = message.get("body") # pyright: ignore[reportAny] # ASGI Message body is untyped + if not isinstance(raw_body, bytes) or not raw_body: + return + # Bound the work: a JSON-RPC protocol-error envelope is tiny. + # Skip the full deserialization for anything larger than the cap + # (a permitted tool result) and only parse when the bytes look + # like an error envelope. This avoids re-deserializing every + # successful response body after it has already been sent. + if len(raw_body) > _MCP_PROTOCOL_ERROR_PEEK_MAX_BYTES: + return + if b'"error"' not in raw_body or b'"result"' in raw_body: + return + payload: Final = _parse_jsonrpc_error_response_for_logging(raw_body) + if payload is None: + return + already_logged["done"] = True + await _log_mcp_protocol_rejection( + request_method=request_method, + params=params, + jsonrpc_error=payload["error"], + request_id=request_id, + user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + start_time=start_time, + ) + except (TypeError, ValueError): + # Non-JSON / streaming chunk — nothing to log here. + return + except Exception as wrap_exc: # noqa: BLE001 # best-effort: never break the response on logging + verbose_logger.debug( + "MCP protocol-error send wrapper failed (continuing): %s", + wrap_exc, + ) + + return _logging_send diff --git a/litellm/proxy/mcp_server/server.py b/litellm/proxy/mcp_server/server.py index bbbbd09e527..ccb91e55238 100644 --- a/litellm/proxy/mcp_server/server.py +++ b/litellm/proxy/mcp_server/server.py @@ -80,6 +80,11 @@ from litellm.proxy.mcp_server.oauth_utils import ( redact_mcp_resource_url, well_known_root_suffix, ) +from litellm.proxy.mcp_server.protocol_error_logging import ( + EMPTY_MCP_PROTOCOL_PARAMS, + parse_jsonrpc_request_for_logging, + wrap_send_for_protocol_error_logging, +) from litellm.proxy.mcp_server.scope_challenge import ( SCOPE_RESPONSE_KEY, OAuthScopeResponse, @@ -2753,6 +2758,19 @@ if MCP_AVAILABLE: ) -> None: _increment_active_request_session(initialized_session_id) + # Parse the (already-peeked) JSON-RPC request so we can emit a + # standard_logging_object if the MCP SDK rejects it at the protocol + # layer (unknown method / malformed params). The SDK writes that + # error directly to ``send``, bypassing call_mcp_tool's logging, so + # we wrap ``send`` below. See issue #28929. + protocol_log_context: Final = ( + parse_jsonrpc_request_for_logging(body) + if body and request_method == "POST" and not is_jsonrpc_response + else (None, EMPTY_MCP_PROTOCOL_PARAMS, None) + ) + _protocol_log_method, _protocol_log_params, _protocol_log_id = protocol_log_context + _protocol_log_start_time: Final = datetime.now() # noqa: DTZ005 # logging pipeline uses naive datetimes + async def _dispatch() -> None: _otel_publish_transport_span_on_scope(scope) _otel_publish_request_destinations_on_scope(scope) @@ -2769,6 +2787,18 @@ if MCP_AVAILABLE: copy_existing_session_auth_context=is_initialize or is_cancellation, ) local_send = scope_response or send + # Additively observe protocol-level JSON-RPC rejections (#28929) + # without touching the success path. Applied first so the + # stateful-session wrapper (below) still wraps the outermost send. + local_send = wrap_send_for_protocol_error_logging( # rebind-ok: layers the logging wrapper onto the per-request send + local_send, + request_method=_protocol_log_method, + params=_protocol_log_params, + request_id=_protocol_log_id, + user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + start_time=_protocol_log_start_time, + ) if use_stateful and is_initialize: local_send = _wrap_send_with_stateful_session_auth_context( local_send, diff --git a/tests/unit/proxy/mcp_server/test_mcp_logging.py b/tests/unit/proxy/mcp_server/test_mcp_logging.py index b12cb916ecb..b269eb53603 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_logging.py +++ b/tests/unit/proxy/mcp_server/test_mcp_logging.py @@ -186,6 +186,7 @@ async def test_sdk_rejection_sync_logging_allows_concurrent_ping( loop: Final = asyncio.get_running_loop() ping_responses: Final = SimpleQueue[EmptyResult]() async with Client(server, mode="legacy") as client: + def sync_failure( kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime ) -> None: @@ -326,6 +327,7 @@ def _call_tool_params(name, arguments=None): return CallToolRequestParams(name=name, arguments=arguments) + class TestMCPLogger(CustomLogger): def __init__(self): self.standard_logging_payload = None @@ -363,9 +365,7 @@ def _set_authorized_user(server_ids): async def test_mcp_cost_tracking(): # Create a mock tool call result litellm.logging_callback_manager._reset_all_callbacks() - mock_result = CallToolResult( - content=[TextContent(type="text", text="Test response")], isError=False - ) + mock_result = CallToolResult(content=[TextContent(type="text", text="Test response")], isError=False) # Create a mock MCPClient mock_client = AsyncMock() @@ -426,7 +426,6 @@ async def test_mcp_cost_tracking(): local_mcp_server_manager, ), ): - _set_authorized_user(local_mcp_server_manager.get_all_mcp_server_ids()) print( @@ -435,12 +434,10 @@ async def test_mcp_cost_tracking(): ) # Manually add the tool mapping to ensure it's available (since mocking might not capture it properly) - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[ - "add_tools" - ] = "zapier_gmail_server" - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[ - "zapier_gmail_server-add_tools" - ] = "zapier_gmail_server" + local_mcp_server_manager.tool_name_to_mcp_server_name_mapping["add_tools"] = "zapier_gmail_server" + local_mcp_server_manager.tool_name_to_mcp_server_name_mapping["zapier_gmail_server-add_tools"] = ( + "zapier_gmail_server" + ) # Call mcp tool response = await mcp_server_tool_call( @@ -452,9 +449,7 @@ async def test_mcp_cost_tracking(): await asyncio.sleep(2) logged_standard_logging_payload = ( - test_logger.mcp_tool_call_payloads[-1] - if test_logger.mcp_tool_call_payloads - else None + test_logger.mcp_tool_call_payloads[-1] if test_logger.mcp_tool_call_payloads else None ) print("logged_standard_logging_payload", logged_standard_logging_payload) @@ -464,9 +459,7 @@ async def test_mcp_cost_tracking(): if isinstance(response, CallToolResult): response_list = response.content else: - response_list = list( - response - ) # Convert iterable to list for backward compatibility + response_list = list(response) # Convert iterable to list for backward compatibility assert len(response_list) == 1 assert isinstance(response_list[0], TextContent) assert response_list[0].text == "Test response" @@ -477,9 +470,7 @@ async def test_mcp_cost_tracking(): ###### # verify response cost is 1.2 as set on default_cost_per_query # Critical - the cost is tracked as $1.2 - assert ( - logged_standard_logging_payload is not None - ), "Standard logging payload should not be None" + assert logged_standard_logging_payload is not None, "Standard logging payload should not be None" assert logged_standard_logging_payload["response_cost"] == 1.2 @@ -488,9 +479,7 @@ async def test_mcp_cost_tracking_per_tool(): """Test that individual tool costs are tracked correctly when tool_name_to_cost_per_query is configured""" # Create a mock tool call result litellm.logging_callback_manager._reset_all_callbacks() - mock_result = CallToolResult( - content=[TextContent(type="text", text="Test response")], isError=False - ) + mock_result = CallToolResult(content=[TextContent(type="text", text="Test response")], isError=False) # Create a mock MCPClient mock_client = AsyncMock() @@ -553,18 +542,10 @@ async def test_mcp_cost_tracking_per_tool(): await local_mcp_server_manager._initialize_tool_name_to_mcp_server_name_mapping() # Manually add the tool mapping to ensure it's available (since mocking might not capture it properly) - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[ - "expensive_tool" - ] = "test_server" - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[ - "test_server-expensive_tool" - ] = "test_server" - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping["cheap_tool"] = ( - "test_server" - ) - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[ - "test_server-cheap_tool" - ] = "test_server" + local_mcp_server_manager.tool_name_to_mcp_server_name_mapping["expensive_tool"] = "test_server" + local_mcp_server_manager.tool_name_to_mcp_server_name_mapping["test_server-expensive_tool"] = "test_server" + local_mcp_server_manager.tool_name_to_mcp_server_name_mapping["cheap_tool"] = "test_server" + local_mcp_server_manager.tool_name_to_mcp_server_name_mapping["test_server-cheap_tool"] = "test_server" # Patch the global manager in both modules where it's used with ( @@ -577,7 +558,6 @@ async def test_mcp_cost_tracking_per_tool(): local_mcp_server_manager, ), ): - _set_authorized_user(local_mcp_server_manager.get_all_mcp_server_ids()) print( @@ -595,18 +575,12 @@ async def test_mcp_cost_tracking_per_tool(): await asyncio.sleep(2) logged_standard_logging_payload_1 = ( - test_logger.mcp_tool_call_payloads[-1] - if test_logger.mcp_tool_call_payloads - else None - ) - print( - "logged_standard_logging_payload_1", logged_standard_logging_payload_1 + test_logger.mcp_tool_call_payloads[-1] if test_logger.mcp_tool_call_payloads else None ) + print("logged_standard_logging_payload_1", logged_standard_logging_payload_1) # Verify expensive tool cost - assert ( - logged_standard_logging_payload_1 is not None - ), "Standard logging payload 1 should not be None" + assert logged_standard_logging_payload_1 is not None, "Standard logging payload 1 should not be None" assert logged_standard_logging_payload_1["response_cost"] == 5.0 # Reset logger for second test @@ -623,18 +597,12 @@ async def test_mcp_cost_tracking_per_tool(): await asyncio.sleep(2) logged_standard_logging_payload_2 = ( - test_logger.mcp_tool_call_payloads[-1] - if test_logger.mcp_tool_call_payloads - else None - ) - print( - "logged_standard_logging_payload_2", logged_standard_logging_payload_2 + test_logger.mcp_tool_call_payloads[-1] if test_logger.mcp_tool_call_payloads else None ) + print("logged_standard_logging_payload_2", logged_standard_logging_payload_2) # Verify cheap tool cost - assert ( - logged_standard_logging_payload_2 is not None - ), "Standard logging payload 2 should not be None" + assert logged_standard_logging_payload_2 is not None, "Standard logging payload 2 should not be None" assert logged_standard_logging_payload_2["response_cost"] == 0.1 # Add basic response assertions @@ -669,9 +637,7 @@ class MCPLoggerHook(TestMCPLogger): async def test_mcp_tool_call_hook(): # Create a mock tool call result litellm.logging_callback_manager._reset_all_callbacks() - mock_result = CallToolResult( - content=[TextContent(type="text", text="Test response")], isError=False - ) + mock_result = CallToolResult(content=[TextContent(type="text", text="Test response")], isError=False) # Create a mock MCPClient mock_client = AsyncMock() @@ -717,12 +683,10 @@ async def test_mcp_tool_call_hook(): await local_mcp_server_manager._initialize_tool_name_to_mcp_server_name_mapping() # Manually add the tool mapping to ensure it's available (since mocking might not capture it properly) - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping["add_tools"] = ( + local_mcp_server_manager.tool_name_to_mcp_server_name_mapping["add_tools"] = "zapier_gmail_server" + local_mcp_server_manager.tool_name_to_mcp_server_name_mapping["zapier_gmail_server-add_tools"] = ( "zapier_gmail_server" ) - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[ - "zapier_gmail_server-add_tools" - ] = "zapier_gmail_server" # Patch the global manager in both modules where it's used with ( @@ -735,7 +699,6 @@ async def test_mcp_tool_call_hook(): local_mcp_server_manager, ), ): - _set_authorized_user(local_mcp_server_manager.get_all_mcp_server_ids()) print( @@ -754,18 +717,666 @@ async def test_mcp_tool_call_hook(): # check logged standard logging payload logged_standard_logging_payload = ( - test_logger.mcp_tool_call_payloads[-1] - if test_logger.mcp_tool_call_payloads - else None + test_logger.mcp_tool_call_payloads[-1] if test_logger.mcp_tool_call_payloads else None ) print("logged_standard_logging_payload", logged_standard_logging_payload) - assert ( - logged_standard_logging_payload is not None - ), "Standard logging payload should not be None" + assert logged_standard_logging_payload is not None, "Standard logging payload should not be None" assert logged_standard_logging_payload["response_cost"] == 1.42 -_QUEUED_LOGGING_OUTLIVES_TEST = ''' +# --------------------------------------------------------------------------- +# Tests for #28929: standard_logging_object on /mcp/ JSON-RPC protocol-level +# rejections (unknown method / malformed params). These exercise the additive +# send-wrapper + helper added to handle_streamable_http_mcp's dispatch path. +# +# The load-bearing assertion is that a REAL standard_logging_object is emitted +# to a registered failure callback (not merely that an internal hook was +# called) — the rejection path builds a genuine LiteLLMLoggingObj and drives +# its async_failure_handler, mirroring list_mcp_tools / call_mcp_tool. +# --------------------------------------------------------------------------- +from datetime import datetime + +from litellm.proxy.mcp_server.protocol_error_logging import ( + _jsonrpc_rejection_reason, + _log_mcp_protocol_rejection, + _parse_jsonrpc_error_response_for_logging, +) +from litellm.proxy.mcp_server.protocol_error_logging import ( + parse_jsonrpc_request_for_logging as _parse_jsonrpc_request_for_logging, +) +from litellm.proxy.mcp_server.protocol_error_logging import ( + wrap_send_for_protocol_error_logging as _wrap_send_for_protocol_error_logging, +) +from litellm.types.utils import CallTypes + + +class _ProtocolRejectionLogger(CustomLogger): + """Captures the standard_logging_object dispatched on the failure path.""" + + def __init__(self): + self.failure_payloads = [] + super().__init__() + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + slo = kwargs.get("standard_logging_object", None) + self.failure_payloads.append(slo) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + # Should never fire for a protocol rejection. + self.failure_payloads.append(("UNEXPECTED_SUCCESS", kwargs.get("standard_logging_object"))) + + +def _set_protocol_test_callbacks(callbacks): + litellm.callbacks = callbacks # test-quality-ok: integration tests require LiteLLM's public callback seam and every caller restores it + + +def _jsonrpc_error_body(code, message="boom", _id=1): + return json.dumps({"jsonrpc": "2.0", "id": _id, "error": {"code": code, "message": message}}).encode("utf-8") + + +def _jsonrpc_result_body(_id=1): + return json.dumps({"jsonrpc": "2.0", "id": _id, "result": {"ok": True}}).encode("utf-8") + + +def test_mcp_jsonrpc_rejection_reason_mapping(): + assert _jsonrpc_rejection_reason(-32601) == "unknown_method" + assert _jsonrpc_rejection_reason(-32602) == "malformed_params" + assert _jsonrpc_rejection_reason(-32700) == "parse_error" + assert _jsonrpc_rejection_reason(-32600) == "invalid_request" + # Unknown / None codes degrade to a generic, low-cardinality reason. + assert _jsonrpc_rejection_reason(12345) == "protocol_error" + assert _jsonrpc_rejection_reason(None) == "protocol_error" + + +@pytest.mark.asyncio +async def test_mcp_protocol_rejection_emits_standard_logging_object(): + """A JSON-RPC error response (no result) must cause a REAL + standard_logging_object to be dispatched to a registered failure callback, + carrying the rejection metadata (#28929).""" + litellm.logging_callback_manager._reset_all_callbacks() + capture = _ProtocolRejectionLogger() + _set_protocol_test_callbacks([capture]) + + user_auth = UserAPIKeyAuth(api_key="test", user_id="u1") + sent_messages = [] + + async def fake_send(message): + sent_messages.append(message) + + try: + wrapped = _wrap_send_for_protocol_error_logging( + fake_send, + request_method="tools/call", + params={"name": "totally_made_up", "arguments": {}}, + request_id=7, + user_api_key_auth=user_auth, + raw_headers={}, + start_time=datetime.now(), + ) + body = _jsonrpc_error_body(-32602, "Invalid request parameters", _id=7) + await wrapped({"type": "http.response.start", "status": 200, "headers": []}) + await wrapped({"type": "http.response.body", "body": body}) + # async_failure_handler may dispatch on the loop; give it a tick. + await asyncio.sleep(1) + finally: + _set_protocol_test_callbacks([]) + + # Response forwarded byte-for-byte and unchanged. + assert any(m.get("type") == "http.response.body" and m.get("body") == body for m in sent_messages), ( + "original response bytes must be forwarded unchanged" + ) + + # A genuine, non-None standard_logging_object reached the callback. + non_none = [p for p in capture.failure_payloads if isinstance(p, dict)] + assert non_none, ( + f"expected a non-None standard_logging_object on the failure callback, got: {capture.failure_payloads!r}" + ) + slo = non_none[0] + assert slo.get("call_type") == CallTypes.call_mcp_tool.value + # status is a failure-bearing value (StandardLoggingPayloadStatus). + assert slo.get("status") == "failure" + # Rejection metadata threaded through spend_logs_metadata. + spend_meta = (slo.get("metadata") or {}).get("spend_logs_metadata") or {} + assert spend_meta.get("mcp_operation") == "tools/call" + assert spend_meta.get("rejection_reason") == "malformed_params" + assert spend_meta.get("jsonrpc_error_code") == -32602 + assert spend_meta.get("mcp_tool_name") == "totally_made_up" + + +@pytest.mark.asyncio +async def test_mcp_protocol_unknown_method_emits_slo(): + """An unknown-method (-32601) rejection emits an SLO with the right + reason and no tool name (empty params).""" + litellm.logging_callback_manager._reset_all_callbacks() + capture = _ProtocolRejectionLogger() + _set_protocol_test_callbacks([capture]) + + user_auth = UserAPIKeyAuth(api_key="test", user_id="u1") + + async def fake_send(message): + pass + + try: + wrapped = _wrap_send_for_protocol_error_logging( + fake_send, + request_method="tools/totally_made_up", + params={}, + request_id=1, + user_api_key_auth=user_auth, + raw_headers={}, + start_time=datetime.now(), + ) + await wrapped( + { + "type": "http.response.body", + "body": _jsonrpc_error_body(-32601, "Method not found"), + } + ) + await asyncio.sleep(1) + finally: + _set_protocol_test_callbacks([]) + + non_none = [p for p in capture.failure_payloads if isinstance(p, dict)] + assert non_none, f"expected an SLO, got: {capture.failure_payloads!r}" + spend_meta = (non_none[0].get("metadata") or {}).get("spend_logs_metadata") or {} + assert spend_meta.get("rejection_reason") == "unknown_method" + assert "mcp_tool_name" not in spend_meta + + +@pytest.mark.asyncio +async def test_mcp_protocol_batch_rejection_emits_slo(): + """A JSON-RPC batch (array body) the MCP SDK rejects as one error must + still produce one SLO, attributed to mcp_operation='batch' (#28929, + Codex batch finding).""" + litellm.logging_callback_manager._reset_all_callbacks() + capture = _ProtocolRejectionLogger() + _set_protocol_test_callbacks([capture]) + + user_auth = UserAPIKeyAuth(api_key="test", user_id="u1") + + async def fake_send(message): + pass + + try: + # Mirrors how handle_streamable_http_mcp maps a top-level array to a + # synthetic 'batch' operation before wrapping send. + wrapped = _wrap_send_for_protocol_error_logging( + fake_send, + request_method="batch", + params={}, + request_id=None, + user_api_key_auth=user_auth, + raw_headers={}, + start_time=datetime.now(), + ) + await wrapped( + { + "type": "http.response.body", + "body": _jsonrpc_error_body(-32602, "Validation error", _id="server-error"), + } + ) + await asyncio.sleep(1) + finally: + _set_protocol_test_callbacks([]) + + non_none = [p for p in capture.failure_payloads if isinstance(p, dict)] + assert non_none, f"expected a batch SLO, got: {capture.failure_payloads!r}" + spend_meta = (non_none[0].get("metadata") or {}).get("spend_logs_metadata") or {} + assert spend_meta.get("mcp_operation") == "batch" + assert spend_meta.get("rejection_reason") == "malformed_params" + + +@pytest.mark.asyncio +async def test_mcp_protocol_success_path_emits_no_failure_record(): + """A JSON-RPC success response (has result) must NOT emit a failure + record, and the response must still be forwarded unchanged.""" + litellm.logging_callback_manager._reset_all_callbacks() + capture = _ProtocolRejectionLogger() + _set_protocol_test_callbacks([capture]) + + user_auth = UserAPIKeyAuth(api_key="test", user_id="u1") + sent_messages = [] + + async def fake_send(message): + sent_messages.append(message) + + try: + wrapped = _wrap_send_for_protocol_error_logging( + fake_send, + request_method="tools/call", + params={"name": "add", "arguments": {}}, + request_id=1, + user_api_key_auth=user_auth, + raw_headers={}, + start_time=datetime.now(), + ) + body = _jsonrpc_result_body(_id=1) + await wrapped({"type": "http.response.body", "body": body}) + await asyncio.sleep(1) + finally: + _set_protocol_test_callbacks([]) + + assert capture.failure_payloads == [], ( + f"success path must not dispatch any failure record, got: {capture.failure_payloads!r}" + ) + assert any(m.get("body") == body for m in sent_messages) + + +@pytest.mark.asyncio +async def test_mcp_protocol_wrapper_noop_when_not_a_request(): + """When the POST carries no JSON-RPC method (request_method is None) the + wrapper returns the original send untouched (no logging seam at all).""" + + async def fake_send(message): + pass + + returned = _wrap_send_for_protocol_error_logging( + fake_send, + request_method=None, + params={}, + request_id=None, + user_api_key_auth=UserAPIKeyAuth(api_key="test", user_id="u1"), + raw_headers={}, + start_time=datetime.now(), + ) + assert returned is fake_send, "non-request POSTs must not be wrapped" + + +@pytest.mark.asyncio +async def test_mcp_protocol_logging_failure_does_not_break_response(): + """If logging-object construction raises, the wrapper must swallow it and + still forward the client response (observability must never drop a + response).""" + litellm.logging_callback_manager._reset_all_callbacks() + _set_protocol_test_callbacks([]) + + user_auth = UserAPIKeyAuth(api_key="test", user_id="u1") + sent_messages = [] + + async def fake_send(message): + sent_messages.append(message) + + try: + with patch( # test-quality-ok: fault injection verifies logging failure cannot interrupt the already-forwarded MCP response + "litellm.proxy.mcp_server.protocol_error_logging.function_setup", + side_effect=RuntimeError("logging init down"), + ): + wrapped = _wrap_send_for_protocol_error_logging( + fake_send, + request_method="tools/call", + params={"name": "add"}, + request_id=1, + user_api_key_auth=user_auth, + raw_headers={}, + start_time=datetime.now(), + ) + body = _jsonrpc_error_body(-32602) + # Must not raise. + await wrapped({"type": "http.response.body", "body": body}) + finally: + _set_protocol_test_callbacks([]) + + assert any(m.get("body") == body for m in sent_messages), "response must be forwarded even when logging init raises" + + +@pytest.mark.asyncio +async def test_mcp_log_protocol_rejection_no_auth_is_noop(): + """The helper is a no-op (no exception, no SLO) when there is no user + auth.""" + litellm.logging_callback_manager._reset_all_callbacks() + capture = _ProtocolRejectionLogger() + _set_protocol_test_callbacks([capture]) + try: + await _log_mcp_protocol_rejection( + request_method="tools/call", + params={"name": "x"}, + jsonrpc_error={"code": -32602, "message": "bad"}, + request_id=1, + user_api_key_auth=None, + raw_headers={}, + start_time=datetime.now(), + ) + await asyncio.sleep(0.2) + finally: + _set_protocol_test_callbacks([]) + assert capture.failure_payloads == [] + + +@pytest.mark.asyncio +async def test_mcp_protocol_wrapper_skips_oversized_body(): + """V2 guard: a response body larger than the parse cap is never + deserialized (a permitted tool result), so no SLO is emitted even if it + happens to contain an error-shaped substring.""" + litellm.logging_callback_manager._reset_all_callbacks() + capture = _ProtocolRejectionLogger() + _set_protocol_test_callbacks([capture]) + + sent = [] + + async def fake_send(message): + sent.append(message) + + # Build a body > 64 KiB that still contains the error marker bytes. + big = b'{"result": {"x": "' + (b"a" * (70 * 1024)) + b'"}, "error": 1}' + try: + wrapped = _wrap_send_for_protocol_error_logging( + fake_send, + request_method="tools/call", + params={"name": "add"}, + request_id=1, + user_api_key_auth=UserAPIKeyAuth(api_key="test", user_id="u1"), + raw_headers={}, + start_time=datetime.now(), + ) + await wrapped({"type": "http.response.body", "body": big}) + await asyncio.sleep(0.3) + finally: + _set_protocol_test_callbacks([]) + + assert capture.failure_payloads == [], "oversized body must not be parsed/logged" + assert any(m.get("body") == big for m in sent) + + +@pytest.mark.asyncio +async def test_mcp_protocol_wrapper_skips_non_error_body(): + """V2 guard: the cheap byte pre-check skips bodies that don't look like an + error envelope (no '"error"', or a '"result"' present) without json.loads.""" + litellm.logging_callback_manager._reset_all_callbacks() + capture = _ProtocolRejectionLogger() + _set_protocol_test_callbacks([capture]) + + async def fake_send(message): + pass + + try: + wrapped = _wrap_send_for_protocol_error_logging( + fake_send, + request_method="tools/call", + params={"name": "add"}, + request_id=1, + user_api_key_auth=UserAPIKeyAuth(api_key="test", user_id="u1"), + raw_headers={}, + start_time=datetime.now(), + ) + # No "error" marker at all. + await wrapped({"type": "http.response.body", "body": b'{"jsonrpc":"2.0","id":1}'}) + # Both "error" and "result" present -> not a pure rejection. + await wrapped( + { + "type": "http.response.body", + "body": b'{"jsonrpc":"2.0","id":1,"result":{},"error":null}', + } + ) + await asyncio.sleep(0.3) + finally: + _set_protocol_test_callbacks([]) + + assert capture.failure_payloads == [], "non-error bodies must not be logged" + + +@pytest.mark.asyncio +async def test_mcp_protocol_wrapper_ignores_non_body_messages_and_empty(): + """The wrapper forwards http.response.start and empty bodies without + attempting to log.""" + litellm.logging_callback_manager._reset_all_callbacks() + capture = _ProtocolRejectionLogger() + _set_protocol_test_callbacks([capture]) + + sent = [] + + async def fake_send(message): + sent.append(message) + + try: + wrapped = _wrap_send_for_protocol_error_logging( + fake_send, + request_method="tools/call", + params={"name": "add"}, + request_id=1, + user_api_key_auth=UserAPIKeyAuth(api_key="test", user_id="u1"), + raw_headers={}, + start_time=datetime.now(), + ) + await wrapped({"type": "http.response.start", "status": 200, "headers": []}) + await wrapped({"type": "http.response.body", "body": b""}) + await asyncio.sleep(0.2) + finally: + _set_protocol_test_callbacks([]) + + assert capture.failure_payloads == [] + assert len(sent) == 2 # both forwarded + + +@pytest.mark.asyncio +async def test_mcp_protocol_wrapper_handles_malformed_error_body(): + """An error-looking but invalid-JSON body trips the JSONDecodeError branch + and is swallowed (no SLO, response still forwarded).""" + litellm.logging_callback_manager._reset_all_callbacks() + capture = _ProtocolRejectionLogger() + _set_protocol_test_callbacks([capture]) + + sent = [] + + async def fake_send(message): + sent.append(message) + + try: + wrapped = _wrap_send_for_protocol_error_logging( + fake_send, + request_method="tools/call", + params={"name": "add"}, + request_id=1, + user_api_key_auth=UserAPIKeyAuth(api_key="test", user_id="u1"), + raw_headers={}, + start_time=datetime.now(), + ) + # passes the byte pre-check ('"error"' present, no '"result"') but is + # not valid JSON. + body = b'{"error": this-is-not-json' + await wrapped({"type": "http.response.body", "body": body}) + await asyncio.sleep(0.2) + finally: + _set_protocol_test_callbacks([]) + + assert capture.failure_payloads == [] + assert any(m.get("body") == body for m in sent) + + +@pytest.mark.asyncio +async def test_mcp_log_protocol_rejection_logging_obj_none_is_noop(monkeypatch): + """If function_setup yields no logging object, the helper returns without + raising (covers the None-guard branch).""" + litellm.logging_callback_manager._reset_all_callbacks() + capture = _ProtocolRejectionLogger() + _set_protocol_test_callbacks([capture]) + + import litellm.proxy.mcp_server.protocol_error_logging as mcp_server_mod + + def _fake_function_setup(*args, **kwargs): + return None, {} + + monkeypatch.setattr(mcp_server_mod, "function_setup", _fake_function_setup) + try: + await _log_mcp_protocol_rejection( + request_method="tools/call", + params={"name": "x"}, + jsonrpc_error={"code": -32602, "message": "bad"}, + request_id=1, + user_api_key_auth=UserAPIKeyAuth(api_key="test", user_id="u1"), + raw_headers={}, + start_time=datetime.now(), + ) + await asyncio.sleep(0.2) + finally: + _set_protocol_test_callbacks([]) + + assert capture.failure_payloads == [] + + +@pytest.mark.asyncio +async def test_mcp_protocol_wrapper_logs_only_once_per_request(): + """already_logged short-circuit: a second error body on the same request + does not produce a second SLO (covers the early-return guard).""" + litellm.logging_callback_manager._reset_all_callbacks() + capture = _ProtocolRejectionLogger() + _set_protocol_test_callbacks([capture]) + + async def fake_send(message): + pass + + try: + wrapped = _wrap_send_for_protocol_error_logging( + fake_send, + request_method="tools/call", + params={"name": "add"}, + request_id=1, + user_api_key_auth=UserAPIKeyAuth(api_key="test", user_id="u1"), + raw_headers={}, + start_time=datetime.now(), + ) + await wrapped({"type": "http.response.body", "body": _jsonrpc_error_body(-32602)}) + await wrapped({"type": "http.response.body", "body": _jsonrpc_error_body(-32601)}) + await asyncio.sleep(1) + finally: + _set_protocol_test_callbacks([]) + + non_none = [p for p in capture.failure_payloads if isinstance(p, dict)] + assert len(non_none) == 1, f"exactly one SLO expected, got {len(non_none)}" + + +@pytest.mark.asyncio +async def test_mcp_protocol_wrapper_swallows_unexpected_error(monkeypatch): + """A non-JSON unexpected error raised inside the detection block is caught + by the broad guard and never breaks the response (covers the BLE001 + branch).""" + litellm.logging_callback_manager._reset_all_callbacks() + _set_protocol_test_callbacks([]) + + import litellm.proxy.mcp_server.protocol_error_logging as mcp_server_mod + + async def _boom(**kwargs): + raise RuntimeError("unexpected") + + monkeypatch.setattr(mcp_server_mod, "_log_mcp_protocol_rejection", _boom) + + sent = [] + + async def fake_send(message): + sent.append(message) + + wrapped = _wrap_send_for_protocol_error_logging( + fake_send, + request_method="tools/call", + params={"name": "add"}, + request_id=1, + user_api_key_auth=UserAPIKeyAuth(api_key="test", user_id="u1"), + raw_headers={}, + start_time=datetime.now(), + ) + body = _jsonrpc_error_body(-32602) + # Must not raise despite _log_mcp_protocol_rejection blowing up. + await wrapped({"type": "http.response.body", "body": body}) + assert any(m.get("body") == body for m in sent) + + +def test_parse_jsonrpc_request_for_logging_single(): + method, params, rid = _parse_jsonrpc_request_for_logging( + b'{"jsonrpc":"2.0","method":"tools/call","id":7,"params":{"name":"x"}}' + ) + assert method == "tools/call" + assert params == {"name": "x"} + assert rid == 7 + + +def test_parse_jsonrpc_request_for_logging_non_string_method_skipped(): + # method present but not a string -> not loggable as a request. + method, params, rid = _parse_jsonrpc_request_for_logging(b'{"jsonrpc":"2.0","method":123,"id":1}') + assert method is None + assert params == {} + assert rid is None + + +def test_parse_jsonrpc_request_for_logging_non_dict_params(): + # params present but not a dict -> params stays empty, method still parsed. + method, params, rid = _parse_jsonrpc_request_for_logging( + b'{"jsonrpc":"2.0","method":"tools/call","id":1,"params":[1,2]}' + ) + assert method == "tools/call" + assert params == {} + + +def test_parse_jsonrpc_request_for_logging_missing_jsonrpc_version(): + method, params, rid = _parse_jsonrpc_request_for_logging(b'{"method":"tools/call","id":1}') + assert method is None + + +def test_parse_jsonrpc_request_for_logging_batch(): + method, params, rid = _parse_jsonrpc_request_for_logging(b'[{"jsonrpc":"2.0","method":"tools/call","id":1}]') + assert method == "batch" + assert params == {} + assert rid is None + + +def test_parse_jsonrpc_request_for_logging_truncated_or_non_json(): + # Truncated/oversized peek or plain non-JSON -> no attribution. + assert _parse_jsonrpc_request_for_logging(b'{"jsonrpc":"2.0","meth') == ( + None, + {}, + None, + ) + assert _parse_jsonrpc_request_for_logging(b"not json at all") == (None, {}, None) + # A bare JSON scalar (neither dict nor list). + assert _parse_jsonrpc_request_for_logging(b'"a string"') == (None, {}, None) + + +def test_parse_jsonrpc_error_response_for_logging_accepts_sse(): + payload = _parse_jsonrpc_error_response_for_logging( + b'event: message\r\ndata: {"jsonrpc":"2.0","id":1,"error":{"code":-32602,"message":"bad"}}\r\n\r\n' + ) + assert payload is not None + assert payload["error"]["code"] == -32602 + + +@pytest.mark.asyncio +async def test_mcp_protocol_sse_rejection_emits_slo(): + litellm.logging_callback_manager._reset_all_callbacks() + capture = _ProtocolRejectionLogger() + _set_protocol_test_callbacks([capture]) + + async def fake_send(message): + pass + + try: + wrapped = _wrap_send_for_protocol_error_logging( + fake_send, + request_method="tools/not-real", + params={}, + request_id=1, + user_api_key_auth=UserAPIKeyAuth(api_key="test", user_id="u1"), + raw_headers={}, + start_time=datetime.now(), + ) + await wrapped( + { + "type": "http.response.body", + "body": ( + b'event: message\r\ndata: {"jsonrpc":"2.0","id":1,' + b'"error":{"code":-32602,"message":"Invalid request parameters"}}\r\n\r\n' + ), + } + ) + await asyncio.sleep(1) + finally: + _set_protocol_test_callbacks([]) + + payloads = [payload for payload in capture.failure_payloads if isinstance(payload, dict)] + assert len(payloads) == 1 + spend_meta = (payloads[0].get("metadata") or {}).get("spend_logs_metadata") or {} + assert spend_meta.get("mcp_operation") == "tools/not-real" + assert spend_meta.get("rejection_reason") == "malformed_params" + + +_QUEUED_LOGGING_OUTLIVES_TEST = """ import time from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -788,7 +1399,7 @@ async def test_2_starts_after_the_previous_tests_logging_ran(): GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(_record_run()) await GLOBAL_LOGGING_WORKER.flush() assert [t < started_at for t in ran_at] == [True, False] -''' +""" def test_logging_queued_by_one_test_is_drained_before_the_next(tmp_path: Path): @@ -805,3 +1416,12 @@ def test_logging_queued_by_one_test_is_drained_before_the_next(tmp_path: Path): timeout=120, ) assert result.returncode == 0, result.stdout + result.stderr + + +def test_jsonrpc_rejection_reason_maps_internal_error(): + assert _jsonrpc_rejection_reason(-32603) == "internal_error" + + +def test_parse_jsonrpc_error_response_rejects_multiple_sse_data_lines(): + body = b'data: {"error":{"code":-32602}}\ndata: {"error":{"code":-32601}}\n' + assert _parse_jsonrpc_error_response_for_logging(body) is None