fix(vertex_ai/agent_engine): don't terminate stream on inner-action STOP (#19121)

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 <mateo-berri@users.noreply.github.com>

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,

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"