ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #22744 from BerriAI/litellm_mcp_streaming_fix

Browse Source 
Add mcp streaming events Fix and consistent response ID
This commit is contained in:
Sameer Kankute 2026-03-04 18:29:48 +05:30 committed by GitHub
commit ece7fdb213
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 278 additions and 52 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

62
litellm/proxy/hooks/responses_id_security.py
View File

@ -188,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
@ -211,31 +212,41 @@ 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 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,
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}"
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):
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 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,
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}"
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
@ -258,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
@ -267,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)
@ -274,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

10
litellm/responses/litellm_completion_transformation/streaming_iterator.py
View File

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

115
litellm/responses/mcp/mcp_streaming_iterator.py
View File

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

143
tests/mcp_tests/test_aresponses_api_with_mcp.py
View File

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