From f1795e6d447a4e5a3660de0d2fde5e1f44f52929 Mon Sep 17 00:00:00 2001 From: Shivam Rawat Date: Wed, 28 Oct 2026 00:51:10 -0700 Subject: [PATCH] fix(responses): keep every MCP tool round in streaming follow-up input (#47698) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: ryan-crabbe-berri --- .../responses/mcp/mcp_streaming_iterator.py | 7 +- ...t_responses_mcp_streaming_round_history.py | 174 ++++++++++++++++++ .../mcp/test_mcp_streaming_iterator.py | 53 +++++- 3 files changed, 232 insertions(+), 2 deletions(-) create mode 100644 tests/integration/mcp/test_responses_mcp_streaming_round_history.py diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index cc79d1c88bb..189366feab4 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -20,6 +20,7 @@ from litellm.types.llms.openai import ( MCPListToolsCompletedEvent, MCPListToolsFailedEvent, MCPListToolsInProgressEvent, + ResponseInputParam, ResponsesAPIResponse, ResponsesAPIStreamEvents, ResponsesAPIStreamingResponse, @@ -290,6 +291,7 @@ 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 {} + self._follow_up_history: str | ResponseInputParam | None = self.original_request_params.get("input") configured_rounds: Final = validated_max_agentic_loops( self.original_request_params.get("max_agentic_loops"), field="litellm_params.max_agentic_loops" ) @@ -1003,9 +1005,12 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): follow_up_input: Final = LiteLLM_Proxy_MCP_Handler.create_follow_up_input( response=self.collected_response, tool_results=self.tool_results, - original_input=self.original_request_params.get("input"), + original_input=self._follow_up_history, preserve_reasoning=persistence_disabled, ) + self._follow_up_history = cast( # cast-ok: helper returns valid input items + "ResponseInputParam", follow_up_input + ) # Make follow-up call with streaming follow_up_params: Final = self.original_request_params.copy() diff --git a/tests/integration/mcp/test_responses_mcp_streaming_round_history.py b/tests/integration/mcp/test_responses_mcp_streaming_round_history.py new file mode 100644 index 00000000000..8313d96e23c --- /dev/null +++ b/tests/integration/mcp/test_responses_mcp_streaming_round_history.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +import json +import uuid +from collections.abc import Mapping +from typing import Final + +from integration._support.client import Gateway +from integration._support.mcp import mcp_peer, register_mcp, tool_calls +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +TOOL_ARGUMENTS: Final = '{"a":2,"b":3}' +ANSWER: Final = "All three tool rounds are done." + + +def _event_frame(event: Mapping[str, JsonValue]) -> bytes: + return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + + +def _streamed_response(round_number: int, model: str, tool_name: str) -> Reply: + response_id: Final = f"resp_mcp_round_{round_number}" + call_id: Final = f"call_{round_number}" + is_tool_call: Final = round_number <= 3 + output_item: Final[dict[str, JsonValue]] = ( + { + "type": "function_call", + "id": f"fc_{round_number}", + "call_id": call_id, + "name": tool_name, + "arguments": TOOL_ARGUMENTS, + "status": "completed", + } + if is_tool_call + else { + "type": "message", + "id": f"msg_{round_number}", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": ANSWER, "annotations": []}], + } + ) + added_item: Final[dict[str, JsonValue]] = ( + {**output_item, "arguments": "", "status": "in_progress"} + if is_tool_call + else {**output_item, "content": [], "status": "in_progress"} + ) + item_events: Final[tuple[dict[str, JsonValue], ...]] = ( + ( + {"type": "response.output_item.added", "output_index": 0, "item": added_item}, + { + "type": "response.function_call_arguments.delta", + "item_id": f"fc_{round_number}", + "output_index": 0, + "delta": TOOL_ARGUMENTS, + }, + { + "type": "response.function_call_arguments.done", + "item_id": f"fc_{round_number}", + "output_index": 0, + "arguments": TOOL_ARGUMENTS, + }, + {"type": "response.output_item.done", "output_index": 0, "item": output_item}, + ) + if is_tool_call + else ( + {"type": "response.output_item.added", "output_index": 0, "item": added_item}, + { + "type": "response.output_text.delta", + "item_id": f"msg_{round_number}", + "output_index": 0, + "content_index": 0, + "delta": ANSWER, + }, + {"type": "response.output_item.done", "output_index": 0, "item": output_item}, + ) + ) + response: Final[dict[str, JsonValue]] = { + "id": response_id, + "object": "response", + "created_at": 1, + "status": "completed", + "model": model, + "output": [output_item], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + events: Final[tuple[dict[str, JsonValue], ...]] = ( + { + "type": "response.created", + "response": {**response, "status": "in_progress", "output": []}, + }, + { + "type": "response.in_progress", + "response": {**response, "status": "in_progress", "output": []}, + }, + *item_events, + {"type": "response.completed", "response": response}, + ) + chunks: Final = tuple( + _event_frame({**event, "sequence_number": sequence_number}) for sequence_number, event in enumerate(events) + ) + return Reply(content_type="text/event-stream", chunks=chunks) + + +def _tool_call_history(body: Mapping[str, JsonValue]) -> tuple[tuple[str, str], ...]: + input_items: Final = body["input"] + assert isinstance(input_items, list) + return tuple( + (str(item["type"]), str(item["call_id"])) + for item in input_items + if isinstance(item, dict) and item.get("type") in ("function_call", "function_call_output") + ) + + +def test_streaming_responses_mcp_follow_up_keeps_every_tool_round(gateway: Gateway) -> None: + alias: Final = "hist" + uuid.uuid4().hex[:8] + upstream_requests: Final[list[dict[str, JsonValue]]] = [] + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target.endswith("/models"): + return Reply(body=b'{"object":"list","data":[]}') + assert request.target.endswith("/responses"), request.target + request_body: Final = JSON_OBJECT.validate_json(request.body) + upstream_requests.append(request_body) + return _streamed_response(len(upstream_requests), str(request_body["model"]), f"{alias}-add") + + with mcp_peer() as peer, wire_server(respond) as wire, gateway.scenario() as scenario: + server_id: Final = register_mcp(scenario, peer, alias) + model: Final = scenario.model(model="openai/responses/gpt-4o-mini", api_base=wire.url + "/v1") + key: Final = scenario.key(object_permission={"mcp_servers": [server_id]}) + peer.drain() + response: Final = gateway.client.post( + "/v1/responses", + headers={"Authorization": f"Bearer {key}"}, + json={ + "model": model, + "input": "run each tool round in order", + "stream": True, + "tools": [ + { + "type": "mcp", + "server_url": "litellm_proxy", + "server_label": "litellm", + "require_approval": "never", + } + ], + }, + timeout=90, + ) + assert response.status_code == 200, response.text + assert "response.completed" in response.text and ANSWER in response.text, response.text + + assert len(upstream_requests) == 4, upstream_requests + assert _tool_call_history(upstream_requests[2]) == ( + ("function_call", "call_1"), + ("function_call_output", "call_1"), + ("function_call", "call_2"), + ("function_call_output", "call_2"), + ) + assert _tool_call_history(upstream_requests[3]) == ( + ("function_call", "call_1"), + ("function_call_output", "call_1"), + ("function_call", "call_2"), + ("function_call_output", "call_2"), + ("function_call", "call_3"), + ("function_call_output", "call_3"), + ) + third_input: Final = upstream_requests[2]["input"] + fourth_input: Final = upstream_requests[3]["input"] + assert isinstance(third_input, list) and isinstance(fourth_input, list) + assert sum(1 for item in third_input if isinstance(item, dict) and item.get("role") == "user") == 1 + assert sum(1 for item in fourth_input if isinstance(item, dict) and item.get("role") == "user") == 1 + assert len(tool_calls(peer.drain())) == 3 diff --git a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py index 78b0e3d4e5b..f745b5bfadf 100644 --- a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py +++ b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py @@ -2,7 +2,7 @@ import sys import types from contextlib import nullcontext from types import SimpleNamespace -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock import pytest @@ -66,6 +66,14 @@ def _item_type(item: dict[str, object] | BaseLiteLLMOpenAIResponseObject) -> str return str(item["type"]) if isinstance(item, dict) else str(item.type) +def _follow_up_tool_history(input_items: list[object]) -> tuple[tuple[str, str], ...]: + return tuple( + (str(item.get("type")), str(item.get("call_id"))) + for item in input_items + if isinstance(item, dict) and item.get("type") in ("function_call", "function_call_output") + ) + + def _tool_call_stream(call_id: str, tool_name: str, response_id: str = "resp-1") -> _FakeAsyncStream: return _FakeAsyncStream([_completed_chunk([_function_call(call_id, tool_name)], response_id=response_id)]) @@ -157,6 +165,49 @@ async def test_second_round_tool_call_is_executed_and_reaches_final_text(monkeyp assert final_output[-1]["content"][0]["text"] == "Here's what I found after retrying." +@pytest.mark.asyncio +async def test_follow_up_input_keeps_every_previous_tool_round(monkeypatch: pytest.MonkeyPatch) -> None: + _mock_mcp_environment(monkeypatch) + + aresponses_mock: Final = AsyncMock( + side_effect=[ + _tool_call_stream("call_2", "read_wiki_contents"), + _tool_call_stream("call_3", "read_wiki_contents"), + _text_only_stream("All three tool rounds are done."), + ] + ) + monkeypatch.setattr(responses_main_module, "aresponses", aresponses_mock) + iterator: Final = _make_iterator( + [ + _output_item_added_chunk(), + _completed_chunk([_function_call("call_1", "read_wiki_contents")]), + ] + ) + + _ = [chunk async for chunk in iterator] + + second_follow_up_input: Final = cast("list[object]", aresponses_mock.call_args_list[1].kwargs["input"]) + third_follow_up_input: Final = cast("list[object]", aresponses_mock.call_args_list[2].kwargs["input"]) + assert _follow_up_tool_history(second_follow_up_input) == ( + ("function_call", "call_1"), + ("function_call_output", "call_1"), + ("function_call", "call_2"), + ("function_call_output", "call_2"), + ) + assert _follow_up_tool_history(third_follow_up_input) == ( + ("function_call", "call_1"), + ("function_call_output", "call_1"), + ("function_call", "call_2"), + ("function_call_output", "call_2"), + ("function_call", "call_3"), + ("function_call_output", "call_3"), + ) + user_messages: Final = [ + item for item in third_follow_up_input if isinstance(item, dict) and item.get("role") == "user" + ] + assert user_messages == [{"type": "message", "role": "user", "content": "what is berriai/litellm?"}] + + @pytest.mark.asyncio @pytest.mark.parametrize( ("max_agentic_loops", "expected_cap"),