From f1e7ee2bc17d59a23f97b6a79b77dc09bd1b9d57 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 4 May 2026 22:30:25 +0000 Subject: [PATCH] fix(vertex_ai/agent_engine): don't terminate stream on inner-action STOP (#19121) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Vertex Agent Engine emits one SSE event per ADK action — e.g. `transfer_to_agent`, an MCP tool call (`list_cases`), then the final text reply. Each event carries `finish_reason: "STOP"` because STOP is the Gemini-level terminator for that single action, not for the Agent Engine stream. The previous chunk_parser surfaced `finish_reason="stop"` on every chunk that had STOP, so litellm's CustomStreamWrapper closed the stream after the first inner action and the final user-facing text was never delivered to the caller. Fix: - Only surface `finish_reason` when the chunk has user-facing content (text or tool_calls). Inner action / thought-only chunks now pass through with finish_reason=None and don't terminate the stream. - When the chunk is a function call, emit it as an OpenAI-format `tool_calls` delta with `finish_reason="tool_calls"` so callers that consume the agent's tool calls can see them. - Handle both camelCase `functionCall` (Vertex REST) and snake_case `function_call` (Python SDK) shapes. Adds regression tests covering all three cases. Co-authored-by: Mateo Wang --- .../vertex_ai/agent_engine/sse_iterator.py | 119 ++++++++++++------ .../agent_engine/test_transformation.py | 108 ++++++++++++++-- 2 files changed, 179 insertions(+), 48 deletions(-) diff --git a/litellm/llms/vertex_ai/agent_engine/sse_iterator.py b/litellm/llms/vertex_ai/agent_engine/sse_iterator.py index 06fb55e1848..ba93336e75a 100644 --- a/litellm/llms/vertex_ai/agent_engine/sse_iterator.py +++ b/litellm/llms/vertex_ai/agent_engine/sse_iterator.py @@ -4,10 +4,16 @@ SSE Stream Iterator for Vertex AI Agent Engine. Handles Server-Sent Events (SSE) streaming responses from Vertex AI Reasoning Engines. """ -from typing import Any, Union +import json +from typing import Any, List, Optional, Union +from litellm._uuid import uuid from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator -from litellm.types.llms.openai import ChatCompletionUsageBlock +from litellm.types.llms.openai import ( + ChatCompletionToolCallChunk, + ChatCompletionToolCallFunctionChunk, + ChatCompletionUsageBlock, +) from litellm.types.utils import ( Delta, GenericStreamingChunk, @@ -27,46 +33,81 @@ class VertexAgentEngineResponseIterator(BaseModelResponseIterator): def __init__(self, streaming_response: Any, sync_stream: bool) -> None: super().__init__(streaming_response=streaming_response, sync_stream=sync_stream) + @staticmethod + def _extract_parts_from_chunk( + chunk: dict, + ) -> tuple[Optional[str], List[ChatCompletionToolCallChunk]]: + """ + Walk the ``content.parts`` array and split it into: + - the first text part (if any) + - any function_call parts converted to OpenAI tool_calls + + Vertex Agent Engine returns parts in either ``functionCall`` (camelCase, + REST API) or ``function_call`` (snake_case, Python SDK) form. + """ + text: Optional[str] = None + tool_calls: List[ChatCompletionToolCallChunk] = [] + + content = chunk.get("content") or {} + parts = content.get("parts") or [] + + for part in parts: + if not isinstance(part, dict): + continue + + if text is None and "text" in part: + text = part["text"] + continue + + function_call = part.get("functionCall") or part.get("function_call") + if function_call: + call_id = function_call.get("id") or f"call_{uuid.uuid4()}" + tool_calls.append( + ChatCompletionToolCallChunk( + id=call_id, + type="function", + function=ChatCompletionToolCallFunctionChunk( + name=function_call.get("name", ""), + arguments=json.dumps(function_call.get("args") or {}), + ), + index=len(tool_calls), + ) + ) + + return text, tool_calls + def chunk_parser( self, chunk: dict ) -> Union[GenericStreamingChunk, ModelResponseStream]: """ Parse a Vertex Agent Engine response chunk into ModelResponseStream. - Vertex Agent Engine response format: - { - "content": { - "parts": [{"text": "..."}], - "role": "model" - }, - "finish_reason": "STOP", - "usage_metadata": { - "prompt_token_count": 100, - "candidates_token_count": 50, - "total_token_count": 150 - } - } + Vertex Agent Engine emits one SSE event per ADK action (e.g. an inner + ``transfer_to_agent`` call, an MCP tool call, and the final text reply). + Each event carries ``finish_reason: "STOP"`` because ``STOP`` is the + Gemini-level terminator for that single action — it does NOT mean the + Agent Engine stream is finished. The SSE stream ending is the only true + end-of-response signal. + + We therefore only surface ``finish_reason`` when the chunk has + user-facing content (text or tool_calls). Otherwise downstream stream + handling closes the stream after the first inner action and the actual + response is dropped (see issue #19121). """ - # Extract text from content.parts - text = None - content = chunk.get("content", {}) - parts = content.get("parts", []) - for part in parts: - if isinstance(part, dict) and "text" in part: - text = part["text"] - break + text, tool_calls = self._extract_parts_from_chunk(chunk) - # Extract finish_reason - finish_reason = None + finish_reason: Optional[str] = None raw_finish_reason = chunk.get("finish_reason") - if raw_finish_reason == "STOP": - finish_reason = "stop" - elif raw_finish_reason: - finish_reason = raw_finish_reason.lower() + if raw_finish_reason: + if tool_calls: + finish_reason = "tool_calls" + elif text is not None: + finish_reason = ( + "stop" if raw_finish_reason == "STOP" else raw_finish_reason.lower() + ) - # Extract usage from usage_metadata usage = None - usage_metadata = chunk.get("usage_metadata", {}) + usage_metadata = chunk.get("usage_metadata") or {} if usage_metadata: usage = ChatCompletionUsageBlock( prompt_tokens=usage_metadata.get("prompt_token_count", 0), @@ -74,16 +115,22 @@ class VertexAgentEngineResponseIterator(BaseModelResponseIterator): total_tokens=usage_metadata.get("total_token_count", 0), ) - # Return ModelResponseStream (OpenAI-compatible chunk) + delta_kwargs: dict = {} + if text is not None: + delta_kwargs["content"] = text + delta_kwargs["role"] = "assistant" + elif tool_calls: + delta_kwargs["tool_calls"] = tool_calls + delta_kwargs["role"] = "assistant" + else: + delta_kwargs["content"] = None + return ModelResponseStream( choices=[ StreamingChoices( finish_reason=finish_reason, index=0, - delta=Delta( - content=text, - role="assistant" if text else None, - ), + delta=Delta(**delta_kwargs), ) ], usage=usage, diff --git a/tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py b/tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py index af0faee9e21..26a6d28c4a3 100644 --- a/tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py +++ b/tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py @@ -4,6 +4,7 @@ Tests for Vertex AI Agent Engine transformation. Tests the request transformation and streaming chunk parsing without making real API calls. """ +import json import os import sys @@ -73,15 +74,16 @@ class TestVertexAgentEngineTransformRequest: class TestVertexAgentEngineChunkParser: """Tests for the streaming chunk parser.""" - def test_chunk_parser_with_text_content(self): - """ - Test that chunk_parser correctly extracts text from Vertex Agent Engine response format. - """ - iterator = VertexAgentEngineResponseIterator( + def _iterator(self) -> VertexAgentEngineResponseIterator: + return VertexAgentEngineResponseIterator( streaming_response=iter([]), sync_stream=True, ) + def test_chunk_parser_with_text_content(self): + """ + Test that chunk_parser correctly extracts text from Vertex Agent Engine response format. + """ chunk = { "content": { "parts": [{"text": "Hello! I can help you with financial analysis."}], @@ -95,7 +97,7 @@ class TestVertexAgentEngineChunkParser: }, } - result = iterator.chunk_parser(chunk) + result = self._iterator().chunk_parser(chunk) assert ( result.choices[0].delta.content @@ -111,11 +113,6 @@ class TestVertexAgentEngineChunkParser: """ Test that chunk_parser handles chunks without finish_reason (intermediate chunks). """ - iterator = VertexAgentEngineResponseIterator( - streaming_response=iter([]), - sync_stream=True, - ) - chunk = { "content": { "parts": [{"text": "Partial response..."}], @@ -123,8 +120,95 @@ class TestVertexAgentEngineChunkParser: }, } - result = iterator.chunk_parser(chunk) + result = self._iterator().chunk_parser(chunk) assert result.choices[0].delta.content == "Partial response..." assert result.choices[0].finish_reason is None assert result.usage is None + + def test_chunk_parser_intermediate_function_call_does_not_finish_stream(self): + """ + Multi-agent Agent Engine streams emit one SSE event per inner action + (e.g. transfer_to_agent, MCP tool call) and each carries + ``finish_reason: STOP``. STOP here means "this Gemini turn is done", + not "the Agent Engine stream is done", so we must NOT surface + ``finish_reason="stop"`` on these chunks — doing so closes the + downstream stream wrapper before the final text arrives. + + Regression test for https://github.com/BerriAI/litellm/issues/19121. + """ + chunk = { + "content": { + "parts": [ + { + "function_call": { + "id": "adk-redacted", + "args": {"agent_name": "analyst"}, + "name": "transfer_to_agent", + } + } + ], + "role": "model", + }, + "finish_reason": "STOP", + } + + result = self._iterator().chunk_parser(chunk) + + assert result.choices[0].finish_reason == "tool_calls" + tool_calls = result.choices[0].delta.tool_calls + assert tool_calls is not None and len(tool_calls) == 1 + assert tool_calls[0]["function"]["name"] == "transfer_to_agent" + assert json.loads(tool_calls[0]["function"]["arguments"]) == { + "agent_name": "analyst" + } + assert result.choices[0].delta.content is None + + def test_chunk_parser_intermediate_chunk_no_content_drops_finish_reason(self): + """ + Some Agent Engine chunks have ``finish_reason: STOP`` but neither text + nor function_call parts (e.g. thought-only chunks). These must not + terminate the stream — drop the finish_reason. + """ + chunk = { + "content": { + "parts": [{"thought_signature": "..redacted.."}], + "role": "model", + }, + "finish_reason": "STOP", + } + + result = self._iterator().chunk_parser(chunk) + + assert result.choices[0].finish_reason is None + assert result.choices[0].delta.content is None + assert result.choices[0].delta.tool_calls is None + + def test_chunk_parser_camelcase_function_call(self): + """ + Vertex's REST API uses ``functionCall`` (camelCase) — make sure we + handle that as well as the SDK's ``function_call``. + """ + chunk = { + "content": { + "parts": [ + { + "functionCall": { + "id": "adk-1", + "args": {}, + "name": "list_cases", + } + } + ], + "role": "model", + }, + "finish_reason": "STOP", + } + + result = self._iterator().chunk_parser(chunk) + + assert result.choices[0].finish_reason == "tool_calls" + tool_calls = result.choices[0].delta.tool_calls + assert tool_calls is not None and len(tool_calls) == 1 + assert tool_calls[0]["function"]["name"] == "list_cases" + assert tool_calls[0]["id"] == "adk-1"