From e0880c3deee1731a13586b977b15f09c078782a9 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 4 Mar 2026 12:11:16 +0530 Subject: [PATCH 1/3] Fix order of streaming for mcp responses --- litellm/proxy/hooks/responses_id_security.py | 55 +++++---- .../streaming_iterator.py | 10 +- .../responses/mcp/mcp_streaming_iterator.py | 115 ++++++++++++++---- 3 files changed, 129 insertions(+), 51 deletions(-) diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index d196a68d36..15798e1547 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -29,7 +29,8 @@ if TYPE_CHECKING: class ResponsesIDSecurity(CustomLogger): def __init__(self): - pass + # Cache for encrypted response IDs to ensure consistency across streaming chunks + self._encrypted_response_id_cache: dict[str, str] = {} async def async_pre_call_hook( self, @@ -211,31 +212,39 @@ class ResponsesIDSecurity(CustomLogger): and isinstance(response_id, str) and response_id.startswith("resp_") ): - encrypted_response_id = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( - response_id, - user_api_key_dict.user_id or "", - user_api_key_dict.team_id or "", - ) + # Check cache first to ensure consistency across streaming chunks + if response_id in self._encrypted_response_id_cache: + setattr(response, "id", self._encrypted_response_id_cache[response_id]) + else: + encrypted_response_id = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( + response_id, + user_api_key_dict.user_id or "", + user_api_key_dict.team_id or "", + ) - encoded_user_id_and_response_id = encrypt_value_helper( - value=encrypted_response_id - ) - setattr( - response, "id", f"resp_{encoded_user_id_and_response_id}" - ) # maintain the 'resp_' prefix for the responses api response id + encoded_user_id_and_response_id = encrypt_value_helper( + value=encrypted_response_id + ) + encrypted_id = f"resp_{encoded_user_id_and_response_id}" + self._encrypted_response_id_cache[response_id] = encrypted_id + setattr(response, "id", encrypted_id) elif response_obj and isinstance(response_obj, ResponsesAPIResponse): - encrypted_response_id = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( - response_obj.id, - user_api_key_dict.user_id or "", - user_api_key_dict.team_id or "", - ) - encoded_user_id_and_response_id = encrypt_value_helper( - value=encrypted_response_id - ) - setattr( - response_obj, "id", f"resp_{encoded_user_id_and_response_id}" - ) # maintain the 'resp_' prefix for the responses api response id + # Check cache first to ensure consistency across streaming chunks + if response_obj.id in self._encrypted_response_id_cache: + setattr(response_obj, "id", self._encrypted_response_id_cache[response_obj.id]) + else: + encrypted_response_id = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( + response_obj.id, + user_api_key_dict.user_id or "", + user_api_key_dict.team_id or "", + ) + encoded_user_id_and_response_id = encrypt_value_helper( + value=encrypted_response_id + ) + encrypted_id = f"resp_{encoded_user_id_and_response_id}" + self._encrypted_response_id_cache[response_obj.id] = encrypted_id + setattr(response_obj, "id", encrypted_id) setattr(response, "response", response_obj) return response diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 6e32a0d48d..e7866ae0f0 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -344,8 +344,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._pending_tool_events.append(item_done_event) def _default_response_created_event_data(self) -> dict: + # Use cached response ID if available, otherwise generate a new one + if self._cached_response_id is None: + self._cached_response_id = f"resp_{str(uuid.uuid4())}" + response_created_event_data = { - "id": f"resp_{str(uuid.uuid4())}", + "id": self._cached_response_id, "object": "response", "created_at": int(time.time()), "status": "in_progress", @@ -1074,6 +1078,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): responses_api_request=self.responses_api_request, ) + # Use the cached response ID to ensure consistency across all events + if self._cached_response_id: + responses_api_response.id = self._cached_response_id + # Encode the response ID to match non-streaming behavior encoded_response = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( responses_api_response=responses_api_response, diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 731aa5c692..0ada532c41 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -269,7 +269,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.should_auto_execute = self._should_auto_execute_tools() # Streaming state management - self.phase = "mcp_discovery" # mcp_discovery -> initial_response -> tool_execution -> follow_up_response -> finished + self.phase = "initial_response" # initial_response -> mcp_discovery -> tool_execution -> follow_up_response -> finished self.finished = False # Event queues and generation flags @@ -305,6 +305,12 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Mark as async iterator self.is_async = True + + # Track if we've emitted initial OpenAI lifecycle events + self.initial_events_emitted = False + + # Cache the response ID to ensure consistency across all events + self._cached_response_id: Optional[str] = None def _extract_mcp_headers_from_params(self) -> None: """Extract MCP headers from original request params to pass to tool calls""" @@ -388,38 +394,43 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): async def __anext__(self) -> ResponsesAPIStreamingResponse: """ Phase-based streaming: - 1. mcp_discovery - Emit MCP discovery events - 2. initial_response - Stream the first LLM response - 3. tool_execution - Emit tool execution events - 4. follow_up_response - Stream the follow-up response - 5. finished - End iteration + 1. initial_response - Stream the first LLM response (includes response.created, response.in_progress, response.output_item.added) + 2. mcp_discovery - Emit MCP discovery events (after response.output_item.added) + 3. continue_initial_response - Continue streaming the initial response content + 4. tool_execution - Emit tool execution events + 5. follow_up_response - Stream the follow-up response + 6. finished - End iteration """ - # Phase 1: MCP Discovery Events - if self.phase == "mcp_discovery": - # Generate MCP discovery events if not already done - # MCP discovery events are already generated and available - - # Emit MCP discovery events - if self.mcp_discovery_events: - return self.mcp_discovery_events.pop(0) - - # All MCP discovery events emitted, move to next phase - verbose_logger.debug( - "MCP discovery phase complete, transitioning to initial_response" - ) - self.phase = "initial_response" - await self._create_initial_response_iterator() - # Fall through to process the initial response immediately - - # Phase 2: Initial Response Stream + # Phase 1: Initial Response Stream (emit standard OpenAI events first) if self.phase == "initial_response": + # Create the initial response iterator if not already created + if self.base_iterator is None: + await self._create_initial_response_iterator() + if self.base_iterator: # Check if base_iterator is actually iterable if hasattr(self.base_iterator, "__anext__"): try: chunk = await cast(Any, self.base_iterator).__anext__() # type: ignore[attr-defined] + # Capture the response ID from the first event to ensure consistency + if self._cached_response_id is None and hasattr(chunk, 'response'): + response_obj = getattr(chunk, 'response', None) + if response_obj and hasattr(response_obj, 'id'): + self._cached_response_id = response_obj.id + verbose_logger.debug(f"Cached response ID: {self._cached_response_id}") + + # After emitting response.output_item.added, transition to MCP discovery + # Check if this is the output_item.added event + if not self.initial_events_emitted and hasattr(chunk, 'type'): + chunk_type = getattr(chunk, 'type', None) + if chunk_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED: + self.initial_events_emitted = True + # Transition to MCP discovery phase after returning this chunk + self.phase = "mcp_discovery" + return chunk + # If auto-execution is enabled, check for completed responses if self.should_auto_execute and self._is_response_completed( chunk @@ -454,7 +465,28 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.phase = "finished" raise StopAsyncIteration - # Phase 3: Tool Execution Events + # Phase 2: MCP Discovery Events (after response.output_item.added) + if self.phase == "mcp_discovery": + # Emit MCP discovery events + if self.mcp_discovery_events: + return self.mcp_discovery_events.pop(0) + self.phase = "continue_initial_response" + # Fall through to continue processing the initial response + + # Phase 3: Continue Initial Response (after MCP discovery events) + if self.phase == "continue_initial_response": + try: + return await self._process_base_iterator_chunk() + except StopAsyncIteration: + # Initial response ended, move to next phase + if self.should_auto_execute and self.collected_response: + self.phase = "tool_execution" + await self._generate_tool_execution_events() + else: + self.phase = "finished" + raise + + # Phase 4: Tool Execution Events if self.phase == "tool_execution": # Emit any queued tool execution events if self.tool_execution_events: @@ -464,7 +496,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.phase = "follow_up_response" await self._create_follow_up_iterator() - # Phase 4: Follow-up Response Stream + # Phase 5: Follow-up Response Stream if self.phase == "follow_up_response": if self.follow_up_iterator: try: @@ -476,7 +508,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.phase = "finished" raise StopAsyncIteration - # Phase 5: Finished + # Phase 6: Finished if self.phase == "finished": raise StopAsyncIteration @@ -491,6 +523,35 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): getattr(chunk, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED ) + async def _process_base_iterator_chunk(self) -> ResponsesAPIStreamingResponse: + """ + Process a chunk from the base iterator with response ID consistency enforcement. + """ + if not self.base_iterator or not hasattr(self.base_iterator, "__anext__"): + raise StopAsyncIteration + + chunk = await cast(Any, self.base_iterator).__anext__() # type: ignore[attr-defined] + + # Ensure response ID consistency - update chunk if needed + if self._cached_response_id and hasattr(chunk, 'response'): + response_obj = getattr(chunk, 'response', None) + if response_obj and hasattr(response_obj, 'id'): + if response_obj.id != self._cached_response_id: + verbose_logger.debug(f"Updating response ID from {response_obj.id} to {self._cached_response_id}") + response_obj.id = self._cached_response_id + + # If auto-execution is enabled, check for completed responses + if self.should_auto_execute and self._is_response_completed(chunk): + # Collect the response for tool execution + response_obj = getattr(chunk, "response", None) + if isinstance(response_obj, ResponsesAPIResponse): + self.collected_response = response_obj + # Move to tool execution phase after emitting this chunk + self.phase = "tool_execution" + await self._generate_tool_execution_events() + + return chunk + async def _create_initial_response_iterator(self) -> None: """Create the initial response iterator by making the first LLM call""" try: From f878923d260f8fe5eb46993c0602a7d409451c1d Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 4 Mar 2026 12:19:54 +0530 Subject: [PATCH 2/3] Add test for correct streaming chunks and responses id consistency --- .../mcp_tests/test_aresponses_api_with_mcp.py | 143 ++++++++++++++++++ 1 file changed, 143 insertions(+) diff --git a/tests/mcp_tests/test_aresponses_api_with_mcp.py b/tests/mcp_tests/test_aresponses_api_with_mcp.py index c22c3537af..8c8582a35d 100644 --- a/tests/mcp_tests/test_aresponses_api_with_mcp.py +++ b/tests/mcp_tests/test_aresponses_api_with_mcp.py @@ -1250,4 +1250,147 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e(): } +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["gpt-4o-mini"]) +async def test_streaming_mcp_event_order_and_response_id_consistency( + model: str, caplog: pytest.LogCaptureFixture +): + """ + Test that: + 1. Streaming events are emitted in correct order (response.created, response.in_progress, response.output_item.added before MCP events) + 2. All response lifecycle events share the same response ID within a cycle + """ + if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv("OPENAI_API_KEY"): + pytest.skip("OPENAI_API_KEY not set, skipping openai model test") + + from unittest.mock import AsyncMock, patch + + mock_mcp_tools = [ + type('MCPTool', (), { + 'name': 'get_weather', + 'description': 'Get weather for a city', + 'inputSchema': { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"} + }, + "required": ["city"] + } + })() + ] + + with caplog.at_level(logging.ERROR): + with patch.object( + LiteLLM_Proxy_MCP_Handler, + '_get_mcp_tools_from_manager', + new_callable=AsyncMock, + ) as mock_get_tools, patch.object( + LiteLLM_Proxy_MCP_Handler, + '_execute_tool_calls', + new_callable=AsyncMock, + ) as mock_execute_tools: + mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"]) + + def mock_execute_side_effect(tool_calls, user_api_key_auth, **kwargs): + results = [] + for tool_call in tool_calls: + call_id = None + if isinstance(tool_call, dict): + call_id = tool_call.get("call_id") or tool_call.get("id") + elif hasattr(tool_call, 'call_id'): + call_id = tool_call.call_id + elif hasattr(tool_call, 'id'): + call_id = tool_call.id + if call_id: + results.append({ + "tool_call_id": call_id, + "result": "Sunny, 72°F", + }) + return results + + mock_execute_tools.side_effect = mock_execute_side_effect + + mcp_tool_config = cast(Any, { + "type": "mcp", + "server_url": "litellm_proxy", + "require_approval": "never", + }) + + response = await litellm.aresponses( + model=model, + tools=[mcp_tool_config], + input=[{ + "role": "user", + "type": "message", + "content": "What's the weather in San Francisco?" + }], + stream=True, + ) + + events = [] + async for chunk in response: + events.append(chunk) + + assert len(events) > 0, "Should receive streaming events" + + created_idx = next((i for i, e in enumerate(events) if getattr(e, 'type', None) == 'response.created'), None) + in_progress_idx = next((i for i, e in enumerate(events) if getattr(e, 'type', None) == 'response.in_progress'), None) + output_item_added_idx = next((i for i, e in enumerate(events) if getattr(e, 'type', None) == 'response.output_item.added'), None) + mcp_in_progress_idx = next((i for i, e in enumerate(events) if 'mcp_list_tools.in_progress' in str(getattr(e, 'type', ''))), None) + completed_idx = next((i for i, e in enumerate(events) if getattr(e, 'type', None) == 'response.completed'), None) + + assert created_idx is not None, "response.created event should be present" + assert in_progress_idx is not None, "response.in_progress event should be present" + assert output_item_added_idx is not None, "response.output_item.added event should be present" + + assert created_idx < in_progress_idx, "response.created should come before response.in_progress" + assert in_progress_idx < output_item_added_idx, "response.in_progress should come before response.output_item.added" + + if mcp_in_progress_idx is not None: + assert output_item_added_idx < mcp_in_progress_idx, "response.output_item.added should come before response.mcp_list_tools.in_progress" + + response_ids = [] + for i, event in enumerate(events): + event_type = getattr(event, 'type', None) + if hasattr(event, 'response'): + response_obj = getattr(event, 'response', None) + if response_obj and hasattr(response_obj, 'id'): + event_type_value = event_type.value if hasattr(event_type, 'value') else str(event_type) + if any(x in event_type_value for x in ['response.created', 'response.in_progress', 'response.completed']): + response_ids.append((i, event_type_value, response_obj.id)) + + assert len(response_ids) >= 2, f"Should have at least 2 response lifecycle events. Found {len(response_ids)}" + + cycles = [] + current_cycle = [] + current_id = None + + for idx, event_type, resp_id in response_ids: + if current_id is None or resp_id == current_id: + current_cycle.append((idx, event_type, resp_id)) + current_id = resp_id + else: + if current_cycle: + cycles.append(current_cycle) + current_cycle = [(idx, event_type, resp_id)] + current_id = resp_id + if current_cycle: + cycles.append(current_cycle) + + for cycle_num, cycle in enumerate(cycles): + cycle_ids = set(resp_id for _, _, resp_id in cycle) + assert len(cycle_ids) == 1, f"Cycle {cycle_num + 1} should have consistent response ID. Found {len(cycle_ids)} unique IDs" + + assert completed_idx is not None, "response.completed event should be present" + + lite_errors = [ + record for record in caplog.records + if record.levelno >= logging.ERROR + and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage()) + ] + assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join( + record.getMessage() for record in lite_errors + ) + + From 9907c635efd894c6ab0bf797213bda9f65ca820a Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 4 Mar 2026 12:30:49 +0530 Subject: [PATCH 3/3] Fix: Removed the process-level _encrypted_response_id_cache from __init__ --- litellm/proxy/hooks/responses_id_security.py | 29 ++++++++++++-------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index 15798e1547..39f33ade38 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -29,8 +29,7 @@ if TYPE_CHECKING: class ResponsesIDSecurity(CustomLogger): def __init__(self): - # Cache for encrypted response IDs to ensure consistency across streaming chunks - self._encrypted_response_id_cache: dict[str, str] = {} + pass async def async_pre_call_hook( self, @@ -189,6 +188,7 @@ class ResponsesIDSecurity(CustomLogger): self, response: BaseLiteLLMOpenAIResponseObject, user_api_key_dict: "UserAPIKeyAuth", + request_cache: Optional[dict[str, str]] = None, ) -> BaseLiteLLMOpenAIResponseObject: # encrypt the response id using the symmetric key # encrypt the response id, and encode the user id and response id in base64 @@ -212,9 +212,9 @@ class ResponsesIDSecurity(CustomLogger): and isinstance(response_id, str) and response_id.startswith("resp_") ): - # Check cache first to ensure consistency across streaming chunks - if response_id in self._encrypted_response_id_cache: - setattr(response, "id", self._encrypted_response_id_cache[response_id]) + # Check request-scoped cache first (for streaming consistency) + if request_cache is not None and response_id in request_cache: + setattr(response, "id", request_cache[response_id]) else: encrypted_response_id = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( response_id, @@ -226,13 +226,14 @@ class ResponsesIDSecurity(CustomLogger): value=encrypted_response_id ) encrypted_id = f"resp_{encoded_user_id_and_response_id}" - self._encrypted_response_id_cache[response_id] = encrypted_id + if request_cache is not None: + request_cache[response_id] = encrypted_id setattr(response, "id", encrypted_id) elif response_obj and isinstance(response_obj, ResponsesAPIResponse): - # Check cache first to ensure consistency across streaming chunks - if response_obj.id in self._encrypted_response_id_cache: - setattr(response_obj, "id", self._encrypted_response_id_cache[response_obj.id]) + # Check request-scoped cache first (for streaming consistency) + if request_cache is not None and response_obj.id in request_cache: + setattr(response_obj, "id", request_cache[response_obj.id]) else: encrypted_response_id = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( response_obj.id, @@ -243,7 +244,8 @@ class ResponsesIDSecurity(CustomLogger): value=encrypted_response_id ) encrypted_id = f"resp_{encoded_user_id_and_response_id}" - self._encrypted_response_id_cache[response_obj.id] = encrypted_id + if request_cache is not None: + request_cache[response_obj.id] = encrypted_id setattr(response_obj, "id", encrypted_id) setattr(response, "response", response_obj) return response @@ -267,7 +269,7 @@ class ResponsesIDSecurity(CustomLogger): if isinstance(response, ResponsesAPIResponse): response = cast( ResponsesAPIResponse, - self._encrypt_response_id(response, user_api_key_dict), + self._encrypt_response_id(response, user_api_key_dict, request_cache=None), ) return response @@ -276,6 +278,9 @@ class ResponsesIDSecurity(CustomLogger): ) -> AsyncGenerator[BaseLiteLLMOpenAIResponseObject, None]: from litellm.proxy.proxy_server import general_settings + # Create a request-scoped cache for consistent encryption across streaming chunks. + request_encryption_cache: dict[str, str] = {} + async for chunk in response: if ( isinstance(chunk, BaseLiteLLMOpenAIResponseObject) @@ -283,5 +288,5 @@ class ResponsesIDSecurity(CustomLogger): == "/v1/responses" # only encrypt the response id for the responses api and not general_settings.get("disable_responses_id_security", False) ): - chunk = self._encrypt_response_id(chunk, user_api_key_dict) + chunk = self._encrypt_response_id(chunk, user_api_key_dict, request_encryption_cache) yield chunk