From e33a2b486a0b5366fbea356b04bb8195109e9dec Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 27 Feb 2026 04:48:24 -0800 Subject: [PATCH] fixed mcp streaming ordering --- .../responses/mcp/mcp_streaming_iterator.py | 198 ++++++++++-------- .../mcp_tests/test_aresponses_api_with_mcp.py | 102 ++++++++- 2 files changed, 210 insertions(+), 90 deletions(-) diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 731aa5c692b..e4c3fa5da75 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -269,7 +269,9 @@ 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 + # get_response_created: emit response.created first (OpenAI SDK expects it before other events) + # mcp_discovery -> initial_response -> tool_execution -> follow_up_response -> finished + self.phase = "get_response_created" self.finished = False # Event queues and generation flags @@ -388,101 +390,123 @@ 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. get_response_created - Emit response.created first (OpenAI SDK expects it before other events) + 2. mcp_discovery - Emit MCP discovery events + 3. initial_response - Stream the first LLM response + 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 - if self.phase == "initial_response": - 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] - - # 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 - 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 - else: - # base_iterator is not async iterable (likely a ResponsesAPIResponse) - # Collect it for tool execution if needed - if self.should_auto_execute and isinstance( - self.base_iterator, ResponsesAPIResponse - ): - self.collected_response = self.base_iterator - self.phase = "tool_execution" - await self._generate_tool_execution_events() - else: - self.phase = "finished" - raise StopAsyncIteration - - # Phase 3: Tool Execution Events - if self.phase == "tool_execution": - # Emit any queued tool execution events - if self.tool_execution_events: - return self.tool_execution_events.pop(0) - - # Move to follow-up response phase - self.phase = "follow_up_response" - await self._create_follow_up_iterator() - - # Phase 4: Follow-up Response Stream - if self.phase == "follow_up_response": - if self.follow_up_iterator: - try: - return await cast(Any, self.follow_up_iterator).__anext__() # type: ignore[attr-defined] - except StopAsyncIteration: - self.phase = "finished" - raise + while True: + if self.phase == "get_response_created": + result = await self._handle_get_response_created_phase() + elif self.phase == "mcp_discovery": + result = await self._handle_mcp_discovery_phase() + elif self.phase == "initial_response": + result = await self._handle_initial_response_phase() + elif self.phase == "tool_execution": + result = await self._handle_tool_execution_phase() + elif self.phase == "follow_up_response": + result = await self._handle_follow_up_response_phase() else: - self.phase = "finished" raise StopAsyncIteration - # Phase 5: Finished + if result is not None: + return result + + async def _handle_get_response_created_phase( + self, + ) -> Optional[ResponsesAPIStreamingResponse]: + """Emit response.created first (OpenAI SDK expects it before other events).""" + await self._create_initial_response_iterator() if self.phase == "finished": raise StopAsyncIteration + if self.base_iterator and hasattr(self.base_iterator, "__anext__"): + try: + first_chunk = await cast(Any, self.base_iterator).__anext__() # type: ignore[attr-defined] + self.phase = "mcp_discovery" + return first_chunk + except StopAsyncIteration: + self.phase = "finished" + raise + self.phase = "mcp_discovery" + return None - # Should not reach here + async def _handle_mcp_discovery_phase( + self, + ) -> Optional[ResponsesAPIStreamingResponse]: + """Emit MCP discovery events (after response.created).""" + if self.mcp_discovery_events: + return self.mcp_discovery_events.pop(0) + verbose_logger.debug( + "MCP discovery phase complete, transitioning to initial_response" + ) + self.phase = "initial_response" + return None + + async def _handle_initial_response_phase( + self, + ) -> Optional[ResponsesAPIStreamingResponse]: + """Stream the first LLM response.""" + if not self.base_iterator: + self.phase = "finished" + raise StopAsyncIteration + if not hasattr(self.base_iterator, "__anext__"): + return await self._handle_initial_response_non_iterable() + try: + chunk = await cast(Any, self.base_iterator).__anext__() # type: ignore[attr-defined] + if self.should_auto_execute and self._is_response_completed(chunk): + response_obj = getattr(chunk, "response", None) + if isinstance(response_obj, ResponsesAPIResponse): + self.collected_response = response_obj + self.phase = "tool_execution" + await self._generate_tool_execution_events() + return chunk + except StopAsyncIteration: + if self.should_auto_execute and self.collected_response: + self.phase = "tool_execution" + await self._generate_tool_execution_events() + return None + self.phase = "finished" + raise + + async def _handle_initial_response_non_iterable( + self, + ) -> Optional[ResponsesAPIStreamingResponse]: + """Handle base_iterator that is not async iterable (e.g. ResponsesAPIResponse).""" + if self.should_auto_execute and isinstance( + self.base_iterator, ResponsesAPIResponse + ): + self.collected_response = self.base_iterator + self.phase = "tool_execution" + await self._generate_tool_execution_events() + return None + self.phase = "finished" raise StopAsyncIteration + async def _handle_tool_execution_phase( + self, + ) -> Optional[ResponsesAPIStreamingResponse]: + """Emit tool execution events.""" + if self.tool_execution_events: + return self.tool_execution_events.pop(0) + self.phase = "follow_up_response" + await self._create_follow_up_iterator() + return None + + async def _handle_follow_up_response_phase( + self, + ) -> Optional[ResponsesAPIStreamingResponse]: + """Stream the follow-up response.""" + if not self.follow_up_iterator: + self.phase = "finished" + raise StopAsyncIteration + try: + return await cast(Any, self.follow_up_iterator).__anext__() # type: ignore[attr-defined] + except StopAsyncIteration: + self.phase = "finished" + raise + def _is_response_completed(self, chunk: ResponsesAPIStreamingResponse) -> bool: """Check if this chunk indicates the response is completed""" from litellm.types.llms.openai import ResponsesAPIStreamEvents diff --git a/tests/mcp_tests/test_aresponses_api_with_mcp.py b/tests/mcp_tests/test_aresponses_api_with_mcp.py index bae0b15dfec..7ce87dc571f 100644 --- a/tests/mcp_tests/test_aresponses_api_with_mcp.py +++ b/tests/mcp_tests/test_aresponses_api_with_mcp.py @@ -467,15 +467,111 @@ async def test_mcp_allowed_tools_filtering(): print("✓ MCP allowed_tools filtering test completed successfully!") +@pytest.mark.asyncio +async def test_sse_event_ordering_response_created_first(): + """ + Test that response.created is emitted before mcp_list_tools events. + OpenAI Node SDK expects response.created as the first SSE event. + """ + from unittest.mock import AsyncMock, patch + from litellm.types.llms.openai import ( + ResponsesAPIStreamEvents, + ResponseCreatedEvent, + ResponsesAPIResponse, + ) + from litellm.responses.mcp.mcp_streaming_iterator import ( + MCPEnhancedStreamingIterator, + create_mcp_list_tools_events, + ) + + mock_mcp_tools = [ + type("MCPTool", (), { + "name": "test_tool", + "description": "Test", + "inputSchema": {"type": "object", "properties": {}}, + })(), + ] + + mcp_discovery_events = await create_mcp_list_tools_events( + mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/test"}], + user_api_key_auth=None, + base_item_id="mcp_test123", + pre_processed_mcp_tools=mock_mcp_tools, + ) + + async def mock_stream(): + yield ResponseCreatedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_CREATED, + response=ResponsesAPIResponse( + id="resp_123", + object="response", + created_at=1234567890, + status="in_progress", + error=None, + incomplete_details=None, + instructions=None, + max_output_tokens=None, + model="gpt-4o-mini", + output=[], + parallel_tool_calls=True, + previous_response_id=None, + reasoning=None, + store=True, + temperature=1.0, + text={"format": {"type": "text"}}, + tool_choice="auto", + tools=[], + top_p=1.0, + truncation="disabled", + usage=None, + user=None, + metadata={}, + ), + ) + + mock_stream_obj = mock_stream() + + with patch( + "litellm.responses.main.aresponses", + new_callable=AsyncMock, + return_value=mock_stream_obj, + ): + iterator = MCPEnhancedStreamingIterator( + base_iterator=None, + mcp_events=mcp_discovery_events, + tool_server_map={"test": "test_server"}, + mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy"}], + user_api_key_auth=None, + original_request_params={ + "model": "gpt-4o-mini", + "stream": True, + "tools": [{"type": "mcp", "server_url": "litellm_proxy"}], + }, + ) + + event_types = [] + async for chunk in iterator: + event_types.append(getattr(chunk, "type", "unknown")) + if len(event_types) >= 5: + break + + assert len(event_types) >= 1, "Should have at least one event" + assert event_types[0] == ResponsesAPIStreamEvents.RESPONSE_CREATED, ( + f"First event must be response.created for OpenAI SDK compatibility, got {event_types[0]}. " + f"Order: {event_types}" + ) + + @pytest.mark.asyncio async def test_streaming_mcp_events_validation(): """ Test that MCP streaming events are properly emitted when using streaming with MCP tools. This test validates: - 1. MCP discovery events are emitted first - 2. Regular streaming response events follow - 3. Tool execution events are emitted when tools are auto-executed + 1. response.created is emitted first (OpenAI SDK requirement) + 2. MCP discovery events follow + 3. Regular streaming response events follow + 4. Tool execution events are emitted when tools are auto-executed """ from unittest.mock import AsyncMock, patch from litellm.types.llms.openai import ResponsesAPIStreamEvents