diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 4e3b92edc89d..95654629eda3 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -881,6 +881,7 @@ async def call_tool( host_progress_callback: Callable | None = None, raise_on_error: bool = False, allow_input_required: bool = False, + on_dispatch: Callable[[], None] | None = None, ) -> MCPCallToolResult | InputRequiredResult: """ Call an MCP Tool. @@ -893,6 +894,8 @@ async def call_tool( allow_input_required: When True, a 2026-07-28 upstream may answer with an interim ``InputRequiredResult`` and it is returned as is. The SDK rejects it otherwise, so callers only opt in when the downstream side can carry it. + on_dispatch: Called once the session is ready, right before the tool call is sent, + so a caller can tell a call the server may have run from one that never left. """ verbose_logger.info("MCP client calling tool '%s'", call_tool_request_params.name) @@ -911,6 +914,8 @@ async def on_progress(progress: float, total: float | None, message: str | None) async def _call_tool_operation(session: ClientSession): verbose_logger.debug("MCP client sending tool call to session") + if on_dispatch is not None: + on_dispatch() return await session.call_tool( name=call_tool_request_params.name, arguments=call_tool_request_params.arguments, diff --git a/litellm/llms/anthropic/pass_through/messages/mcp_handler.py b/litellm/llms/anthropic/pass_through/messages/mcp_handler.py index ab722a60ca5d..8d86e5cff035 100644 --- a/litellm/llms/anthropic/pass_through/messages/mcp_handler.py +++ b/litellm/llms/anthropic/pass_through/messages/mcp_handler.py @@ -13,6 +13,7 @@ from litellm._logging import verbose_logger from litellm.responses.mcp.request_context import MCPRequestContext +from litellm.router_utils.mcp_tool_execution import MCPToolReplayGuard from litellm.types.llms.anthropic import ( AnthropicMessagesTool, AnthropicMessagesToolResultParam, @@ -137,6 +138,7 @@ async def anthropic_messages_with_mcp( messages=list(working_messages), stream=False, **base_call_args ) + replay_guard: Final = MCPToolReplayGuard() for _ in range(MAX_MCP_TOOL_USE_ITERATIONS): if _get_stop_reason(response) != "tool_use": break @@ -157,6 +159,7 @@ async def anthropic_messages_with_mcp( litellm_trace_id=context.litellm_trace_id, request_tags=list(context.request_tags) if context.request_tags else None, guardrail_context=context.guardrail_context, + on_tool_dispatched=replay_guard.record_tool_dispatch, ) # Every tool call was skipped, so there is nothing to feed back; a @@ -169,9 +172,10 @@ async def anthropic_messages_with_mcp( {"role": "assistant", "content": list(_get_response_content(response))}, _build_tool_result_message(tool_results), ) - response = await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn( - messages=list(working_messages), stream=False, **base_call_args - ) + with replay_guard: + response = await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn( + messages=list(working_messages), stream=False, **base_call_args + ) else: verbose_logger.warning( "MCP tool loop hit its %s iteration cap for model %s; returning the last response", diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 3e3b53f387df..da75968c3166 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -5954,6 +5954,7 @@ async def _obo_call_tool_with_retry( raw_headers: Mapping[str, str] | None = None, client_ip: str | None = None, allow_input_required: bool = False, + on_dispatch: Callable[[], None] | None = None, ) -> CallToolResult | InputRequiredResult: """Call a token_exchange (OBO) tool; on an upstream 401/403 re-mint the token once and retry. @@ -5968,6 +5969,7 @@ async def _obo_call_tool_with_retry( host_progress_callback=host_progress_callback, raise_on_error=True, allow_input_required=allow_input_required, + on_dispatch=on_dispatch, ) except Exception as exc: if _extract_upstream_auth_failure(exc) is None: @@ -5989,6 +5991,7 @@ async def _obo_call_tool_with_retry( call_tool_params, host_progress_callback=host_progress_callback, allow_input_required=allow_input_required, + on_dispatch=on_dispatch, ) async def _call_regular_mcp_tool( @@ -6007,6 +6010,7 @@ async def _call_regular_mcp_tool( user_api_key_auth: UserAPIKeyAuth | None = None, client_ip: str | None = None, allow_input_required: bool = False, + on_dispatch: Callable[[], None] | None = None, ) -> CallToolResult | InputRequiredResult: """ Call a regular MCP tool using the MCP client. @@ -6170,6 +6174,7 @@ async def _obo_call_tool_limited(): client=client, call_tool_params=call_tool_params, host_progress_callback=host_progress_callback, + on_dispatch=on_dispatch, mcp_server=mcp_server, server_auth_header=server_auth_header, extra_headers=extra_headers, @@ -6196,6 +6201,7 @@ async def _call_tool_via_client(client, params): params, host_progress_callback=host_progress_callback, allow_input_required=allow_input_required, + on_dispatch=on_dispatch, ) # The client-forwarded modes carry the caller's own upstream token, so an upstream # 401 (expired/invalid token) is the caller's to resolve: relay it as @@ -6212,6 +6218,7 @@ async def _call_tool_via_client(client, params): host_progress_callback=host_progress_callback, raise_on_error=True, allow_input_required=allow_input_required, + on_dispatch=on_dispatch, ) except Exception as e: auth_info: Final = _extract_upstream_auth_failure(e) @@ -6469,6 +6476,7 @@ async def call_tool( guardrail_context: Mapping[str, object] | None = None, client_ip: str | None = None, wire_compat: WireCompat = WireCompat.LEGACY, + on_dispatch: Callable[[], None] | None = None, ) -> CallToolResult | InputRequiredResult: """ Call a tool with the given name and arguments @@ -6484,6 +6492,8 @@ async def call_tool( litellm_logging_obj: Optional request logger the guardrail hooks record their evaluations onto, so MCP guardrail activity reaches the Guardrails Monitor. See ``pre_call_tool_check`` + on_dispatch: Optional callback run right before the call is sent to the server, + after its headers and client are ready, even if the server never returns a result Returns: @@ -6578,6 +6588,8 @@ async def _call_openapi_via_handler(): resolved_token: Final = _request_resolved_auth_headers.set(resolved_auth_headers) try: async with self._limit_outbound_concurrency(mcp_server): + if on_dispatch is not None: + on_dispatch() return await self._call_openapi_tool_handler(mcp_server, name, arguments, wire_compat) finally: _request_auth_header.reset(auth_token) @@ -6601,6 +6613,7 @@ async def _call_openapi_via_handler(): hook_extra_headers=hook_result.get("extra_headers"), user_api_key_auth=user_api_key_auth, allow_input_required=wire_compat is WireCompat.MODERN, + on_dispatch=on_dispatch, ) return await self._gather_openapi_tool_tasks(tasks, proxy_logging_obj) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f1ed8d3e9b38..a0d55d2391ad 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -35,6 +35,7 @@ ) from litellm.responses.mcp.request_context import MCPRequestContext from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.router_utils.mcp_tool_execution import MCPToolReplayGuard from litellm.types.llms.openai import ( PromptObject, Reasoning, @@ -337,6 +338,7 @@ async def aresponses_api_with_mcp( tools=tools, ) + replay_guard: Final = MCPToolReplayGuard() tool_results: Final = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( tool_server_map=tool_server_map, tool_calls=tool_calls, @@ -351,6 +353,7 @@ async def aresponses_api_with_mcp( guardrail_context=MCPRequestContext.resolve_guardrail_context( MappingProxyType({**kwargs, "metadata": metadata, "model": model}) ), + on_tool_dispatched=replay_guard.record_tool_dispatch, ) if tool_results: @@ -375,13 +378,14 @@ async def aresponses_api_with_mcp( tool_calls=tool_calls, tool_results=tool_results ) - final_response = await LiteLLM_Proxy_MCP_Handler._make_follow_up_call( - follow_up_input=follow_up_input, - model=model, - all_tools=all_tools, - response_id=previous_response_id if persistence_disabled else response.id, - **follow_up_call_params, - ) + with replay_guard: + final_response = await LiteLLM_Proxy_MCP_Handler._make_follow_up_call( + follow_up_input=follow_up_input, + model=model, + all_tools=all_tools, + response_id=previous_response_id if persistence_disabled else response.id, + **follow_up_call_params, + ) # If streaming and we have tool execution events, wrap the response if ( @@ -403,17 +407,18 @@ async def aresponses_api_with_mcp( # Add custom output elements to the final response (for non-streaming) elif isinstance(final_response, ResponsesAPIResponse): # Fetch MCP tools again for output elements (without OpenAI transformation) - ( - mcp_tools_for_output, - _, - ) = await LiteLLM_Proxy_MCP_Handler._process_mcp_tools_without_openai_transform( - user_api_key_auth=user_api_key_auth, - mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - request_tags=LiteLLM_Proxy_MCP_Handler._get_parent_request_tags(kwargs), - raw_headers=discovery_raw_headers, - ) + with replay_guard: + ( + mcp_tools_for_output, + _, + ) = await LiteLLM_Proxy_MCP_Handler._process_mcp_tools_without_openai_transform( + user_api_key_auth=user_api_key_auth, + mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + request_tags=LiteLLM_Proxy_MCP_Handler._get_parent_request_tags(kwargs), + raw_headers=discovery_raw_headers, + ) final_response = LiteLLM_Proxy_MCP_Handler._add_mcp_output_elements_to_response( response=final_response, mcp_tools_fetched=mcp_tools_for_output, diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index b17e0befba64..928ba9b5876a 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -10,6 +10,7 @@ LiteLLM_Proxy_MCP_Handler, ) from litellm.responses.mcp.request_context import MCPRequestContext +from litellm.router_utils.mcp_tool_execution import MCPToolReplayGuard from litellm.types.utils import Message, ModelResponse from litellm.utils import CustomStreamWrapper @@ -244,6 +245,7 @@ def __init__( self.follow_up_stream = None self.follow_up_iterator = None self.follow_up_exhausted = False + self.replay_guard = MCPToolReplayGuard() def __aiter__(self): return self @@ -382,7 +384,8 @@ async def __anext__(self): verbose_logger.debug("Follow-up stream iterator created") try: - chunk = await self.follow_up_iterator.__anext__() + with self.replay_guard: + chunk = await self.follow_up_iterator.__anext__() from litellm._logging import verbose_logger verbose_logger.debug("Follow-up chunk yielded: %s", chunk) @@ -445,6 +448,7 @@ async def _process_tool_calls(self): litellm_trace_id=self.litellm_trace_id, request_tags=self.request_tags, guardrail_context=context.guardrail_context, + on_tool_dispatched=self.replay_guard.record_tool_dispatch, ) async def _prepare_follow_up_call(self): @@ -606,6 +610,7 @@ def __next__(self): return initial_response # Execute tool calls + replay_guard: Final = MCPToolReplayGuard() tool_results: Final = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( tool_server_map=tool_server_map, tool_calls=tool_calls, @@ -618,6 +623,7 @@ def __next__(self): litellm_trace_id=context.litellm_trace_id, request_tags=request_tags, guardrail_context=context.guardrail_context, + on_tool_dispatched=replay_guard.record_tool_dispatch, ) if not tool_results: @@ -640,7 +646,8 @@ def __next__(self): follow_up_call_args["messages"] = follow_up_messages follow_up_call_args["stream"] = stream - response = await litellm_acompletion(**follow_up_call_args) + with replay_guard: + response = await litellm_acompletion(**follow_up_call_args) if isinstance(response, (ModelResponse, CustomStreamWrapper)): _add_mcp_metadata_to_response( response=response, diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 71f610791540..bf6a83345aea 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -696,8 +696,14 @@ async def _execute_tool_calls( litellm_trace_id: str | None = None, request_tags: list[str] | None = None, guardrail_context: Mapping[str, object] | None = None, + on_tool_dispatched: Callable[[], None] | None = None, ) -> list[MCPToolResult]: - """Execute tool calls and return results.""" + """Execute tool calls and return results. + + ``on_tool_dispatched`` runs for every call that is sent to its MCP server, even one that then + times out. Calls rejected or failing before that, for example on a pre-call guardrail or a + credential the caller has not set, do not trigger it. + """ from fastapi import HTTPException from litellm._uuid import uuid @@ -860,6 +866,7 @@ async def _execute_tool_calls( proxy_logging_obj=proxy_logging_obj, litellm_logging_obj=litellm_logging_obj, guardrail_context=guardrail_context, + on_dispatch=on_tool_dispatched, ) if proxy_logging_obj: diff --git a/litellm/router.py b/litellm/router.py index cfed5dc81c0a..8ecef3c46a36 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -224,6 +224,7 @@ send_llm_exception_alert, ) from litellm.router_utils.health_state_cache import DeploymentHealthCache +from litellm.router_utils.mcp_tool_execution import mcp_tools_executed from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, warn_on_unknown_model_group_affinity_flags, @@ -7249,7 +7250,12 @@ async def async_function_with_fallbacks_common_utils( hop_depth: Final = kwargs.get("fallback_depth") nested_fallback_hop: Final = isinstance(hop_depth, int) and hop_depth > 0 - if disable_fallbacks is True or original_model_group is None or is_guardrail_intervention(e): + if ( + disable_fallbacks is True + or original_model_group is None + or is_guardrail_intervention(e) + or mcp_tools_executed(e) + ): raise e input_kwargs: Final = { @@ -7663,7 +7669,7 @@ async def async_function_with_retries(self, *args, **kwargs): response = add_retry_headers_to_response(response=response, attempted_retries=0, max_retries=None) return response except Exception as e: - if is_guardrail_intervention(e): + if is_guardrail_intervention(e) or mcp_tools_executed(e): raise current_attempt = None original_exception = e @@ -7764,6 +7770,8 @@ async def async_function_with_retries(self, *args, **kwargs): # Always track the latest error so we raise the most # recent exception instead of the first one. original_exception = e + if mcp_tools_executed(e): + raise ## LOGGING kwargs = self.log_retry(kwargs=kwargs, e=e) diff --git a/litellm/router_utils/mcp_tool_execution.py b/litellm/router_utils/mcp_tool_execution.py new file mode 100644 index 000000000000..15f29fd1520f --- /dev/null +++ b/litellm/router_utils/mcp_tool_execution.py @@ -0,0 +1,62 @@ +from collections.abc import Iterator +from types import TracebackType +from typing import Final + +_MCP_TOOLS_EXECUTED_ATTR: Final = "_litellm_mcp_tools_executed" + +_MAX_EXCEPTION_CHAIN_DEPTH: Final = 20 + + +def mark_mcp_tools_executed(exception: BaseException) -> None: + """Tag an error raised after an MCP gateway loop already executed tool calls. + + Retrying or falling back re-runs the whole loop, which executes those tool calls + again, so the router surfaces a tagged error instead of replaying the request. The + exception is tagged rather than wrapped so the client sees the same error. + """ + setattr(exception, _MCP_TOOLS_EXECUTED_ATTR, True) + + +def _exception_chain(exception: BaseException) -> Iterator[BaseException]: + current = exception # rebind-ok: advances one link per iteration of the bounded walk + for _ in range(_MAX_EXCEPTION_CHAIN_DEPTH): + yield current + following = current.__cause__ or current.__context__ + if following is None: + return + current = following + + +def mcp_tools_executed(exception: BaseException | None) -> bool: + """Whether ``exception`` or an exception in its chain was tagged by ``mark_mcp_tools_executed``. + + The chain matters because a caller can re-raise a tagged error as a new exception, as + ``exception_type`` does for anything that is not already a LiteLLM exception. + """ + return exception is not None and any( + getattr(link, _MCP_TOOLS_EXECUTED_ATTR, False) is True for link in _exception_chain(exception) + ) + + +class MCPToolReplayGuard: + """Tags errors raised inside ``with guard:`` once the request has sent an MCP tool call to a server.""" + + __slots__ = ("_tool_dispatched",) + + def __init__(self) -> None: + self._tool_dispatched: bool = False + + def record_tool_dispatch(self) -> None: + self._tool_dispatched = True + + def __enter__(self) -> None: + return None + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: TracebackType | None, + ) -> None: + if self._tool_dispatched and isinstance(exc, Exception): + mark_mcp_tools_executed(exc) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py index 0dc7ac5ecd90..2b47a46dbe9c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py @@ -47,7 +47,9 @@ def _make_server(server_id: str, max_concurrent_requests: Optional[int]) -> MCPS def _patch_client_with_tracker(manager: MCPServerManager, tracker: _ConcurrencyTracker): async def fake_create_mcp_client(server, **kwargs): class _ProbeClient: - async def call_tool(self, params, host_progress_callback=None, allow_input_required=False): + async def call_tool( + self, params, host_progress_callback=None, allow_input_required=False, on_dispatch=None + ): tracker.enter(server.server_id) try: await asyncio.sleep(HOLD_SECONDS) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 0a7ea012b981..b8e52235afb5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -6911,7 +6911,7 @@ async def test_call_tool_without_broken_pipe_error(self): # Create mock client that tracks call_tool usage mock_client = AsyncMock() - async def mock_call_tool(params, host_progress_callback=None, allow_input_required=False): + async def mock_call_tool(params, host_progress_callback=None, allow_input_required=False, on_dispatch=None): # Return a mock CallToolResult result = MagicMock(spec=CallToolResult) result.content = [{"type": "text", "text": "Tool executed successfully"}] @@ -6952,6 +6952,39 @@ async def mock_call_tool(params, host_progress_callback=None, allow_input_requir # Verify the MCP client call was awaited exactly once assert mock_client.call_tool.await_count == 1 + @pytest.mark.asyncio + async def test_call_tool_hands_on_dispatch_to_the_mcp_client(self): + """The MCP client fires on_dispatch when it sends the call, so call_tool must hand it down.""" + from unittest.mock import AsyncMock, MagicMock + + from mcp.types import CallToolResult + + manager = MCPServerManager() + manager.registry = { + "test-server": MCPServer( + server_id="test-server", name="test-server", transport=MCPTransport.http, url="http://test-server.com" + ) + } + manager.tool_name_to_mcp_server_name_mapping["test-server-test_tool"] = "test-server" + manager._create_mcp_client = AsyncMock(return_value=_RetryFakeClient(result=CallToolResult(content=[]))) + proxy_logging_obj = MagicMock() + proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={}) + proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={}) + proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + dispatched = MagicMock() + + await manager.call_tool( + server_name="test-server", + name="test_tool", + arguments={}, + user_api_key_auth=MagicMock(object_permission=None, object_permission_id=None), + proxy_logging_obj=proxy_logging_obj, + on_dispatch=dispatched, + ) + + dispatched.assert_called_once_with() + @pytest.mark.asyncio async def test_get_allowed_mcp_servers_with_user_api_key_auth(self): """ @@ -10302,8 +10335,10 @@ def __init__(self, *, raises=None, result=None) -> None: self._MCPClient = MCPClient self.attempts = 0 - async def call_tool(self, params, host_progress_callback=None, raise_on_error=False, allow_input_required=False): + async def call_tool(self, params, host_progress_callback=None, raise_on_error=False, allow_input_required=False, on_dispatch=None): self.attempts += 1 + if on_dispatch is not None: + on_dispatch() if self._raises is not None: if raise_on_error: raise self._raises @@ -10438,6 +10473,56 @@ async def test_call_regular_routes_id_jag_through_the_retry_path(self): manager._cred_provider.invalidate_credentials.assert_awaited_once() assert first.attempts == 1 and retry.attempts == 1 + @pytest.mark.asyncio + @pytest.mark.parametrize("path", ["plain", "relay", "obo_retry"]) + async def test_call_regular_reports_every_sent_attempt_through_on_dispatch(self, path): + """on_dispatch tells the caller a tool may have run, so every attempt that is sent must reach + it: once on the plain and relay paths, twice when the OBO path retries after an upstream 401.""" + manager = self._manager() + success = CallToolResult(content=[], isError=False) + clients = [_RetryFakeClient(raises=_UpstreamAuthError(401)), _RetryFakeClient(result=success)] + manager._create_mcp_client = AsyncMock(side_effect=clients if path == "obo_retry" else clients[1:]) + servers = { + "plain": MCPServer( + server_id="plain-srv", name="plain", url="https://upstream.example/mcp", transport=MCPTransport.sse + ), + "relay": MCPServer( + server_id="relay-srv", + name="relay", + url="https://upstream.example/mcp", + transport=MCPTransport.sse, + auth_type=MCPAuth.true_passthrough, + ), + "obo_retry": MCPServer( + server_id="id-jag-srv", + name="id-jag", + url="https://upstream.example/mcp", + transport=MCPTransport.sse, + auth_type=MCPAuth.oauth2_id_jag, + client_id="gateway-client", + client_secret="gateway-secret", + token_exchange_endpoint="https://org-idp.example/oauth2/token", + id_jag_resource_token_endpoint="https://resource-as.example/oauth2/token", + ), + } + dispatched = MagicMock() + + result = await manager._call_regular_mcp_tool( + mcp_server=servers[path], + original_tool_name="tool", + arguments={}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={"Authorization": "Bearer caller-id-token"} if path == "obo_retry" else None, + raw_headers={"authorization": "Bearer caller-token"} if path == "relay" else None, + proxy_logging_obj=None, + on_dispatch=dispatched, + ) + + assert result is success + assert dispatched.call_count == (2 if path == "obo_retry" else 1) + @pytest.mark.asyncio async def test_non_auth_error_does_not_retry(self): manager = self._manager() @@ -10514,7 +10599,7 @@ async def test_obo_dispatch_respects_max_concurrent_requests(self): inflight = {"current": 0, "peak": 0} class _ConcurrencyRecordingClient: - async def call_tool(self, params, host_progress_callback=None, raise_on_error=False, allow_input_required=False): + async def call_tool(self, params, host_progress_callback=None, raise_on_error=False, allow_input_required=False, on_dispatch=None): inflight["current"] += 1 inflight["peak"] = max(inflight["peak"], inflight["current"]) try: diff --git a/tests/unit/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py index 1a56227b008b..29fdc216c6b1 100644 --- a/tests/unit/experimental_mcp_client/test_mcp_client.py +++ b/tests/unit/experimental_mcp_client/test_mcp_client.py @@ -777,6 +777,36 @@ async def test_call_tool_does_not_log_arguments(): assert "gpt-5-mini" not in logged +@pytest.mark.asyncio +@pytest.mark.parametrize("session_opens", [True, False]) +async def test_call_tool_reports_dispatch_only_when_the_call_is_sent(session_opens): + """ + on_dispatch marks the moment a tool call can reach the server, so it fires right before the + session sends the call and never when the session could not be opened. + """ + from mcp.types import CallToolRequestParams + + events: list[str] = [] + + async def send(**_): + events.append("sent") + return MagicMock() + + async def run_with_session(operation, quiet_on_error=False): + if not session_opens: + raise ConnectionError("cannot reach the server") + return await operation(MagicMock(call_tool=send)) + + client = MCPClient(server_url="http://test-server") + client.run_with_session = run_with_session + + await client.call_tool( + CallToolRequestParams(name="search_tool", arguments={}), on_dispatch=lambda: events.append("dispatched") + ) + + assert events == (["dispatched", "sent"] if session_opens else []) + + @pytest.mark.asyncio async def test_get_prompt_does_not_log_arguments(): from mcp.types import GetPromptRequestParams diff --git a/tests/unit/test_router_retry_non_retryable_errors.py b/tests/unit/test_router_retry_non_retryable_errors.py index 98a7db7a0792..ed9bc0af4bb6 100644 --- a/tests/unit/test_router_retry_non_retryable_errors.py +++ b/tests/unit/test_router_retry_non_retryable_errors.py @@ -12,14 +12,19 @@ import asyncio import datetime -from collections.abc import Awaitable, Callable -from typing import Final +import json +from collections.abc import AsyncIterator, Awaitable, Callable +from typing import TYPE_CHECKING, Final, NamedTuple from unittest.mock import AsyncMock, patch +import httpx import pytest +import respx import litellm from litellm import Router +from litellm.exceptions import GuardrailRaisedException +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( _union_duration_ms, @@ -27,8 +32,14 @@ ) from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.rules import Rules +from litellm.router_utils.mcp_tool_execution import mark_mcp_tools_executed +from litellm.types.router import RetryPolicy from litellm.utils import function_setup +if TYPE_CHECKING: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + from litellm.proxy._types import UserAPIKeyAuth + def _make_rate_limit_error(message="Rate limited"): """Create a RateLimitError for testing.""" @@ -66,7 +77,7 @@ def _make_not_found_error(message="Model not found"): ) -def _create_router(num_retries=2): +def _create_router(num_retries=2, retry_policy=None): """Create a Router with two deployments for testing.""" return Router( model_list=[ @@ -86,6 +97,7 @@ def _create_router(num_retries=2): }, ], num_retries=num_retries, + retry_policy=retry_policy, ) @@ -357,3 +369,534 @@ async def invoke(original_function: Callable[..., Awaitable[str]], *args: object assert timing_metrics["litellm_overhead_time_ms"] == pytest.approx( round(total_response_time_ms - union_duration_ms, 4) ) + + +def _make_follow_up_error(tools_executed: bool) -> litellm.InternalServerError: + error: Final = litellm.InternalServerError(message="follow-up failed", llm_provider="openai", model="gpt-4") + if tools_executed: + mark_mcp_tools_executed(error) + return error + + +@pytest.mark.asyncio +@pytest.mark.parametrize("retry_policy", [None, RetryPolicy(InternalServerErrorRetries=2)]) +@pytest.mark.parametrize("tools_executed", [True, False]) +async def test_500_after_mcp_tool_execution_is_not_retried(retry_policy, tools_executed): + """ + Regression test for #43153: a retry re-runs the whole MCP gateway loop, so a + failure raised after the loop executed tool calls must surface instead of + executing those tools again. A retry policy that retries 500s must not override + that, and a 500 raised before any tool ran must still be retried. + """ + router: Final = _create_router(num_retries=2, retry_policy=retry_policy) + error: Final = _make_follow_up_error(tools_executed) + attempt: Final = AsyncMock(side_effect=error) + + with pytest.raises(litellm.InternalServerError) as exc_info: + await router.async_function_with_retries(**{**_base_kwargs(), "original_function": attempt, "num_retries": 2}) + + assert exc_info.value is error + assert attempt.await_count == (1 if tools_executed else 3) + + +@pytest.mark.asyncio +async def test_500_after_mcp_tool_execution_stops_an_ongoing_retry_loop(): + """ + The first attempt fails before any tool ran and is retried; the retry executes + the tools and then fails, so the loop must stop there instead of running them + a second time. + """ + router: Final = _create_router(num_retries=3) + after_tools: Final = _make_follow_up_error(tools_executed=True) + attempt: Final = AsyncMock(side_effect=[_make_follow_up_error(tools_executed=False), after_tools]) + + with pytest.raises(litellm.InternalServerError) as exc_info: + await router.async_function_with_retries(**{**_base_kwargs(), "original_function": attempt, "num_retries": 3}) + + assert exc_info.value is after_tools + assert attempt.await_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tools_executed", [True, False]) +async def test_500_after_mcp_tool_execution_does_not_fall_back(tools_executed): + """ + A fallback re-runs the whole MCP gateway loop on the fallback group, so it must + not start once the primary attempt executed tool calls. + """ + router: Final = Router( + model_list=[ + {"model_name": "test-model", "litellm_params": {"model": "openai/gpt-4", "api_key": "fake-key-1"}}, + {"model_name": "backup-model", "litellm_params": {"model": "openai/gpt-4", "api_key": "fake-key-2"}}, + ], + num_retries=0, + fallbacks=[{"test-model": ["backup-model"]}], + ) + error: Final = _make_follow_up_error(tools_executed) + fallback_response: Final = litellm.ModelResponse() + + async def primary_fails(**kwargs): + if kwargs["model"] == "test-model": + raise error + return fallback_response + + attempt: Final = AsyncMock(side_effect=primary_fails) + call: Final = router.async_function_with_fallbacks( + **{**_base_kwargs(), "original_function": attempt, "num_retries": 0} + ) + + if tools_executed: + with pytest.raises(litellm.InternalServerError) as exc_info: + await call + assert exc_info.value is error + else: + assert await call is fallback_response + + attempted_models: Final = [awaited.kwargs["model"] for awaited in attempt.await_args_list] + assert attempted_models == (["test-model"] if tools_executed else ["test-model", "backup-model"]) + + +_UPSTREAM: Final = "https://upstream.example/v1" +_TOOL_NAME: Final = "ledger-commit_write" + + +def _ledger_tools() -> list[dict[str, str]]: + return [{"type": "mcp", "server_url": "litellm_proxy/ledger", "require_approval": "never"}] + + +def _ledger_caller() -> "UserAPIKeyAuth": + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth + + return UserAPIKeyAuth( + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="test", mcp_servers=["ledger"]) + ) + + +class _LedgerGateway(NamedTuple): + tool: AsyncMock + list_tools: AsyncMock + manager: "MCPServerManager" + + +@pytest.fixture +def ledger_gateway(monkeypatch: pytest.MonkeyPatch) -> _LedgerGateway: + """Serve one auto-approved MCP tool through the real MCP manager and tool registry.""" + from mcp.types import Tool + + from litellm.caching.caching import DualCache + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import mcp_server_manager, server, tool_registry + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing + from litellm.proxy.utils import ProxyLogging + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + manager: Final = mcp_server_manager.MCPServerManager() + manager.registry = { + "ledger": MCPServer( + server_id="ledger", + name="ledger", + server_name="ledger", + transport="http", + url="https://ledger.example/mcp", + spec_path="ledger.json", + auth_type="none", + ) + } + manager.tool_name_to_mcp_server_name_mapping = {_TOOL_NAME: "ledger"} + tool: Final = AsyncMock(return_value={"written": True}) + registry: Final = tool_registry.MCPToolRegistry() + registry.register_tool(_TOOL_NAME, "Write one ledger entry", {"type": "object"}, tool) + list_tools: Final = AsyncMock( + return_value=AggregateToolListing(tools=[Tool(name=_TOOL_NAME, inputSchema={"type": "object"})], outcomes={}) + ) + monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry) + monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache())) + monkeypatch.setattr(server, "_get_tools_from_mcp_servers", list_tools) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + return _LedgerGateway(tool=tool, list_tools=list_tools, manager=manager) + + +def _is_follow_up(request: httpx.Request) -> bool: + body: Final = json.loads(request.content) + messages: Final = body.get("messages") or [] + items: Final = body["input"] if isinstance(body.get("input"), list) else [] + return any(message.get("role") == "tool" for message in messages) or any( + item.get("type") == "function_call_output" for item in items + ) + + +def _chat_completion(message: dict[str, object], finish_reason: str) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 0, + "model": "gpt-4.1-mini", + "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + +def _responses_api_response(output_item: dict[str, object]) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "resp_1", + "object": "response", + "created_at": 0, + "status": "completed", + "model": "gpt-4.1-mini", + "output": [output_item], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + ) + + +def _tool_call(responses_api: bool) -> httpx.Response: + arguments: Final = json.dumps({"entry": "alpha"}) + if responses_api: + return _responses_api_response( + {"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": _TOOL_NAME, "arguments": arguments} + ) + return _chat_completion( + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": _TOOL_NAME, "arguments": arguments}} + ], + }, + "tool_calls", + ) + + +def _final_answer(responses_api: bool) -> httpx.Response: + if responses_api: + return _responses_api_response( + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Recorded alpha", "annotations": []}], + } + ) + return _chat_completion({"role": "assistant", "content": "Recorded alpha"}, "stop") + + +class _FakeModel: + """Answers both OpenAI routes: a tool call, then a final answer once the tool result arrives.""" + + def __init__(self, failing_initial_calls: int = 0, failing_follow_ups: int = 0) -> None: + self.calls: list[str] = [] + self._failures: Final = {"initial": failing_initial_calls, "follow_up": failing_follow_ups} + + def __call__(self, request: httpx.Request) -> httpx.Response: + kind: Final = "follow_up" if _is_follow_up(request) else "initial" + self.calls.append(kind) + if self.calls.count(kind) <= self._failures[kind]: + return httpx.Response(500, json={"error": {"message": "upstream failed", "type": "server_error"}}) + responses_api: Final = request.url.path.endswith("/responses") + return _tool_call(responses_api) if kind == "initial" else _final_answer(responses_api) + + +class _StreamBrokenBeforeFirstChunk(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + raise httpx.ReadError("connection dropped before the first chunk") + yield b"" + + +def _chat_tool_call_stream() -> httpx.Response: + tool_call: Final = { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": _TOOL_NAME, "arguments": json.dumps({"entry": "alpha"})}, + } + chunks: Final = ( + {"role": "assistant", "content": None, "tool_calls": [tool_call]}, + {}, + ) + body: Final = "".join( + "data: " + + json.dumps( + { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 0, + "model": "gpt-4.1-mini", + "choices": [{"index": 0, "delta": delta, "finish_reason": None if delta else "tool_calls"}], + } + ) + + "\n\n" + for delta in chunks + ) + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=f"{body}data: [DONE]\n\n") + + +def _serve(model: _FakeModel) -> None: + respx.post(f"{_UPSTREAM}/chat/completions").mock(side_effect=model) + respx.post(f"{_UPSTREAM}/responses").mock(side_effect=model) + + +def _mcp_router(num_retries: int, fallbacks: list[dict[str, list[str]]] | None = None) -> litellm.Router: + deployment: Final = {"model": "openai/gpt-4.1-mini", "api_key": "sk-test", "api_base": _UPSTREAM} + return litellm.Router( + model_list=[ + {"model_name": "repro-model", "litellm_params": deployment}, + {"model_name": "repro-model", "litellm_params": deployment}, + {"model_name": "backup-model", "litellm_params": deployment}, + ], + num_retries=num_retries, + fallbacks=fallbacks, + ) + + +def _chat(router: litellm.Router) -> Awaitable[object]: + return router.acompletion( + model="repro-model", + messages=[{"role": "user", "content": "Record entry alpha"}], + tools=_ledger_tools(), + metadata={"user_api_key_auth": _ledger_caller()}, + ) + + +def _responses(router: litellm.Router) -> Awaitable[object]: + return router.aresponses( + model="repro-model", + input="Record entry alpha", + tools=_ledger_tools(), + litellm_metadata={"user_api_key_auth": _ledger_caller()}, + ) + + +def _messages(router: litellm.Router) -> Awaitable[object]: + return router.aanthropic_messages( + model="repro-model", + max_tokens=64, + messages=[{"role": "user", "content": "Record entry alpha"}], + tools=_ledger_tools(), + litellm_metadata={"user_api_key_auth": _ledger_caller()}, + ) + + +_ENDPOINTS: Final = pytest.mark.parametrize( + "send", [_chat, _responses, _messages], ids=["chat_completions", "responses", "messages"] +) + + +@pytest.mark.asyncio +@_ENDPOINTS +@pytest.mark.parametrize( + "router_kwargs", + [{"num_retries": 2}, {"num_retries": 0, "fallbacks": [{"repro-model": ["backup-model"]}]}], + ids=["retries", "fallbacks"], +) +@respx.mock +async def test_router_does_not_replay_an_executed_mcp_tool( + ledger_gateway: _LedgerGateway, + send: Callable[[litellm.Router], Awaitable[object]], + router_kwargs: dict[str, object], +): + """ + Regression test for #43153: an auto-approved MCP tool must run once per request. + + Given: The model asks for the tool, the tool runs, and every follow-up model call returns a 500 + When: The router has retries or a fallback group configured + Then: The client gets the 500 after one tool execution and one follow-up call + + A retry or fallback re-runs the whole MCP loop, so before the fix every replay executed + the tool again (three writes with num_retries=2, two with one fallback). + """ + model: Final = _FakeModel(failing_follow_ups=3) + _serve(model) + + with pytest.raises(litellm.InternalServerError): + await send(_mcp_router(**router_kwargs)) + + assert ledger_gateway.tool.await_count == 1 + assert model.calls == ["initial", "follow_up"] + + +@pytest.mark.asyncio +@_ENDPOINTS +@respx.mock +async def test_router_retries_when_a_guardrail_blocked_every_mcp_tool( + ledger_gateway: _LedgerGateway, + send: Callable[[litellm.Router], Awaitable[object]], + monkeypatch: pytest.MonkeyPatch, +): + """ + Given: A pre-call guardrail blocks the tool call, so no tool runs + When: The follow-up model call fails once with a 500 + Then: The router still retries, because replaying the request cannot repeat a side effect + """ + + class BlockEveryToolCall(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + raise GuardrailRaisedException(message="tool calls are blocked", blocked_content=True) + + monkeypatch.setattr( + litellm, "callbacks", [BlockEveryToolCall(guardrail_name="block", event_hook="pre_mcp_call", default_on=True)] + ) + model: Final = _FakeModel(failing_follow_ups=1) + _serve(model) + + await send(_mcp_router(num_retries=2)) + + assert ledger_gateway.tool.await_count == 0 + assert model.calls == ["initial", "follow_up", "initial", "follow_up"] + + +@pytest.mark.asyncio +@_ENDPOINTS +@respx.mock +async def test_router_does_not_replay_a_tool_call_that_failed_after_it_was_sent( + ledger_gateway: _LedgerGateway, + send: Callable[[litellm.Router], Awaitable[object]], + monkeypatch: pytest.MonkeyPatch, +): + """ + Given: The tool runs, but its call fails afterwards, here because a during-call guardrail + rejects the result, which leaves the router in the same spot as a timeout after the write + When: Every follow-up model call returns a 500 + Then: The router does not replay the request, because the tool may already have written + """ + tool_ran: Final = asyncio.Event() + + def write(**_: object) -> dict[str, bool]: + tool_ran.set() + return {"written": True} + + class RejectToolResult(CustomGuardrail): + async def async_moderation_hook(self, data, user_api_key_dict, call_type): + await tool_ran.wait() + raise GuardrailRaisedException(message="tool result rejected", blocked_content=True) + + ledger_gateway.tool.side_effect = write + monkeypatch.setattr( + litellm, "callbacks", [RejectToolResult(guardrail_name="reject", event_hook="during_mcp_call", default_on=True)] + ) + model: Final = _FakeModel(failing_follow_ups=3) + _serve(model) + + with pytest.raises(litellm.InternalServerError): + await send(_mcp_router(num_retries=2)) + + assert ledger_gateway.tool.await_count == 1 + assert model.calls == ["initial", "follow_up"] + + +@pytest.mark.asyncio +@respx.mock +async def test_router_retries_when_an_mcp_tool_call_never_reached_its_server( + ledger_gateway: _LedgerGateway, + monkeypatch: pytest.MonkeyPatch, +): + """ + Given: The MCP server needs a per-user header the caller has not set, so the tool call fails + before anything is sent to the server + When: The follow-up model call fails once with a 500 + Then: The router still retries, because replaying the request cannot repeat a side effect + """ + registry: Final = ledger_gateway.manager.registry + monkeypatch.setitem( + registry, + "ledger", + registry["ledger"].model_copy( + update={ + "spec_path": None, + "static_headers": {"Authorization": "Bearer ${LEDGER_TOKEN}"}, + "env_vars": [{"name": "LEDGER_TOKEN", "scope": "user"}], + } + ), + ) + model: Final = _FakeModel(failing_follow_ups=1) + _serve(model) + + await _chat(_mcp_router(num_retries=2)) + + assert model.calls == ["initial", "follow_up", "initial", "follow_up"] + + +@pytest.mark.asyncio +@respx.mock +async def test_router_does_not_fall_back_when_a_streamed_follow_up_breaks_after_a_tool_ran( + ledger_gateway: _LedgerGateway, +): + """ + Given: A streaming Chat Completions request whose tool ran + When: The follow-up stream breaks before its first chunk and a fallback group is configured + Then: The client gets the error instead of a fallback that runs the tool again + """ + calls: Final[list[str]] = [] + + def stream_model(request: httpx.Request) -> httpx.Response: + kind: Final = "follow_up" if _is_follow_up(request) else "initial" + calls.append(kind) + if kind == "follow_up": + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=_StreamBrokenBeforeFirstChunk() + ) + return _chat_tool_call_stream() + + respx.post(f"{_UPSTREAM}/chat/completions").mock(side_effect=stream_model) + stream: Final = await _mcp_router(num_retries=0, fallbacks=[{"repro-model": ["backup-model"]}]).acompletion( + model="repro-model", + messages=[{"role": "user", "content": "Record entry alpha"}], + tools=_ledger_tools(), + metadata={"user_api_key_auth": _ledger_caller()}, + stream=True, + ) + + with pytest.raises(litellm.APIConnectionError): + async for _ in stream: + pass + + assert ledger_gateway.tool.await_count == 1 + assert calls == ["initial", "follow_up"] + + +@pytest.mark.asyncio +@respx.mock +async def test_router_still_retries_a_failure_before_any_mcp_tool_ran(ledger_gateway: _LedgerGateway): + """ + Given: The initial model call fails once with a 500, then asks for the tool + When: The router retries + Then: The retry succeeds and the tool runs once + """ + model: Final = _FakeModel(failing_initial_calls=1) + _serve(model) + + await _chat(_mcp_router(num_retries=2)) + + assert ledger_gateway.tool.await_count == 1 + assert model.calls == ["initial", "initial", "follow_up"] + + +@pytest.mark.asyncio +@respx.mock +async def test_responses_tool_listing_failure_after_a_tool_ran_is_not_replayed(ledger_gateway: _LedgerGateway): + """ + Given: A Responses request whose tool ran and whose follow-up succeeded + When: Listing the MCP tools again for the output items fails + Then: The router surfaces that failure instead of replaying the request and the tool + """ + listing: Final = ledger_gateway.list_tools.return_value + listing_error: Final = RuntimeError("tool listing failed") + ledger_gateway.list_tools.side_effect = [listing, listing_error, listing, listing_error, listing, listing_error] + model: Final = _FakeModel() + _serve(model) + + with pytest.raises(Exception, match="tool listing failed"): + await _responses(_mcp_router(num_retries=2)) + + assert ledger_gateway.tool.await_count == 1 + assert ledger_gateway.list_tools.await_count == 2 + assert model.calls == ["initial", "follow_up"]