diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 030911238ec..9e51e2533ea 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Final, cast from litellm._logging import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.agentic_loop_settings import validated_max_agentic_loops from litellm.responses.litellm_completion_transformation.streaming_iterator import LiteLLMCompletionStreamingIterator from litellm.responses.mcp.request_context import MCPRequestContext from litellm.responses.mcp.stream_round_billing import StreamedRoundBilling @@ -289,6 +290,10 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.mcp_tools_with_litellm_proxy = mcp_tools_with_litellm_proxy or [] self.user_api_key_auth = user_api_key_auth self.original_request_params = original_request_params or {} + configured_rounds: Final = validated_max_agentic_loops( + self.original_request_params.get("max_agentic_loops"), field="litellm_params.max_agentic_loops" + ) + self.max_tool_call_rounds: int = MAX_MCP_TOOL_CALL_ROUNDS if configured_rounds is None else configured_rounds self.should_auto_execute = self._should_auto_execute_tools() # Streaming state management @@ -509,7 +514,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): 2. mcp_discovery - Emit MCP discovery events (after response.output_item.added) 3. continue_initial_response - Stream the current round's response (initial or follow-up). On completion, if auto-execute is on and the response contains tool calls, loops back - through tool_execution/follow-up instead of ending, up to MAX_MCP_TOOL_CALL_ROUNDS. + through tool_execution/follow-up instead of ending, up to self.max_tool_call_rounds. 4. tool_execution - Emit tool execution events 5. finished - End iteration """ @@ -1015,14 +1020,14 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Remove tool_choice to avoid forcing more tool calls follow_up_params.pop("tool_choice", None) - if self.tool_call_round >= MAX_MCP_TOOL_CALL_ROUNDS: + if self.tool_call_round >= self.max_tool_call_rounds: # Hit the round cap: drop tools entirely so the model must # answer in text instead of emitting another (unexecuted) # tool call that would otherwise end the stream in silence. follow_up_params.pop("tools", None) verbose_logger.warning( - "MCP auto-execute hit MAX_MCP_TOOL_CALL_ROUNDS=%s; forcing a text-only follow-up.", - MAX_MCP_TOOL_CALL_ROUNDS, + "MCP auto-execute hit max_agentic_loops=%s; forcing a text-only follow-up.", + self.max_tool_call_rounds, ) follow_up_response: Final = await aresponses(**follow_up_params) diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index e88d951ff09..285d90ef8de 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -29,6 +29,7 @@ from litellm.proxy._experimental.mcp_server import mcp_server_manager from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ServerListOk from litellm.proxy.utils import _estimate_dispatched_failure_usage, _failure_fields_to_lift from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler +from litellm.responses.mcp.mcp_streaming_iterator import MCPEnhancedStreamingIterator from litellm.types.llms.openai import ( ResponsesAPIResponse, ResponsesAPIStreamingResponse, @@ -2384,8 +2385,11 @@ def _first_index(events: Sequence[object], event_type: str) -> int: async def _scripted_mcp_stream( - model: str, upstream_url: str, rounds: Sequence[bytes] -) -> tuple[tuple[object, ...], respx.Route, AsyncMock, AsyncMock]: + model: str, + upstream_url: str, + rounds: Sequence[bytes], + max_agentic_loops: int | None = None, +) -> tuple[tuple[object, ...], respx.Route, AsyncMock, AsyncMock, MCPEnhancedStreamingIterator]: upstream: Final = respx.post(upstream_url).mock( side_effect=[httpx.Response(200, content=body, headers=SSE_HEADERS) for body in rounds] ) @@ -2402,9 +2406,11 @@ async def _scripted_mcp_stream( tool_choice="required", input=[{"role": "user", "type": "message", "content": "What's the weather in San Francisco?"}], stream=True, + **({"max_agentic_loops": max_agentic_loops} if max_agentic_loops is not None else {}), ) + assert isinstance(response, MCPEnhancedStreamingIterator) events: Final = tuple([event async for event in response]) - return events, upstream, fetch, execute + return events, upstream, fetch, execute, response @pytest.mark.asyncio @@ -2417,7 +2423,7 @@ async def test_streaming_responses_api_with_mcp_tools( model, upstream_url, tool_round, answer_round = SCRIPTED_PROVIDERS[provider] with caplog.at_level(logging.ERROR): - events, upstream, fetch, execute = await _scripted_mcp_stream( + events, upstream, fetch, execute, _ = await _scripted_mcp_stream( model, upstream_url, (tool_round(), answer_round()) ) @@ -2436,6 +2442,22 @@ async def test_streaming_responses_api_with_mcp_tools( assert not [record.getMessage() for record in caplog.records if record.levelno >= logging.ERROR] +@pytest.mark.asyncio +@respx.mock +async def test_streaming_mcp_iterator_uses_max_agentic_loops_from_request(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + model, upstream_url, tool_round, answer_round = SCRIPTED_PROVIDERS["openai"] + + _, _, _, _, iterator = await _scripted_mcp_stream( + model, + upstream_url, + (tool_round(), answer_round()), + max_agentic_loops=7, + ) + + assert iterator.max_tool_call_rounds == 7 + + @pytest.mark.asyncio @respx.mock async def test_streaming_mcp_event_order_and_response_id_consistency( @@ -2445,7 +2467,7 @@ async def test_streaming_mcp_event_order_and_response_id_consistency( model, upstream_url, tool_round, answer_round = SCRIPTED_PROVIDERS["openai"] with caplog.at_level(logging.ERROR): - events, _, _, _ = await _scripted_mcp_stream(model, upstream_url, (tool_round(), answer_round())) + events, _, _, _, _ = await _scripted_mcp_stream(model, upstream_url, (tool_round(), answer_round())) created: Final = _first_index(events, "response.created") in_progress: Final = _first_index(events, "response.in_progress") diff --git a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py index 8eb73b7c6e1..ed39da1b7a9 100644 --- a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py +++ b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py @@ -2,6 +2,7 @@ import sys import types from contextlib import nullcontext from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -9,7 +10,6 @@ from mcp.types import CallToolResult, TextContent import litellm # noqa: F401 - ensures litellm.responses.main is registered in sys.modules from litellm.responses.mcp.mcp_streaming_iterator import ( - MAX_MCP_TOOL_CALL_ROUNDS, MCPEnhancedStreamingIterator, ) from litellm.types.llms.openai import ( @@ -96,7 +96,7 @@ def _mock_mcp_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: return call_tool -def _make_iterator(initial_chunks) -> MCPEnhancedStreamingIterator: +def _make_iterator(initial_chunks, max_agentic_loops: int | str | None = None) -> MCPEnhancedStreamingIterator: return MCPEnhancedStreamingIterator( base_iterator=_FakeAsyncStream(initial_chunks), mcp_events=[], @@ -107,6 +107,7 @@ def _make_iterator(initial_chunks) -> MCPEnhancedStreamingIterator: "model": "gpt-4", "input": "what is berriai/litellm?", "tools": [{"type": "mcp"}], + **({"max_agentic_loops": max_agentic_loops} if max_agentic_loops is not None else {}), }, ) @@ -157,35 +158,43 @@ async def test_second_round_tool_call_is_executed_and_reaches_final_text(monkeyp @pytest.mark.asyncio -async def test_tool_call_rounds_are_capped(monkeypatch): +@pytest.mark.parametrize( + ("max_agentic_loops", "expected_cap"), + ( + pytest.param(None, 5, id="unset"), + pytest.param(7, 7, id="integer"), + pytest.param("7", 7, id="integer-string"), + pytest.param(2, 2, id="custom-cap"), + ), +) +async def test_tool_call_rounds_are_capped(monkeypatch, max_agentic_loops, expected_cap): """ A model that keeps calling tools every round must not loop forever — - auto-execution stops at MAX_MCP_TOOL_CALL_ROUNDS, and the follow-up made + auto-execution stops at the configured round cap, and the follow-up made at the cap drops "tools" from the request so the model is forced to answer in text instead of the stream just hanging. """ call_tool = _mock_mcp_environment(monkeypatch) - # Rounds 2..MAX_MCP_TOOL_CALL_ROUNDS keep calling the tool; the call made - # once the cap is hit returns a text-only response. - tool_call_streams = [ - _tool_call_stream(f"call_{i}", "read_wiki_contents") for i in range(2, MAX_MCP_TOOL_CALL_ROUNDS + 1) - ] - aresponses_mock = AsyncMock(side_effect=[*tool_call_streams, _text_only_stream("giving up on tools")]) + tool_call_streams: Final = tuple( + _tool_call_stream(f"call_{i}", "read_wiki_contents") for i in range(2, expected_cap + 1) + ) + aresponses_mock: Final = AsyncMock(side_effect=(*tool_call_streams, _text_only_stream("giving up on tools"))) monkeypatch.setattr(responses_main_module, "aresponses", aresponses_mock) iterator = _make_iterator( [ _output_item_added_chunk(), _completed_chunk([_function_call("call_1", "read_wiki_contents")]), - ] + ], + max_agentic_loops=max_agentic_loops, ) _ = [chunk async for chunk in iterator] - assert iterator.tool_call_round == MAX_MCP_TOOL_CALL_ROUNDS - assert call_tool.call_count == MAX_MCP_TOOL_CALL_ROUNDS - assert aresponses_mock.call_count == MAX_MCP_TOOL_CALL_ROUNDS + assert iterator.tool_call_round == expected_cap + assert call_tool.call_count == expected_cap + assert aresponses_mock.call_count == expected_cap # Every follow-up before the cap still offered tools; only the capped one drops them. for call in aresponses_mock.call_args_list[:-1]: @@ -193,6 +202,12 @@ async def test_tool_call_rounds_are_capped(monkeypatch): assert "tools" not in aresponses_mock.call_args_list[-1].kwargs +@pytest.mark.parametrize(("max_agentic_loops", "expected_error"), ((0, ValueError), ("three", TypeError))) +def test_invalid_max_agentic_loops_raises_when_creating_iterator(max_agentic_loops, expected_error): + with pytest.raises(expected_error, match=r"litellm_params\.max_agentic_loops"): + _make_iterator([], max_agentic_loops=max_agentic_loops) + + @pytest.mark.asyncio async def test_continuation_id_is_final_round_not_interim_tool_call(monkeypatch): """