Python: Honor the configured model in streamed Gemini chunks

python/packages/gemini/agent_framework_gemini/_chat_client.py

@@ -492,8 +492,13 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]:
                     # The SDK stream owns the HTTP response; close it as soon as the consumer stops early
                     # instead of leaving the response open until the generator is garbage collected.
                     try:
+                        # Once the server names a version, later chunks that omit it keep that version
+                        # instead of reverting to the requested alias.
+                        fallback_model = model
                         async for chunk in response_stream:
-                            yield self._process_chunk(chunk)
+                            if chunk.model_version:
+                                fallback_model = chunk.model_version
+                            yield self._process_chunk(chunk, fallback_model)
                     finally:
                         aclose = getattr(response_stream, "aclose", None)
                         if aclose is not None:
@@ -1071,13 +1076,15 @@ def _process_generate_response(
             raw_representation=response,
         )
 
-    def _process_chunk(self, chunk: types.GenerateContentResponse) -> ChatResponseUpdate:
+    def _process_chunk(self, chunk: types.GenerateContentResponse, model: str) -> ChatResponseUpdate:
         """Convert a single streaming chunk to a framework ChatResponseUpdate.
 
         Usage details are attached only to the final chunk, identified by a non-None finish reason.
 
         Args:
             chunk: A streaming ``GenerateContentResponse`` chunk from the Gemini API.
+            model: The model to report when the chunk carries no model version: the last version the
+                server reported in this stream, or the model resolved for this request.
 
         Returns:
             A ``ChatResponseUpdate`` with parsed contents, finish reason, and model ID.
@@ -1096,7 +1103,7 @@ def _process_chunk(self, chunk: types.GenerateContentResponse) -> ChatResponseUp
 
         return ChatResponseUpdate(
             contents=contents,
-            model=chunk.model_version,
+            model=chunk.model_version or model,
             finish_reason=finish_reason,
             raw_representation=chunk,
         )

python/packages/gemini/tests/test_gemini_client.py

@@ -547,6 +547,74 @@ async def test_get_response_model_from_response() -> None:
     assert response.model == "gemini-2.5-pro-002"
 
 
+async def test_get_response_streaming_model_from_chunk() -> None:
+    """Populates the update model from the model_version field on each chunk."""
+    client, mock = _make_gemini_client()
+    chunks = [_make_response([_make_part(text="Hi")], model_version="gemini-2.5-pro-002")]
+    mock.aio.models.generate_content_stream = AsyncMock(return_value=_async_iter(chunks))
+
+    stream = client.get_response(
+        messages=[Message(role="user", contents=[Content.from_text("Hi")])],
+        stream=True,
+    )
+    updates = [update async for update in stream]
+
+    assert updates[0].model == "gemini-2.5-pro-002"
+
+
+async def test_get_response_streaming_model_falls_back_to_configured_model() -> None:
+    """Falls back to the configured model when a chunk carries no model_version."""
+    client, mock = _make_gemini_client()
+    chunk = _make_response([_make_part(text="Hi")])
+    chunk.model_version = None
+    mock.aio.models.generate_content_stream = AsyncMock(return_value=_async_iter([chunk]))
+
+    stream = client.get_response(
+        messages=[Message(role="user", contents=[Content.from_text("Hi")])],
+        stream=True,
+    )
+    updates = [update async for update in stream]
+
+    assert updates[0].model == "gemini-2.5-flash"
+
+
+async def test_get_response_streaming_model_falls_back_to_requested_model() -> None:
+    """Falls back to the model requested in options, not the client default."""
+    client, mock = _make_gemini_client(model="gemini-2.5-flash")
+    chunk = _make_response([_make_part(text="Hi")])
+    chunk.model_version = None
+    mock.aio.models.generate_content_stream = AsyncMock(return_value=_async_iter([chunk]))
+
+    stream = client.get_response(
+        messages=[Message(role="user", contents=[Content.from_text("Hi")])],
+        options={"model": "gemini-2.5-pro"},
+        stream=True,
+    )
+    updates = [update async for update in stream]
+
+    assert updates[0].model == "gemini-2.5-pro"
+
+
+async def test_get_response_streaming_keeps_model_version_seen_earlier() -> None:
+    """A later chunk without model_version must not replace a version already seen in the stream."""
+    client, mock = _make_gemini_client(model="gemini-2.5-flash")
+    first = _make_response([_make_part(text="Hi")], finish_reason=None, model_version="gemini-2.5-pro-002")
+    last = _make_response([_make_part(text="!")], finish_reason="STOP")
+    last.model_version = None
+    mock.aio.models.generate_content_stream = AsyncMock(return_value=_async_iter([first, last]))
+
+    stream = client.get_response(
+        messages=[Message(role="user", contents=[Content.from_text("Hi")])],
+        options={"model": "gemini-2.5-pro"},
+        stream=True,
+    )
+    updates = [update async for update in stream]
+    final = await stream.get_final_response()
+
+    assert [update.model for update in updates] == ["gemini-2.5-pro-002", "gemini-2.5-pro-002"]
+    assert final.model == "gemini-2.5-pro-002"
+
+
 async def test_get_response_uses_model_from_options() -> None:
     """Uses the model specified in options, overriding the client's default."""
     client, mock = _make_gemini_client(model="gemini-2.5-flash")