scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

fix(responses): keep every MCP tool round in streaming follow-up input (#47698)

Browse source 
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: ryan-crabbe-berri <ryan@berri.ai>
This commit is contained in:
Shivam Rawat 2026-10-28 00:51:10 -07:00 • committed by GitHub
parent 6fe0263761
commit f1795e6d44
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 232 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/responses/mcp/mcp_streaming_iterator.py
View file

@ -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()

174
tests/integration/mcp/test_responses_mcp_streaming_round_history.py Normal file
View file

@ -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

53
tests/unit/responses/mcp/test_mcp_streaming_iterator.py
View file

@ -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"),