From 03ca9fa9320776757156761d9ffa561d224f4882 Mon Sep 17 00:00:00 2001 From: tarunag10 Date: Thu, 20 Aug 2026 12:45:44 +0530 Subject: [PATCH 1/9] fix: isolate chained MCP spend logs --- .../responses/mcp/chat_completions_handler.py | 8 +++- .../mcp/litellm_proxy_mcp_handler.py | 30 ++++++++++++- .../responses/mcp/mcp_streaming_iterator.py | 4 +- .../mcp/test_litellm_proxy_mcp_handler.py | 43 ++++++++++++++++--- 4 files changed, 74 insertions(+), 11 deletions(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 2a0406f9a4d..bf948df52c2 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -453,7 +453,9 @@ async def acompletion_with_mcp( ) # Make follow-up call with streaming - follow_up_call_args: Final = dict(self.base_call_args) + follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params( + self.base_call_args + ) 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 +627,9 @@ async def acompletion_with_mcp( ) # Make follow-up call with original stream setting - follow_up_call_args: Final = dict(base_call_args) + follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params( + base_call_args + ) follow_up_call_args["messages"] = follow_up_messages follow_up_call_args["stream"] = stream diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 197d0c02ba8..46bb067eeeb 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -71,6 +71,28 @@ class LiteLLM_Proxy_MCP_Handler: This handles when a user passes mcp server_url="litellm_proxy" in their tools. """ + @staticmethod + def _prepare_follow_up_call_params(params: Mapping[str, Any]) -> dict[str, Any]: + """Copy request params without state owned by the previous LLM call. + + MCP auto-execution keeps the trace identifier so chained rounds remain + correlated, but each provider call must create its own logging object and + call identifier. Reusing either makes success dispatch one-shot and drops + spend rows for later rounds. + """ + follow_up_params = dict(params) + follow_up_params.pop("litellm_logging_obj", None) + follow_up_params.pop("litellm_call_id", None) + + nested_params = follow_up_params.get("litellm_params") + if isinstance(nested_params, dict): + nested_params = dict(nested_params) + nested_params.pop("litellm_logging_obj", None) + nested_params.pop("litellm_call_id", None) + follow_up_params["litellm_params"] = nested_params + + return follow_up_params + @staticmethod def _get_parent_request_tags(kwargs: dict[str, Any] | None) -> list[str]: """Tags from the parent LLM request, using the same extraction logic as standard logging (incl. User-Agent).""" @@ -702,13 +724,19 @@ class LiteLLM_Proxy_MCP_Handler: }, } ] - tool_logging_call_id = litellm_call_id or str(uuid.uuid4()) + # SpendLogs.request_id is unique. The parent LLM call ID is + # therefore metadata, not the tool execution's call ID: reusing + # it causes all but the first MCP tool row to be skipped by the + # database's duplicate protection. + tool_logging_call_id = str(uuid.uuid4()) logging_metadata: dict[str, object] = { "tool_call_id": tool_call_id, "tool_name": sanitized_tool_name, "server_name": server_name, "headers": logging_safe_headers, } + if litellm_call_id: + logging_metadata["parent_litellm_call_id"] = litellm_call_id logging_request_data = { "model": f"MCP: {tool_name}", "metadata": logging_metadata, diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index c7471518398..69e698f7f65 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -782,7 +782,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): ) # Make follow-up call with streaming - follow_up_params: Final = self.original_request_params.copy() + follow_up_params: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params( + self.original_request_params + ) follow_up_params.update( { "input": follow_up_input, diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index b33ed3bb581..a1039a20489 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -410,7 +410,7 @@ async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkey @pytest.mark.asyncio -async def test_execute_tool_calls_passes_litellm_call_id_and_trace_id_to_function_setup( +async def test_execute_tool_calls_uses_unique_call_ids_and_preserves_parent_context( monkeypatch, ): """ @@ -420,10 +420,10 @@ async def test_execute_tool_calls_passes_litellm_call_id_and_trace_id_to_functio _setup_proxy_logging(monkeypatch) call_tool_mock = _setup_mcp_call_environment(monkeypatch) - captured = {} + captured = [] def fake_function_setup(*_args, **kwargs): - captured.update(kwargs) + captured.append(kwargs) return None, None # NOTE: Don't patch via dotted string path here because `litellm.responses` @@ -435,7 +435,10 @@ async def test_execute_tool_calls_passes_litellm_call_id_and_trace_id_to_functio monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) tool_name = "deepwiki-read_wiki_structure" - tool_calls = [{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}] + tool_calls = [ + {"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}, + {"id": "call-2", "function": {"name": tool_name, "arguments": "{}"}}, + ] await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( tool_server_map={tool_name: "deepwiki"}, @@ -446,10 +449,36 @@ async def test_execute_tool_calls_passes_litellm_call_id_and_trace_id_to_functio ) # Ensure the tool call was attempted (sanity) - assert call_tool_mock.await_count == 1 + assert call_tool_mock.await_count == 2 + assert len(captured) == 2 + assert captured[0]["litellm_call_id"] != captured[1]["litellm_call_id"] + assert all(item["litellm_call_id"] != "cid" for item in captured) + assert all(item["litellm_trace_id"] == "tid" for item in captured) + assert all(item["metadata"]["parent_litellm_call_id"] == "cid" for item in captured) - assert captured.get("litellm_call_id") == "cid" - assert captured.get("litellm_trace_id") == "tid" + +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", + "litellm_params": { + "litellm_call_id": "nested-parent-call", + "litellm_logging_obj": object(), + "metadata": {"team": "legal"}, + }, + } + + follow_up = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params(original) + + 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["litellm_params"]["metadata"] == {"team": "legal"} + assert original["litellm_call_id"] == "parent-call" @pytest.mark.asyncio From 160f0d440c45641b7c03a7415e77613e5268e5ee Mon Sep 17 00:00:00 2001 From: tarunag10 Date: Thu, 20 Aug 2026 12:50:05 +0530 Subject: [PATCH 2/9] style: format MCP follow-up calls --- .../responses/mcp/chat_completions_handler.py | 136 ++++++++++++------ 1 file changed, 95 insertions(+), 41 deletions(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index bf948df52c2..7df0c34aa47 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -138,7 +138,9 @@ async def acompletion_with_mcp( ) # Combine with other tools - all_tools: Final = openai_tools + other_tools if (openai_tools or other_tools) else None + all_tools: Final = ( + openai_tools + other_tools if (openai_tools or other_tools) else None + ) # Determine if we should auto-execute tools should_auto_execute: Final = LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools( @@ -239,7 +241,9 @@ async def acompletion_with_mcp( def __aiter__(self): return self - def _add_mcp_list_tools_to_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: + def _add_mcp_list_tools_to_chunk( + self, chunk: ModelResponseStream + ) -> ModelResponseStream: """Add mcp_list_tools to the first chunk.""" from litellm.types.utils import ( StreamingChoices, @@ -251,10 +255,19 @@ async def acompletion_with_mcp( if hasattr(chunk, "choices") and chunk.choices: for choice in chunk.choices: - if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: + if ( + isinstance(choice, StreamingChoices) + and hasattr(choice, "delta") + and choice.delta + ): # Get existing provider_specific_fields or create new dict - existing_fields = getattr(choice.delta, "provider_specific_fields", None) or {} - provider_fields = dict(existing_fields) # Create a copy to avoid mutating the original + existing_fields = ( + getattr(choice.delta, "provider_specific_fields", None) + or {} + ) + provider_fields = dict( + existing_fields + ) # Create a copy to avoid mutating the original # Add only mcp_list_tools to first chunk provider_fields["mcp_list_tools"] = self.openai_tools @@ -265,7 +278,9 @@ async def acompletion_with_mcp( return chunk - def _add_mcp_tool_metadata_to_final_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: + def _add_mcp_tool_metadata_to_final_chunk( + self, chunk: ModelResponseStream + ) -> ModelResponseStream: """Add mcp_tool_calls and mcp_call_results to the final chunk.""" from litellm.types.utils import ( StreamingChoices, @@ -274,15 +289,25 @@ async def acompletion_with_mcp( if hasattr(chunk, "choices") and chunk.choices: for choice in chunk.choices: - if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: + if ( + isinstance(choice, StreamingChoices) + and hasattr(choice, "delta") + and choice.delta + ): # Get existing provider_specific_fields or create new dict # Access the attribute directly to handle Pydantic model attributes correctly existing_fields = {} if hasattr(choice.delta, "provider_specific_fields"): - attr_value = getattr(choice.delta, "provider_specific_fields", None) + attr_value = getattr( + choice.delta, "provider_specific_fields", None + ) if attr_value is not None: # Create a copy to avoid mutating the original - existing_fields = dict(attr_value) if isinstance(attr_value, dict) else {} + existing_fields = ( + dict(attr_value) + if isinstance(attr_value, dict) + else {} + ) provider_fields = existing_fields @@ -357,7 +382,9 @@ async def acompletion_with_mcp( # If we have chunks, yield the final one with metadata if self.collected_chunks: final_chunk = self.collected_chunks[-1] - final_chunk = self._add_mcp_tool_metadata_to_final_chunk(final_chunk) + final_chunk = self._add_mcp_tool_metadata_to_final_chunk( + final_chunk + ) # If we have tool results, prepare follow-up call if self.tool_results and self.complete_response: await self._prepare_follow_up_call() @@ -395,7 +422,9 @@ async def acompletion_with_mcp( ): from litellm._logging import verbose_logger - verbose_logger.warning("Follow-up stream was not created despite having tool results") + verbose_logger.warning( + "Follow-up stream was not created despite having tool results" + ) raise StopAsyncIteration @@ -424,17 +453,19 @@ async def acompletion_with_mcp( if self.tool_calls: # Execute tool calls - self.tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( - tool_server_map=self.tool_server_map, - tool_calls=self.tool_calls, - user_api_key_auth=self.user_api_key_auth, - mcp_auth_header=self.mcp_auth_header, - mcp_server_auth_headers=self.mcp_server_auth_headers, - oauth2_headers=self.oauth2_headers, - raw_headers=self.raw_headers, - litellm_call_id=self.litellm_call_id, - litellm_trace_id=self.litellm_trace_id, - request_tags=self.request_tags, + self.tool_results = ( + await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map=self.tool_server_map, + tool_calls=self.tool_calls, + user_api_key_auth=self.user_api_key_auth, + mcp_auth_header=self.mcp_auth_header, + mcp_server_auth_headers=self.mcp_server_auth_headers, + oauth2_headers=self.oauth2_headers, + raw_headers=self.raw_headers, + litellm_call_id=self.litellm_call_id, + litellm_trace_id=self.litellm_trace_id, + request_tags=self.request_tags, + ) ) async def _prepare_follow_up_call(self): @@ -446,15 +477,19 @@ async def acompletion_with_mcp( return # Create follow-up messages with tool results - follow_up_messages: Final = LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( - original_messages=self.messages, - response=self.complete_response, - tool_results=self.tool_results, + follow_up_messages: Final = ( + LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( + original_messages=self.messages, + response=self.complete_response, + tool_results=self.tool_results, + ) ) # Make follow-up call with streaming - follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params( - self.base_call_args + follow_up_call_args: Final = ( + LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params( + self.base_call_args + ) ) follow_up_call_args["messages"] = follow_up_messages follow_up_call_args["stream"] = True @@ -465,7 +500,9 @@ async def acompletion_with_mcp( # This ensures the patch works correctly in tests import litellm - follow_up_response: Final = await litellm.acompletion(**follow_up_call_args) + follow_up_response: Final = await litellm.acompletion( + **follow_up_call_args + ) # Ensure follow-up response is a CustomStreamWrapper if isinstance(follow_up_response, CustomStreamWrapper): @@ -478,7 +515,8 @@ async def acompletion_with_mcp( from litellm._logging import verbose_logger verbose_logger.warning( - "Follow-up response is not a CustomStreamWrapper: %s", type(follow_up_response) + "Follow-up response is not a CustomStreamWrapper: %s", + type(follow_up_response), ) self.follow_up_stream = None @@ -502,16 +540,24 @@ async def acompletion_with_mcp( # Create a wrapper class that delegates to our custom iterator # We'll use a simple approach: just replace the __aiter__ method class MCPStreamWrapper(CustomStreamWrapper): - def __init__(self, original_wrapper: CustomStreamWrapper, custom_iterator: MCPStreamingIterator): + def __init__( + self, + original_wrapper: CustomStreamWrapper, + custom_iterator: MCPStreamingIterator, + ): # Initialize with the same parameters as original wrapper super().__init__( completion_stream=None, model=getattr(original_wrapper, "model", "unknown"), logging_obj=original_wrapper.logging_obj, - custom_llm_provider=getattr(original_wrapper, "custom_llm_provider", None), + custom_llm_provider=getattr( + original_wrapper, "custom_llm_provider", None + ), stream_options=getattr(original_wrapper, "stream_options", None), make_call=getattr(original_wrapper, "make_call", None), - _response_headers=getattr(original_wrapper, "_response_headers", None), + _response_headers=getattr( + original_wrapper, "_response_headers", None + ), ) self._original_wrapper = original_wrapper self._custom_iterator = custom_iterator @@ -535,7 +581,9 @@ async def acompletion_with_mcp( except RuntimeError: self._sync_loop = asyncio.new_event_loop() asyncio.set_event_loop(self._sync_loop) - self._sync_iterator = _SyncIteratorWrapper(self._custom_iterator, self._sync_loop) + self._sync_iterator = _SyncIteratorWrapper( + self._custom_iterator, self._sync_loop + ) return self._sync_iterator def __next__(self): @@ -588,7 +636,11 @@ async def acompletion_with_mcp( return initial_response # Extract tool calls from response - tool_calls: Final = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response(response=initial_response) + tool_calls: Final = ( + LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response( + response=initial_response + ) + ) if not tool_calls: _add_mcp_metadata_to_response( @@ -620,15 +672,17 @@ async def acompletion_with_mcp( return initial_response # Create follow-up messages with tool results - follow_up_messages: Final = LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( - original_messages=messages, - response=initial_response, - tool_results=tool_results, + follow_up_messages: Final = ( + LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( + original_messages=messages, + response=initial_response, + tool_results=tool_results, + ) ) # Make follow-up call with original stream setting - follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params( - base_call_args + follow_up_call_args: Final = ( + LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params(base_call_args) ) follow_up_call_args["messages"] = follow_up_messages follow_up_call_args["stream"] = stream From 9597c6c102155bcd2ef023dae258ad898335d943 Mon Sep 17 00:00:00 2001 From: tarunag10 Date: Thu, 20 Aug 2026 12:51:55 +0530 Subject: [PATCH 3/9] style: apply repository lint rules --- .../responses/mcp/chat_completions_handler.py | 143 ++++++------------ .../mcp/litellm_proxy_mcp_handler.py | 12 +- .../responses/mcp/mcp_streaming_iterator.py | 6 +- .../mcp/test_litellm_proxy_mcp_handler.py | 57 +++---- 4 files changed, 69 insertions(+), 149 deletions(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 7df0c34aa47..2b3c0d79192 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -3,13 +3,14 @@ import logging from typing import TYPE_CHECKING, Any, Final, cast -from litellm.responses.mcp.litellm_proxy_mcp_handler import ( - LiteLLM_Proxy_MCP_Handler, -) from litellm.responses.mcp.request_context import MCPRequestContext from litellm.types.utils import Message, ModelResponse from litellm.utils import CustomStreamWrapper +from litellm.responses.mcp.litellm_proxy_mcp_handler import ( + LiteLLM_Proxy_MCP_Handler, +) + if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth @@ -138,9 +139,7 @@ async def acompletion_with_mcp( ) # Combine with other tools - all_tools: Final = ( - openai_tools + other_tools if (openai_tools or other_tools) else None - ) + all_tools: Final = openai_tools + other_tools if (openai_tools or other_tools) else None # Determine if we should auto-execute tools should_auto_execute: Final = LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools( @@ -241,9 +240,7 @@ async def acompletion_with_mcp( def __aiter__(self): return self - def _add_mcp_list_tools_to_chunk( - self, chunk: ModelResponseStream - ) -> ModelResponseStream: + def _add_mcp_list_tools_to_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: """Add mcp_list_tools to the first chunk.""" from litellm.types.utils import ( StreamingChoices, @@ -255,19 +252,10 @@ async def acompletion_with_mcp( if hasattr(chunk, "choices") and chunk.choices: for choice in chunk.choices: - if ( - isinstance(choice, StreamingChoices) - and hasattr(choice, "delta") - and choice.delta - ): + if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: # Get existing provider_specific_fields or create new dict - existing_fields = ( - getattr(choice.delta, "provider_specific_fields", None) - or {} - ) - provider_fields = dict( - existing_fields - ) # Create a copy to avoid mutating the original + existing_fields = getattr(choice.delta, "provider_specific_fields", None) or {} + provider_fields = dict(existing_fields) # Create a copy to avoid mutating the original # Add only mcp_list_tools to first chunk provider_fields["mcp_list_tools"] = self.openai_tools @@ -278,9 +266,7 @@ async def acompletion_with_mcp( return chunk - def _add_mcp_tool_metadata_to_final_chunk( - self, chunk: ModelResponseStream - ) -> ModelResponseStream: + def _add_mcp_tool_metadata_to_final_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: """Add mcp_tool_calls and mcp_call_results to the final chunk.""" from litellm.types.utils import ( StreamingChoices, @@ -289,25 +275,15 @@ async def acompletion_with_mcp( if hasattr(chunk, "choices") and chunk.choices: for choice in chunk.choices: - if ( - isinstance(choice, StreamingChoices) - and hasattr(choice, "delta") - and choice.delta - ): + if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: # Get existing provider_specific_fields or create new dict # Access the attribute directly to handle Pydantic model attributes correctly existing_fields = {} if hasattr(choice.delta, "provider_specific_fields"): - attr_value = getattr( - choice.delta, "provider_specific_fields", None - ) + attr_value = getattr(choice.delta, "provider_specific_fields", None) if attr_value is not None: # Create a copy to avoid mutating the original - existing_fields = ( - dict(attr_value) - if isinstance(attr_value, dict) - else {} - ) + existing_fields = dict(attr_value) if isinstance(attr_value, dict) else {} provider_fields = existing_fields @@ -382,9 +358,7 @@ async def acompletion_with_mcp( # If we have chunks, yield the final one with metadata if self.collected_chunks: final_chunk = self.collected_chunks[-1] - final_chunk = self._add_mcp_tool_metadata_to_final_chunk( - final_chunk - ) + final_chunk = self._add_mcp_tool_metadata_to_final_chunk(final_chunk) # If we have tool results, prepare follow-up call if self.tool_results and self.complete_response: await self._prepare_follow_up_call() @@ -422,9 +396,7 @@ async def acompletion_with_mcp( ): from litellm._logging import verbose_logger - verbose_logger.warning( - "Follow-up stream was not created despite having tool results" - ) + verbose_logger.warning("Follow-up stream was not created despite having tool results") raise StopAsyncIteration @@ -453,19 +425,17 @@ async def acompletion_with_mcp( if self.tool_calls: # Execute tool calls - self.tool_results = ( - await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( - tool_server_map=self.tool_server_map, - tool_calls=self.tool_calls, - user_api_key_auth=self.user_api_key_auth, - mcp_auth_header=self.mcp_auth_header, - mcp_server_auth_headers=self.mcp_server_auth_headers, - oauth2_headers=self.oauth2_headers, - raw_headers=self.raw_headers, - litellm_call_id=self.litellm_call_id, - litellm_trace_id=self.litellm_trace_id, - request_tags=self.request_tags, - ) + self.tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map=self.tool_server_map, + tool_calls=self.tool_calls, + user_api_key_auth=self.user_api_key_auth, + mcp_auth_header=self.mcp_auth_header, + mcp_server_auth_headers=self.mcp_server_auth_headers, + oauth2_headers=self.oauth2_headers, + raw_headers=self.raw_headers, + litellm_call_id=self.litellm_call_id, + litellm_trace_id=self.litellm_trace_id, + request_tags=self.request_tags, ) async def _prepare_follow_up_call(self): @@ -477,19 +447,15 @@ async def acompletion_with_mcp( return # Create follow-up messages with tool results - follow_up_messages: Final = ( - LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( - original_messages=self.messages, - response=self.complete_response, - tool_results=self.tool_results, - ) + follow_up_messages: Final = LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( + original_messages=self.messages, + response=self.complete_response, + tool_results=self.tool_results, ) # Make follow-up call with streaming - follow_up_call_args: Final = ( - LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params( - self.base_call_args - ) + follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params( + self.base_call_args ) follow_up_call_args["messages"] = follow_up_messages follow_up_call_args["stream"] = True @@ -500,9 +466,7 @@ async def acompletion_with_mcp( # This ensures the patch works correctly in tests import litellm - follow_up_response: Final = await litellm.acompletion( - **follow_up_call_args - ) + follow_up_response: Final = await litellm.acompletion(**follow_up_call_args) # Ensure follow-up response is a CustomStreamWrapper if isinstance(follow_up_response, CustomStreamWrapper): @@ -515,8 +479,7 @@ async def acompletion_with_mcp( from litellm._logging import verbose_logger verbose_logger.warning( - "Follow-up response is not a CustomStreamWrapper: %s", - type(follow_up_response), + "Follow-up response is not a CustomStreamWrapper: %s", type(follow_up_response) ) self.follow_up_stream = None @@ -540,24 +503,16 @@ async def acompletion_with_mcp( # Create a wrapper class that delegates to our custom iterator # We'll use a simple approach: just replace the __aiter__ method class MCPStreamWrapper(CustomStreamWrapper): - def __init__( - self, - original_wrapper: CustomStreamWrapper, - custom_iterator: MCPStreamingIterator, - ): + def __init__(self, original_wrapper: CustomStreamWrapper, custom_iterator: MCPStreamingIterator): # Initialize with the same parameters as original wrapper super().__init__( completion_stream=None, model=getattr(original_wrapper, "model", "unknown"), logging_obj=original_wrapper.logging_obj, - custom_llm_provider=getattr( - original_wrapper, "custom_llm_provider", None - ), + custom_llm_provider=getattr(original_wrapper, "custom_llm_provider", None), stream_options=getattr(original_wrapper, "stream_options", None), make_call=getattr(original_wrapper, "make_call", None), - _response_headers=getattr( - original_wrapper, "_response_headers", None - ), + _response_headers=getattr(original_wrapper, "_response_headers", None), ) self._original_wrapper = original_wrapper self._custom_iterator = custom_iterator @@ -581,9 +536,7 @@ async def acompletion_with_mcp( except RuntimeError: self._sync_loop = asyncio.new_event_loop() asyncio.set_event_loop(self._sync_loop) - self._sync_iterator = _SyncIteratorWrapper( - self._custom_iterator, self._sync_loop - ) + self._sync_iterator = _SyncIteratorWrapper(self._custom_iterator, self._sync_loop) return self._sync_iterator def __next__(self): @@ -636,11 +589,7 @@ async def acompletion_with_mcp( return initial_response # Extract tool calls from response - tool_calls: Final = ( - LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response( - response=initial_response - ) - ) + tool_calls: Final = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response(response=initial_response) if not tool_calls: _add_mcp_metadata_to_response( @@ -672,18 +621,14 @@ async def acompletion_with_mcp( return initial_response # Create follow-up messages with tool results - follow_up_messages: Final = ( - LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( - original_messages=messages, - response=initial_response, - tool_results=tool_results, - ) + follow_up_messages: Final = LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( + original_messages=messages, + response=initial_response, + tool_results=tool_results, ) # Make follow-up call with original stream setting - follow_up_call_args: Final = ( - LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params(base_call_args) - ) + follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params(base_call_args) follow_up_call_args["messages"] = follow_up_messages follow_up_call_args["stream"] = stream diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 46bb067eeeb..d2e9c8f37ba 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -4,9 +4,6 @@ from collections.abc import Iterable, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload -from openai.types.chat import ChatCompletionToolParam -from openai.types.responses.function_tool_param import FunctionToolParam - from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -33,13 +30,14 @@ from litellm.types.utils import ( StandardLoggingMCPToolCall, ) from litellm.utils import Rules, function_setup +from openai.types.chat import ChatCompletionToolParam +from openai.types.responses.function_tool_param import FunctionToolParam if TYPE_CHECKING: - from mcp.types import CallToolResult - from mcp.types import Tool as MCPTool - from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging + from mcp.types import CallToolResult + from mcp.types import Tool as MCPTool else: MCPTool = Any @@ -669,7 +667,6 @@ class LiteLLM_Proxy_MCP_Handler: ) -> list[MCPToolResult]: """Execute tool calls and return results.""" from fastapi import HTTPException - from litellm._uuid import uuid from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -1205,6 +1202,7 @@ class LiteLLM_Proxy_MCP_Handler: List of MCP tool execution events for streaming """ from litellm._uuid import uuid + from litellm.responses.mcp.mcp_streaming_iterator import create_mcp_call_events tool_execution_events: Final[list[ResponsesAPIStreamingResponse]] = [] diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 69e698f7f65..89774568d57 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -22,9 +22,9 @@ from litellm.types.llms.openai import ( ) if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth from mcp.types import Tool as MCPTool - from litellm.proxy._types import UserAPIKeyAuth from litellm.responses.mcp.litellm_proxy_mcp_handler import MCPToolResult else: MCPTool = Any @@ -331,11 +331,10 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): def _extract_mcp_headers_from_params(self) -> None: """Extract MCP headers from original request params to pass to tool calls""" - from starlette.datastructures import Headers - from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) + from starlette.datastructures import Headers # Extract headers from secret_fields in original_request_params raw_headers_from_request: dict[str, str] | None = None @@ -768,6 +767,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): return from litellm.responses.main import aresponses + from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index a1039a20489..429adb2962d 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -1,20 +1,20 @@ +import importlib import subprocess import sys import textwrap import types +from typing import Any, cast from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException -import importlib - from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing +from litellm.types.responses.main import OutputFunctionToolCall +from litellm.types.utils import ModelResponse + from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) -from typing import Any, cast -from litellm.types.utils import ModelResponse -from litellm.types.responses.main import OutputFunctionToolCall class _DummyMCPResult: @@ -105,9 +105,7 @@ def test_extract_tool_calls_from_chat_response_handles_tool_calls(): object="chat.completion", ) - tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response( - response - ) + tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response(response) assert len(tool_calls) == 1 assert tool_calls[0]["function"]["name"] == "foo" @@ -177,9 +175,7 @@ def test_transform_mcp_tools_to_openai_uses_chat_format(monkeypatch): fake_transform_responses, ) - chat_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai( - ["tool"], target_format="chat" - ) + chat_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai(["tool"], target_format="chat") resp_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai(["tool"]) assert chat_tools == [{"chat": True}] @@ -299,9 +295,7 @@ async def test_execute_tool_calls_strips_prefix_when_alias_differs_from_server_n ) from litellm.proxy._experimental.mcp_server import mcp_server_manager as _msm - _msm.global_mcp_server_manager._get_mcp_server_from_tool_name = MagicMock( - return_value=fake_server - ) + _msm.global_mcp_server_manager._get_mcp_server_from_tool_name = MagicMock(return_value=fake_server) tool_name = "my_deepwiki-read_wiki_structure" tool_calls = [ @@ -375,7 +369,7 @@ async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkey fake_manager = types.SimpleNamespace( get_registry=MagicMock(return_value={}), - call_tool=AsyncMock(side_effect=HTTPException(status_code=500, detail="boom")) + call_tool=AsyncMock(side_effect=HTTPException(status_code=500, detail="boom")), ) monkeypatch.setattr( "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", @@ -383,9 +377,7 @@ async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkey ) tool_name = "deepwiki-read_wiki_structure" - tool_calls = [ - {"id": "call-err", "function": {"name": tool_name, "arguments": "{}"}} - ] + tool_calls = [{"id": "call-err", "function": {"name": tool_name, "arguments": "{}"}}] user_auth = types.SimpleNamespace(api_key="test_key", user_id="test_user") @@ -403,10 +395,7 @@ async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkey post_call_failure_hook.assert_awaited_once() assert post_call_failure_hook.await_args is not None - assert ( - post_call_failure_hook.await_args.kwargs.get("route") - == "/responses/mcp/call_tool" - ) + assert post_call_failure_hook.await_args.kwargs.get("route") == "/responses/mcp/call_tool" @pytest.mark.asyncio @@ -429,9 +418,7 @@ async def test_execute_tool_calls_uses_unique_call_ids_and_preserves_parent_cont # NOTE: Don't patch via dotted string path here because `litellm.responses` # is a function attribute on the `litellm` package (shadowing the submodule), # which breaks monkeypatch's importpath resolution. - handler_module = importlib.import_module( - "litellm.responses.mcp.litellm_proxy_mcp_handler" - ) + handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) tool_name = "deepwiki-read_wiki_structure" @@ -544,9 +531,7 @@ async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch user_auth = types.SimpleNamespace(api_key="test_key", user_id="test_user") tools, _server_names = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( user_api_key_auth=user_auth, - mcp_tools_with_litellm_proxy=[ - {"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"} - ], + mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"}], ) assert tools == [] @@ -557,9 +542,7 @@ async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch def test_get_parent_request_tags_from_metadata(): - tags = LiteLLM_Proxy_MCP_Handler._get_parent_request_tags( - {"metadata": {"tags": ["team-a", "prod"]}} - ) + tags = LiteLLM_Proxy_MCP_Handler._get_parent_request_tags({"metadata": {"tags": ["team-a", "prod"]}}) assert tags == ["team-a", "prod"] @@ -596,9 +579,7 @@ async def test_get_mcp_tools_from_manager_forwards_request_tags(monkeypatch): await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( user_api_key_auth=types.SimpleNamespace(api_key="k", user_id="u"), - mcp_tools_with_litellm_proxy=[ - {"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"} - ], + mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"}], request_tags=["team-a"], ) @@ -618,9 +599,7 @@ async def test_execute_tool_calls_exposes_sanitized_client_headers_to_logging(mo captured.update(kwargs) return None, None - handler_module = importlib.import_module( - "litellm.responses.mcp.litellm_proxy_mcp_handler" - ) + handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) tool_name = "deepwiki-read_wiki_structure" @@ -646,9 +625,7 @@ async def test_execute_tool_calls_propagates_request_tags_to_function_setup(monk captured.update(kwargs) return None, None - handler_module = importlib.import_module( - "litellm.responses.mcp.litellm_proxy_mcp_handler" - ) + handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) tool_name = "deepwiki-read_wiki_structure" From 8e57d4d457c8b7dc6bc97646c65864ea3d593857 Mon Sep 17 00:00:00 2001 From: tarunag10 Date: Thu, 20 Aug 2026 12:57:19 +0530 Subject: [PATCH 4/9] style: apply repository Ruff import ordering --- litellm/responses/mcp/chat_completions_handler.py | 7 +++---- litellm/responses/mcp/litellm_proxy_mcp_handler.py | 1 - litellm/responses/mcp/mcp_streaming_iterator.py | 4 +--- 3 files changed, 4 insertions(+), 8 deletions(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 2b3c0d79192..73dba5ba441 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -3,13 +3,12 @@ import logging from typing import TYPE_CHECKING, Any, Final, cast -from litellm.responses.mcp.request_context import MCPRequestContext -from litellm.types.utils import Message, ModelResponse -from litellm.utils import CustomStreamWrapper - from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) +from litellm.responses.mcp.request_context import MCPRequestContext +from litellm.types.utils import Message, ModelResponse +from litellm.utils import CustomStreamWrapper if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index d2e9c8f37ba..6254e60bf20 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -1202,7 +1202,6 @@ class LiteLLM_Proxy_MCP_Handler: List of MCP tool execution events for streaming """ from litellm._uuid import uuid - from litellm.responses.mcp.mcp_streaming_iterator import create_mcp_call_events tool_execution_events: Final[list[ResponsesAPIStreamingResponse]] = [] diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 89774568d57..8534e13aa06 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -23,9 +23,8 @@ from litellm.types.llms.openai import ( if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth - from mcp.types import Tool as MCPTool - from litellm.responses.mcp.litellm_proxy_mcp_handler import MCPToolResult + from mcp.types import Tool as MCPTool else: MCPTool = Any @@ -767,7 +766,6 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): return from litellm.responses.main import aresponses - from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) From cf8a18ee7b86306dcd16469a366ebe0fe32115d2 Mon Sep 17 00:00:00 2001 From: tarunag10 Date: Thu, 20 Aug 2026 13:02:19 +0530 Subject: [PATCH 5/9] style: match full-package Ruff import classification --- litellm/responses/mcp/litellm_proxy_mcp_handler.py | 11 +++++++---- litellm/responses/mcp/mcp_streaming_iterator.py | 6 ++++-- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 6254e60bf20..46bb067eeeb 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -4,6 +4,9 @@ from collections.abc import Iterable, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload +from openai.types.chat import ChatCompletionToolParam +from openai.types.responses.function_tool_param import FunctionToolParam + from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -30,14 +33,13 @@ from litellm.types.utils import ( StandardLoggingMCPToolCall, ) from litellm.utils import Rules, function_setup -from openai.types.chat import ChatCompletionToolParam -from openai.types.responses.function_tool_param import FunctionToolParam if TYPE_CHECKING: - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.utils import ProxyLogging from mcp.types import CallToolResult from mcp.types import Tool as MCPTool + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging else: MCPTool = Any @@ -667,6 +669,7 @@ class LiteLLM_Proxy_MCP_Handler: ) -> list[MCPToolResult]: """Execute tool calls and return results.""" from fastapi import HTTPException + from litellm._uuid import uuid from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 8534e13aa06..69e698f7f65 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -22,9 +22,10 @@ from litellm.types.llms.openai import ( ) if TYPE_CHECKING: + from mcp.types import Tool as MCPTool + from litellm.proxy._types import UserAPIKeyAuth from litellm.responses.mcp.litellm_proxy_mcp_handler import MCPToolResult - from mcp.types import Tool as MCPTool else: MCPTool = Any @@ -330,10 +331,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): def _extract_mcp_headers_from_params(self) -> None: """Extract MCP headers from original request params to pass to tool calls""" + from starlette.datastructures import Headers + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) - from starlette.datastructures import Headers # Extract headers from secret_fields in original_request_params raw_headers_from_request: dict[str, str] | None = None From 05d4996a410cb6750656e246a3f8428a4d28954d Mon Sep 17 00:00:00 2001 From: tarunag10 Date: Thu, 20 Aug 2026 13:07:52 +0530 Subject: [PATCH 6/9] fix: avoid MCP follow-up helper collision --- litellm/responses/mcp/chat_completions_handler.py | 6 ++---- litellm/responses/mcp/litellm_proxy_mcp_handler.py | 2 +- litellm/responses/mcp/mcp_streaming_iterator.py | 2 +- .../responses/mcp/test_litellm_proxy_mcp_handler.py | 4 ++-- 4 files changed, 6 insertions(+), 8 deletions(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 73dba5ba441..b914bb1e4aa 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -453,9 +453,7 @@ async def acompletion_with_mcp( ) # Make follow-up call with streaming - follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params( - self.base_call_args - ) + follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_chained_call_params(self.base_call_args) follow_up_call_args["messages"] = follow_up_messages follow_up_call_args["stream"] = True # Ensure follow-up call doesn't trigger MCP handler again @@ -627,7 +625,7 @@ async def acompletion_with_mcp( ) # Make follow-up call with original stream setting - follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params(base_call_args) + follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_chained_call_params(base_call_args) follow_up_call_args["messages"] = follow_up_messages follow_up_call_args["stream"] = stream diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 46bb067eeeb..d72f335e89c 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -72,7 +72,7 @@ class LiteLLM_Proxy_MCP_Handler: """ @staticmethod - def _prepare_follow_up_call_params(params: Mapping[str, Any]) -> dict[str, Any]: + def _prepare_chained_call_params(params: Mapping[str, Any]) -> dict[str, Any]: """Copy request params without state owned by the previous LLM call. MCP auto-execution keeps the trace identifier so chained rounds remain diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 69e698f7f65..1e1cbd34ce2 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -782,7 +782,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): ) # Make follow-up call with streaming - follow_up_params: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params( + follow_up_params: Final = LiteLLM_Proxy_MCP_Handler._prepare_chained_call_params( self.original_request_params ) follow_up_params.update( diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index 429adb2962d..c552170a599 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -444,7 +444,7 @@ 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_follow_up_call_params_resets_per_call_logging_state(): +def test_prepare_chained_call_params_resets_per_call_logging_state(): original = { "model": "gpt-4", "litellm_call_id": "parent-call", @@ -457,7 +457,7 @@ def test_prepare_follow_up_call_params_resets_per_call_logging_state(): }, } - follow_up = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params(original) + follow_up = LiteLLM_Proxy_MCP_Handler._prepare_chained_call_params(original) assert "litellm_call_id" not in follow_up assert "litellm_logging_obj" not in follow_up From e5ffc4f8db3019c05857bb67888c4575a72f71a7 Mon Sep 17 00:00:00 2001 From: tarunag10 Date: Thu, 20 Aug 2026 13:14:43 +0530 Subject: [PATCH 7/9] fix: preserve chat follow-up normalization --- litellm/responses/mcp/chat_completions_handler.py | 10 ++++++++-- litellm/responses/mcp/litellm_proxy_mcp_handler.py | 6 ++++-- .../responses/mcp/test_litellm_proxy_mcp_handler.py | 7 +++++-- 3 files changed, 17 insertions(+), 6 deletions(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index b914bb1e4aa..764fe003f2c 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -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 diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index d72f335e89c..07e192e68b4 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -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 diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index c552170a599..de889421e1c 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -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" From 887d8c751d9e02154c8faf59cff42e6319f707ce Mon Sep 17 00:00:00 2001 From: tarunag10 Date: Thu, 20 Aug 2026 13:20:06 +0530 Subject: [PATCH 8/9] fix: omit null web search options on chained calls --- litellm/responses/mcp/litellm_proxy_mcp_handler.py | 2 ++ .../responses/mcp/test_litellm_proxy_mcp_handler.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 07e192e68b4..f13bcece7ca 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -85,6 +85,8 @@ class LiteLLM_Proxy_MCP_Handler: follow_up_params = dict(params) follow_up_params.pop("litellm_logging_obj", None) follow_up_params.pop("litellm_call_id", None) + if follow_up_params.get("web_search_options") is None: + follow_up_params.pop("web_search_options", None) nested_params = follow_up_params.get("litellm_params") if isinstance(nested_params, dict): diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index de889421e1c..f2658a1b0a5 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -451,6 +451,7 @@ def test_prepare_follow_up_call_params_resets_per_call_logging_state(): "litellm_logging_obj": object(), "litellm_trace_id": "trace-1", "tool_choice": "auto", + "web_search_options": None, "litellm_params": { "litellm_call_id": "nested-parent-call", "litellm_logging_obj": object(), @@ -467,6 +468,7 @@ def test_prepare_follow_up_call_params_resets_per_call_logging_state(): assert follow_up["litellm_trace_id"] == "trace-1" assert follow_up["stream"] is True assert "tool_choice" not in follow_up + assert "web_search_options" not in follow_up assert follow_up["litellm_params"]["metadata"] == {"team": "legal"} assert original["litellm_call_id"] == "parent-call" From 7768fe9b0099bef182edad20c5103a60f2cfc1e0 Mon Sep 17 00:00:00 2001 From: tarunag10 Date: Thu, 20 Aug 2026 13:28:50 +0530 Subject: [PATCH 9/9] refactor: expose MCP chained-call helpers --- litellm/responses/mcp/chat_completions_handler.py | 4 ++-- litellm/responses/mcp/litellm_proxy_mcp_handler.py | 8 +++++--- litellm/responses/mcp/mcp_streaming_iterator.py | 2 +- .../responses/mcp/test_litellm_proxy_mcp_handler.py | 2 +- 4 files changed, 9 insertions(+), 7 deletions(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 764fe003f2c..cb6db1be9a4 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -453,7 +453,7 @@ async def acompletion_with_mcp( ) # Make follow-up call with streaming - follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params( + follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler.prepare_follow_up_call_params( self.base_call_args, original_stream_setting=True, ) @@ -628,7 +628,7 @@ async def acompletion_with_mcp( ) # Make follow-up call with original stream setting - follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params( + follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler.prepare_follow_up_call_params( base_call_args, original_stream_setting=stream, ) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index f13bcece7ca..ad54796b102 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -72,7 +72,7 @@ class LiteLLM_Proxy_MCP_Handler: """ @staticmethod - def _prepare_chained_call_params( + 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. @@ -1261,14 +1261,14 @@ class LiteLLM_Proxy_MCP_Handler: return initial_params @staticmethod - def _prepare_follow_up_call_params(call_params: dict[str, Any], original_stream_setting: bool) -> dict[str, Any]: + def prepare_follow_up_call_params(call_params: dict[str, Any], original_stream_setting: bool) -> dict[str, Any]: """ Prepare call parameters for the follow-up LLM call after tool execution. Restores the original streaming setting and removes tool_choice since we're now providing tool results, not requesting tool calls. """ - follow_up_params: Final = LiteLLM_Proxy_MCP_Handler._prepare_chained_call_params(call_params) + 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 @@ -1278,6 +1278,8 @@ class LiteLLM_Proxy_MCP_Handler: return follow_up_params + _prepare_follow_up_call_params = prepare_follow_up_call_params + @staticmethod def _add_mcp_output_elements_to_response( response: ResponsesAPIResponse, diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 1e1cbd34ce2..9dae0254d30 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -782,7 +782,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): ) # Make follow-up call with streaming - follow_up_params: Final = LiteLLM_Proxy_MCP_Handler._prepare_chained_call_params( + follow_up_params: Final = LiteLLM_Proxy_MCP_Handler.prepare_chained_call_params( self.original_request_params ) follow_up_params.update( diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index f2658a1b0a5..5006ea58087 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -459,7 +459,7 @@ def test_prepare_follow_up_call_params_resets_per_call_logging_state(): }, } - follow_up = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params(original, original_stream_setting=True) + 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