From 28c7c2401c38882e2c9f2ad39b8d9d2964186269 Mon Sep 17 00:00:00 2001
From: kcncb <128760047+kcncb@users.noreply.github.com>
Date: Wed, 28 Oct 2026 11:28:40 +0000
Subject: [PATCH] fix(gemini): report tool_calls when a streamed call ends with
 MAX_TOKENS

A Gemini reply that calls a function and runs out of budget ends with
`finishReason: "MAX_TOKENS"`.  Non-streaming reports `tool_calls` for it,
streaming reported `length`: the empty-text final chunk is mapped from that
chunk alone, and the override that recovers "the calls came earlier" only
rewrote `stop`.

`length` is a normal termination for a call as well, so it is rewritten too.
A filtered reply keeps `content_filter`, which is what the non-streaming path
reports for it, so the override does not swallow that.

Covers both chunk shapes: the call and the reason in separate chunks, and both
in one chunk.
---
 .../vertex_and_google_ai_studio_gemini.py     | 12 ++-
 ...emini_streaming_tool_call_finish_reason.py | 76 +++++++++++++++++++
 2 files changed, 84 insertions(+), 4 deletions(-)

diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
index 9acbc9c459c..a9535a9a605 100644
--- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
+++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
@@ -3982,12 +3982,16 @@ class ModelResponseIterator:
 
         # Also handle the case where the final chunk has empty
         # content (e.g. text:"") WITH finishReason. In this case
-        # _process_candidates DOES create a choice, but maps
-        # finishReason="STOP" to "stop" because the current chunk
-        # has no tool_calls. Override if we saw tool_calls earlier.
+        # _process_candidates DOES create a choice, but maps the reason
+        # from what this chunk alone shows, which is never a call. Override
+        # if we saw tool_calls earlier.
+        # "length" is a normal termination for a call too: Gemini reports
+        # MAX_TOKENS when the reply with the call ran out of budget, and the
+        # non-streaming path reports "tool_calls" for that same reply.  A
+        # filtered reply keeps its reason, matching the non-streaming path.
         if self.has_seen_tool_calls:
             for choice in model_response.choices:
-                if choice.finish_reason == "stop":
+                if choice.finish_reason in ("stop", "length"):
                     choice.finish_reason = "tool_calls"
 
         VertexGeminiConfig.set_stream_metadata_on_response(
diff --git a/tests/unit/llms/vertex_ai/gemini/test_gemini_streaming_tool_call_finish_reason.py b/tests/unit/llms/vertex_ai/gemini/test_gemini_streaming_tool_call_finish_reason.py
index 92464ae2c31..ce162afbd82 100644
--- a/tests/unit/llms/vertex_ai/gemini/test_gemini_streaming_tool_call_finish_reason.py
+++ b/tests/unit/llms/vertex_ai/gemini/test_gemini_streaming_tool_call_finish_reason.py
@@ -11,6 +11,8 @@ Ref: https://github.com/BerriAI/litellm/issues/21041
 
 from unittest.mock import MagicMock
 
+import pytest
+
 from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
     ModelResponseIterator,
 )
@@ -338,3 +340,77 @@ def test_streaming_metadata_only_chunk_does_not_yield_empty_choices():
     assert len(response.choices) == 1
     assert response.choices[0].finish_reason is None
     assert response.choices[0].delta.content is None
+
+
+@pytest.mark.parametrize(
+    "finish_reason, expected",
+    [
+        ("STOP", "tool_calls"),
+        ("MAX_TOKENS", "tool_calls"),
+        ("SAFETY", "content_filter"),
+    ],
+)
+def test_streaming_tool_call_finish_reason_matches_non_streaming(finish_reason: str, expected: str):
+    """The same reply must report the same reason streamed and not streamed.
+
+    Gemini sends the call in one chunk and the reason in the next.  Only "stop"
+    was rewritten to "tool_calls", so a reply with a call that ran into
+    MAX_TOKENS reported "length" while the non-streaming path reports
+    "tool_calls" for it.  A filtered reply keeps "content_filter" on both paths,
+    so that one is pinned too -- the rewrite must not swallow it.
+    """
+    logging_obj = _make_logging_obj()
+    iterator = ModelResponseIterator(
+        streaming_response=iter([]),
+        sync_stream=True,
+        logging_obj=logging_obj,
+    )
+
+    iterator.chunk_parser(
+        {
+            "candidates": [
+                {
+                    "content": {
+                        "parts": [{"functionCall": {"name": "get_weather", "args": {}}}],
+                        "role": "model",
+                    },
+                    "index": 0,
+                }
+            ]
+        }
+    )
+    final = iterator.chunk_parser({"candidates": [{"finishReason": finish_reason, "index": 0}]})
+
+    assert final is not None
+    assert final.choices[0].finish_reason == expected
+
+
+def test_streaming_tool_call_and_max_tokens_in_one_chunk():
+    """The other chunk shape: the call and the reason arrive together."""
+    logging_obj = _make_logging_obj()
+    iterator = ModelResponseIterator(
+        streaming_response=iter([]),
+        sync_stream=True,
+        logging_obj=logging_obj,
+    )
+
+    response = iterator.chunk_parser(
+        {
+            "candidates": [
+                {
+                    "content": {
+                        "parts": [
+                            {"functionCall": {"name": "get_weather", "args": {}}},
+                            {"text": ""},
+                        ],
+                        "role": "model",
+                    },
+                    "finishReason": "MAX_TOKENS",
+                    "index": 0,
+                }
+            ]
+        }
+    )
+
+    assert response is not None
+    assert response.choices[0].finish_reason == "tool_calls"