scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

fix: preserve chat follow-up normalization

Browse source
This commit is contained in:
tarunag10 2026-08-20 13:14:43 +05:30
parent 05d4996a41
commit e5ffc4f8db
3 changed files with 17 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/responses/mcp/chat_completions_handler.py
View file

@ -453,7 +453,10 @@ async def acompletion_with_mcp(
)
# Make follow-up call with streaming
follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_chained_call_params(self.base_call_args)
follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params(
self.base_call_args,
original_stream_setting=True,
)
follow_up_call_args["messages"] = follow_up_messages
follow_up_call_args["stream"] = True
# Ensure follow-up call doesn't trigger MCP handler again
@ -625,7 +628,10 @@ async def acompletion_with_mcp(
)
# Make follow-up call with original stream setting
follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_chained_call_params(base_call_args)
follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params(
base_call_args,
original_stream_setting=stream,
)
follow_up_call_args["messages"] = follow_up_messages
follow_up_call_args["stream"] = stream

6
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

@ -72,7 +72,9 @@ class LiteLLM_Proxy_MCP_Handler:
"""
@staticmethod
def _prepare_chained_call_params(params: Mapping[str, Any]) -> dict[str, Any]:
def _prepare_chained_call_params(
params: Mapping[str, Any],
) -> dict[str, Any]: # mutable-ok: returns a sanitized copy for the next provider call
"""Copy request params without state owned by the previous LLM call.
MCP auto-execution keeps the trace identifier so chained rounds remain
@ -1264,7 +1266,7 @@ class LiteLLM_Proxy_MCP_Handler:
Restores the original streaming setting and removes tool_choice since
we're now providing tool results, not requesting tool calls.
"""
follow_up_params: Final = call_params.copy()
follow_up_params: Final = LiteLLM_Proxy_MCP_Handler._prepare_chained_call_params(call_params)
# Restore original streaming setting for follow-up call
follow_up_params["stream"] = original_stream_setting

7
tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py
View file

@ -444,12 +444,13 @@ async def test_execute_tool_calls_uses_unique_call_ids_and_preserves_parent_cont
assert all(item["metadata"]["parent_litellm_call_id"] == "cid" for item in captured)
def test_prepare_chained_call_params_resets_per_call_logging_state():
def test_prepare_follow_up_call_params_resets_per_call_logging_state():
original = {
"model": "gpt-4",
"litellm_call_id": "parent-call",
"litellm_logging_obj": object(),
"litellm_trace_id": "trace-1",
"tool_choice": "auto",
"litellm_params": {
"litellm_call_id": "nested-parent-call",
"litellm_logging_obj": object(),
@ -457,13 +458,15 @@ def test_prepare_chained_call_params_resets_per_call_logging_state():
},
}
follow_up = LiteLLM_Proxy_MCP_Handler._prepare_chained_call_params(original)
follow_up = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params(original, original_stream_setting=True)
assert "litellm_call_id" not in follow_up
assert "litellm_logging_obj" not in follow_up
assert "litellm_call_id" not in follow_up["litellm_params"]
assert "litellm_logging_obj" not in follow_up["litellm_params"]
assert follow_up["litellm_trace_id"] == "trace-1"
assert follow_up["stream"] is True
assert "tool_choice" not in follow_up
assert follow_up["litellm_params"]["metadata"] == {"team": "legal"}
assert original["litellm_call_id"] == "parent-call"