fix(gemini): pair streamed server-side tool calls with their responses

litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py

@@ -7,7 +7,8 @@ import threading
 import time
 from collections.abc import Awaitable, Callable, Mapping, Sequence
 from copy import deepcopy
-from functools import lru_cache, partial
+from functools import lru_cache, partial, reduce
+from itertools import chain
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args
 
@@ -79,6 +80,7 @@ from litellm.types.utils import (
     ChatCompletionAudioResponse,
     ChatCompletionTokenLogprob,
     ChoiceLogprobs,
+    Choices,
     CompletionTokensDetailsWrapper,
     PromptTokensDetailsWrapper,
     TopLogprob,
@@ -233,6 +235,49 @@ def _served_model_name(model_version: object) -> str | None:
     return model_version.split("@", 1)[0]
 
 
+def _chunk_field(value: object, name: str) -> object:
+    if isinstance(value, Mapping):
+        return cast(Mapping[str, object], value).get(name)  # cast-ok: streamed chunk dicts are untyped JSON
+    return getattr(value, name, None)
+
+
+def _delta_provider_fields(chunk: object) -> Mapping[str, object] | None:
+    choices: Final = _chunk_field(chunk, "choices")
+    if not isinstance(choices, list) or not choices:
+        return None
+    first_choice: Final = cast(Sequence[object], choices)[0]  # cast-ok: streamed chunk choices are untyped
+    return as_str_mapping(_chunk_field(_chunk_field(first_choice, "delta"), "provider_specific_fields"))
+
+
+def _list_field(fields: Mapping[str, object], key: str) -> Sequence[object]:
+    value: Final = fields.get(key)
+    return cast(Sequence[object], value) if isinstance(value, list) else ()  # cast-ok: items forwarded verbatim
+
+
+def _paired_invocation(first: Mapping[str, object], second: Mapping[str, object]) -> Mapping[str, object]:
+    if ("response" in first) == ("response" in second):
+        return first
+    call, result = (first, second) if "response" in second else (second, first)
+    return {
+        **{key: value for key, value in call.items() if key not in ("response", "response_thought_signature")},
+        "response": result["response"],
+        **(
+            {"response_thought_signature": result["response_thought_signature"]}
+            if "response_thought_signature" in result
+            else {}
+        ),
+    }
+
+
+def _paired_invocations(invocations: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, object], ...]:
+    ids: Final = tuple(dict.fromkeys(invocation.get("id") for invocation in invocations))
+    return tuple(
+        reduce(_paired_invocation, (invocation for invocation in invocations if invocation.get("id") == invocation_id))
+        for invocation_id in ids
+        if invocation_id is not None
+    ) + tuple(invocation for invocation in invocations if invocation.get("id") is None)
+
+
 _GEMINI_SPEECH_CONFIG_KEY_MAP: Final[Mapping[str, str]] = MappingProxyType(
     {
         "multi_speaker_voice_config": "multiSpeakerVoiceConfig",
@@ -2706,6 +2751,26 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
             if merged:
                 setattr(response, field_name, merged)
                 response.hidden_params[field_name] = merged
+        VertexGeminiConfig._merge_streamed_message_provider_fields(response, chunks)
+
+    @staticmethod
+    def _merge_streamed_message_provider_fields(response: ModelResponse, chunks: Sequence[object]) -> None:
+        # a server-side toolCall and its toolResponse stream in different chunks, pair them as the unary path does
+        per_chunk: Final = tuple(fields for chunk in chunks if (fields := _delta_provider_fields(chunk)) is not None)
+        invocations: Final = tuple(
+            mapping
+            for invocation in chain.from_iterable(
+                _list_field(fields, "server_side_tool_invocations") for fields in per_chunk
+            )
+            if (mapping := as_str_mapping(invocation)) is not None
+        )
+        choice: Final = response.choices[0] if response.choices else None
+        if not invocations or not isinstance(choice, Choices):
+            return
+        choice.message.provider_specific_fields = {
+            **(choice.message.provider_specific_fields or {}),
+            "server_side_tool_invocations": list(_paired_invocations(invocations)),
+        }
 
     @staticmethod
     def _convert_grounding_metadata_to_annotations(

tests/unit/llms/vertex_ai/gemini/test_context_circulation.py

@@ -9,11 +9,16 @@ parts for server-side tools (e.g. Google Search). These must be:
 """
 
 import json
-from typing import Any, Dict, List
+from collections.abc import Mapping
+from typing import Any, Dict, Final, List
 from unittest.mock import MagicMock
 
+import httpx
 import pytest
 
+import litellm
+from litellm import stream_chunk_builder
+from litellm.llms.custom_httpx.http_handler import HTTPHandler
 from litellm.llms.vertex_ai.gemini.transformation import (
     gemini_convert_messages_with_history,
 )
@@ -21,6 +26,8 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
     VertexGeminiConfig,
 )
 from litellm.types.llms.vertex_ai import HttpxPartType
+from litellm.types.utils import Choices, ModelResponse
+from litellm.utils import CustomStreamWrapper
 
 # --- Response extraction tests ---
 
@@ -199,6 +206,102 @@ class TestExtractServerSideToolInvocations:
         assert result[0]["response_thought_signature"] == "sig_resp_1"
 
 
+# --- Streaming tests ---
+
+_SEARCH_CALL: Final = {
+    "toolCall": {"toolType": "GOOGLE_SEARCH_WEB", "id": "srv_1", "args": {"queries": ["weather paris"]}},
+    "thoughtSignature": "sig_call",
+}
+_SEARCH_RESPONSE: Final = {
+    "toolResponse": {"toolType": "GOOGLE_SEARCH_WEB", "id": "srv_1", "response": {"results": ["18C, sunny"]}},
+    "thoughtSignature": "sig_resp",
+}
+_ANSWER_TEXT: Final = "It is 18C and sunny in Paris."
+_QUESTION: Final = {"role": "user", "content": "Weather in Paris?"}
+
+
+def _gemini_body(*parts: Mapping[str, object], finish: str | None = None) -> Mapping[str, object]:
+    return {
+        "candidates": [
+            {
+                "content": {"role": "model", "parts": list(parts)},
+                "index": 0,
+                **({"finishReason": finish} if finish is not None else {}),
+            }
+        ],
+        "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 8, "totalTokenCount": 18},
+    }
+
+
+def _gemini_client(response: httpx.Response) -> HTTPHandler:
+    return HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(lambda request: response)))
+
+
+def _provider_fields(response: object) -> Mapping[str, object]:
+    assert isinstance(response, ModelResponse)
+    choice = response.choices[0]
+    assert isinstance(choice, Choices)
+    assert choice.message.provider_specific_fields is not None
+    return choice.message.provider_specific_fields
+
+
+def _streamed_provider_fields() -> Mapping[str, object]:
+    sse = "".join(
+        f"data: {json.dumps(chunk)}\r\n\r\n"
+        for chunk in (
+            _gemini_body(_SEARCH_CALL),
+            _gemini_body(_SEARCH_RESPONSE),
+            _gemini_body({"text": _ANSWER_TEXT}, finish="STOP"),
+        )
+    )
+    stream = litellm.completion(
+        model="gemini/gemini-3.8-flash",
+        messages=[_QUESTION],
+        api_key="fake-gemini-key",
+        stream=True,
+        client=_gemini_client(
+            httpx.Response(200, content=sse.encode(), headers={"content-type": "text/event-stream"})
+        ),
+    )
+    assert isinstance(stream, CustomStreamWrapper)
+    return _provider_fields(stream_chunk_builder(list(stream), messages=[_QUESTION]))
+
+
+class TestStreamedServerSideToolInvocations:
+    """stream_chunk_builder must rebuild invocations split across chunks (#25869)."""
+
+    def test_rebuilt_stream_matches_the_non_streaming_response(self):
+        unary = litellm.completion(
+            model="gemini/gemini-3.8-flash",
+            messages=[_QUESTION],
+            api_key="fake-gemini-key",
+            client=_gemini_client(
+                httpx.Response(
+                    200,
+                    json=_gemini_body(_SEARCH_CALL, _SEARCH_RESPONSE, {"text": _ANSWER_TEXT}, finish="STOP"),
+                )
+            ),
+        )
+
+        assert (
+            _streamed_provider_fields()["server_side_tool_invocations"]
+            == _provider_fields(unary)["server_side_tool_invocations"]
+        )
+
+    def test_next_turn_sends_the_call_with_its_own_args_and_signature(self):
+        contents = gemini_convert_messages_with_history(
+            [
+                _QUESTION,
+                {"role": "assistant", "content": _ANSWER_TEXT, "provider_specific_fields": _streamed_provider_fields()},
+                {"role": "user", "content": "And tomorrow?"},
+            ]
+        )
+
+        parts = [part for content in contents if content.get("role") == "model" for part in content["parts"]]
+        assert [part for part in parts if "toolCall" in part] == [_SEARCH_CALL]
+        assert [part.get("thoughtSignature") for part in parts if "toolResponse" in part] == ["sig_resp"]
+
+
 # --- Input re-injection tests ---