From d83a9ead002f39ae2116f931c3a1701d278c15d4 Mon Sep 17 00:00:00 2001 From: panjianxing <3207880464@qq.com> Date: Thu, 29 Oct 2026 20:39:49 +0800 Subject: [PATCH] fix(mcp): account for virtual tool calls once --- litellm/proxy/mcp_server/operations.py | 135 ++++++++++++------ .../proxy/mcp_server/test_mcp_tool_search.py | 96 +++++++++++-- 2 files changed, 177 insertions(+), 54 deletions(-) diff --git a/litellm/proxy/mcp_server/operations.py b/litellm/proxy/mcp_server/operations.py index d486a69ef1a..bee6418f6a6 100644 --- a/litellm/proxy/mcp_server/operations.py +++ b/litellm/proxy/mcp_server/operations.py @@ -321,39 +321,61 @@ async def _build_virtual_call_logging_obj( raw_headers: Mapping[str, str] | None = None, client_ip: str | None = None, ) -> LiteLLMLoggingObj | None: - """Run the pre-call pipeline (guardrails + logging setup) for a virtual - mcp_tool_call so the SSE path spend-logs like the REST path.""" - from litellm.proxy.common_request_processing import ( - ProxyBaseLLMRequestProcessing, - ) - from litellm.proxy.proxy_server import ( - general_settings, - proxy_config, - proxy_logging_obj, - ) + """Build the virtual tool's logger without running request admission twice. - request: Final = build_synthetic_mcp_request( - path="/mcp/tools/call", - raw_headers=raw_headers, - client_ip=client_ip, + The real tool call owns admission and pre-call guardrails. Running the full + proxy pre-call path here consumes a second RPM slot and parallel slot. + """ + from litellm.proxy.common_request_processing import resolve_litellm_call_id + + headers: Final = dict(raw_headers) if raw_headers else {} + call_id: Final = resolve_litellm_call_id(headers.get("x-litellm-call-id")) + trace_id: Final = get_chain_id_from_headers(headers) + request_data: Final[dict[str, object]] = { # mutable-ok: the proxy metadata helper updates this mapping + "model": f"MCP: {name}", + "call_type": CallTypes.call_mcp_tool.value, + "litellm_call_id": call_id, + "litellm_trace_id": trace_id, + "metadata": { + "spend_logs_metadata": { + "mcp_operation": "call_tool", + "mcp_tool_name": name, + }, + "headers": logging_safe_mcp_headers(raw_headers), + **({"existing_trace_id": trace_id} if trace_id else {}), + }, + "input": [ + { + "role": "system", + "content": { + "mcp_operation": "call_tool", + "mcp_tool_name": name, + "arguments": arguments, + "client_ip": client_ip, + }, + } + ], + } + auth_for_logging: Final = logging_auth or user_api_key_auth + LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( + data=request_data, + user_api_key_dict=auth_for_logging, + _metadata_variable_name="metadata", ) - pre_call_result: Final = await ProxyBaseLLMRequestProcessing( - data={"name": name, "arguments": arguments} - ).common_processing_pre_call_logic( - request=request, - user_api_key_dict=user_api_key_auth, - proxy_config=proxy_config, - route_type=CallTypes.call_mcp_tool.value, - proxy_logging_obj=proxy_logging_obj, - general_settings=general_settings, - logging_auth=logging_auth, - ) - _, logging_candidate = TypeAdapter(tuple[dict[str, object], object]).validate_python(pre_call_result) - if not isinstance(logging_candidate, LiteLLMLoggingObj): - verbose_logger.warning("MCP request processing did not return a logging object") + try: + logging_obj, _ = function_setup( + original_function=CallTypes.call_mcp_tool.value, + is_async_call=True, + rules_obj=Rules(), + start_time=datetime.now(), # noqa: DTZ005 # logging timestamps are naive + **request_data, + ) + except Exception as exc: # noqa: BLE001 # logging setup must not block the tool call + verbose_logger.debug("Failed to initialize logging for MCP virtual tool %s: %s", name, exc) return None - virtual_logging_obj: Final = logging_candidate - return virtual_logging_obj + logging_obj.call_type = CallTypes.call_mcp_tool.value + logging_obj.model = f"MCP: {name}" + return logging_obj async def _log_virtual_mcp_tool_failure( @@ -595,6 +617,7 @@ async def _dispatch_virtual_mcp_tool( tool_name=tool_request.name, mcp_servers=mcp_servers, ) + virtual_call_start: Final = datetime.now() # noqa: DTZ005 # logging timestamps are naive virtual_logging_obj: Final = await _build_virtual_call_logging_obj( name=tool_request.name, arguments=tool_request.arguments or {}, @@ -603,19 +626,45 @@ async def _dispatch_virtual_mcp_tool( raw_headers=raw_headers, client_ip=client_ip, ) - return await handle_mcp_tool_call( - tool_name=tool_request.name, - arguments=tool_request.arguments or {}, - user_api_key_dict=user_api_key_auth, - client_ip=client_ip, - mcp_servers=mcp_servers, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - litellm_logging_obj=virtual_logging_obj, - incoming_bearer_token=incoming_bearer_token, - guardrail_auth=tool_guardrail_auth, + request_data: Final = types.MappingProxyType( + { + "name": tool_request.name, + "arguments": tool_request.arguments or {}, + } + ) + try: + result: Final = await handle_mcp_tool_call( + tool_name=tool_request.name, + arguments=tool_request.arguments or {}, + user_api_key_dict=user_api_key_auth, + client_ip=client_ip, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + litellm_logging_obj=virtual_logging_obj, + incoming_bearer_token=incoming_bearer_token, + guardrail_auth=tool_guardrail_auth, + ) + except Exception as exc: + await fire_mcp_tool_call_failure_logging( + virtual_logging_obj, + exc, + virtual_call_start, + user_api_key_auth, + request_data, + ) + raise + if virtual_logging_obj is None: + return result + return await _fire_mcp_tool_call_logging( + logging_obj=virtual_logging_obj, + result=result, + start_time=virtual_call_start, + end_time=datetime.now(), # noqa: DTZ005 # logging timestamps are naive + user_api_key_auth=user_api_key_auth, + request_data=request_data, ) diff --git a/tests/unit/proxy/mcp_server/test_mcp_tool_search.py b/tests/unit/proxy/mcp_server/test_mcp_tool_search.py index b21202ba227..d05a913c445 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_tool_search.py +++ b/tests/unit/proxy/mcp_server/test_mcp_tool_search.py @@ -179,13 +179,12 @@ FAKE_VECTORS: dict[str, Vector] = { } - - def _paged_params(): from mcp.types import PaginatedRequestParams return PaginatedRequestParams() + class RecordingEmbedder: def __init__(self) -> None: self.calls: list[tuple[str, ...]] = [] @@ -1145,16 +1144,14 @@ class TestDispatchVirtualMcpTool: @pytest.mark.asyncio async def test_call_builds_and_forwards_logging_obj(self) -> None: - """Regression: the SSE dispatch must run the pre-call pipeline and forward - the resulting logging object to handle_mcp_tool_call, otherwise mcp_tool_call - over /mcp/ skips spend logging and guardrails (unlike the REST path).""" - from litellm.proxy.mcp_server import server as srv + """The virtual call forwards one logger and logs after the real tool call.""" uak = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True)) sentinel_logging_obj = object() with ( patch.object( - mcp_operations, "_build_virtual_call_logging_obj", + mcp_operations, + "_build_virtual_call_logging_obj", new_callable=AsyncMock, return_value=sentinel_logging_obj, ) as mock_build, @@ -1163,8 +1160,14 @@ class TestDispatchVirtualMcpTool: new_callable=AsyncMock, return_value="CALL_RESULT", ) as mock_call, + patch.object( + mcp_operations, + "_fire_mcp_tool_call_logging", + new_callable=AsyncMock, + return_value="LOGGED_RESULT", + ) as mock_fire, ): - await mcp_operations._dispatch_virtual_mcp_tool( + result = await mcp_operations._dispatch_virtual_mcp_tool( name=MCP_TOOL_CALL_TOOL_NAME, arguments={"tool_name": "math-add", "arguments": {"a": 1}}, user_api_key_auth=uak, @@ -1173,6 +1176,77 @@ class TestDispatchVirtualMcpTool: assert mock_build.await_count == 1 assert mock_call.await_args.kwargs["litellm_logging_obj"] is sentinel_logging_obj + assert mock_fire.await_count == 1 + assert result == "LOGGED_RESULT" + + @pytest.mark.asyncio + async def test_virtual_logger_does_not_consume_request_admission(self, monkeypatch) -> None: + from litellm.proxy.proxy_server import proxy_logging_obj + + pre_call_hook = AsyncMock() + monkeypatch.setattr(proxy_logging_obj, "pre_call_hook", pre_call_hook) + uak = UserAPIKeyAuth(api_key="k", user_id="user-1") + + logging_obj = await mcp_operations._build_virtual_call_logging_obj( + name="math-add", + arguments={"a": 1}, + user_api_key_auth=uak, + ) + + assert logging_obj is not None + assert logging_obj.model == "MCP: math-add" + pre_call_hook.assert_not_awaited() + + @pytest.mark.asyncio + async def test_virtual_logger_setup_failure_does_not_block_tool(self) -> None: + uak = UserAPIKeyAuth(api_key="k") + + with patch.object( + mcp_operations, + "function_setup", + side_effect=RuntimeError("logging unavailable"), + ): + logging_obj = await mcp_operations._build_virtual_call_logging_obj( + name="math-add", + arguments={"a": 1}, + user_api_key_auth=uak, + ) + + assert logging_obj is None + + @pytest.mark.asyncio + async def test_call_failure_logs_and_propagates(self) -> None: + uak = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True)) + sentinel_logging_obj = object() + failure = RuntimeError("tool failed") + + with ( + patch.object( + mcp_operations, + "_build_virtual_call_logging_obj", + new_callable=AsyncMock, + return_value=sentinel_logging_obj, + ), + patch( + "litellm.proxy.mcp_server.tool_search.handle_mcp_tool_call", + new_callable=AsyncMock, + side_effect=failure, + ), + patch.object( + mcp_operations, + "fire_mcp_tool_call_failure_logging", + new_callable=AsyncMock, + ) as mock_failure_log, + ): + with pytest.raises(RuntimeError, match="tool failed"): + await mcp_operations._dispatch_virtual_mcp_tool( + name=MCP_TOOL_CALL_TOOL_NAME, + arguments={"tool_name": "math-add", "arguments": {"a": 1}}, + user_api_key_auth=uak, + client_ip=None, + ) + + mock_failure_log.assert_awaited_once() @pytest.mark.asyncio async def test_search_coerces_non_int_top_k(self) -> None: @@ -1287,7 +1361,9 @@ class TestCaptureHostProgressCallback: @pytest.mark.parametrize("token", ["tok12345", 12345, 0]) @pytest.mark.parametrize("message", [None, "halfway"]) @pytest.mark.parametrize("delivery_error", [None, RuntimeError("request stream closed")]) - async def test_forwards_progress_on_originating_request(self, _mcp_request_ctx, token, message, delivery_error) -> None: + async def test_forwards_progress_on_originating_request( + self, _mcp_request_ctx, token, message, delivery_error + ) -> None: from mcp.server.connection import Connection from mcp.server.session import ServerSession from mcp.shared.dispatcher import DispatchContext @@ -1467,5 +1543,3 @@ async def test_cold_concurrent_schema_validation_accepts_valid_arguments() -> No finally: if original_affinity is not None: os.sched_setaffinity(0, original_affinity) - -