fix(models): Prevent grounding metadata loss in Gemini 3.1

In Gemini 3.1 Live, grounding metadata is often emitted early alongside partial chunks or immediately after tool execution, causing the runner to discard it because it is associated with partial=True events.

- Retain pending grounding metadata during the turn in GeminiLlmConnection.receive and yield it on the final non-partial event.
- Add test_receive_grounding_metadata_pending unit test to verify the fix.

Change-Id: I4d8ab1f231a3fed63c4375cbfcd8831d6535721e

Author: wukath

src/google/adk/models/gemini_llm_connection.py

@@ -159,7 +159,12 @@ async def send_realtime(self, input: RealtimeInput):
     else:
       raise ValueError('Unsupported input type: %s' % type(input))
 
-  def __build_full_text_response(self, text: str, is_thought: bool = False):
+  def __build_full_text_response(
+      self,
+      text: str,
+      is_thought: bool = False,
+      grounding_metadata: types.GroundingMetadata | None = None,
+  ):
     """Builds a full text response.
 
     The text should not be partial and the returned LlmResponse is not
@@ -168,6 +173,7 @@ def __build_full_text_response(self, text: str, is_thought: bool = False):
     Args:
       text: The text to be included in the response.
       is_thought: Whether the text is a thought.
+      grounding_metadata: The grounding metadata to include.
 
     Returns:
       An LlmResponse containing the full text.
@@ -180,6 +186,7 @@ def __build_full_text_response(self, text: str, is_thought: bool = False):
             role='model',
             parts=[part],
         ),
+        grounding_metadata=grounding_metadata,
         partial=False,
         live_session_id=self._gemini_session.session_id,
     )
@@ -194,6 +201,7 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
     text = ''
     is_thought = False
     tool_call_parts = []
+    pending_grounding_metadata = None
     async with Aclosing(self._gemini_session.receive()) as agen:
       # TODO(b/440101573): Reuse StreamingResponseAggregator to accumulate
       # partial content and emit responses as needed.
@@ -209,6 +217,10 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
           )
         if message.server_content:
           content = message.server_content.model_turn
+          if message.server_content.grounding_metadata:
+            pending_grounding_metadata = (
+                message.server_content.grounding_metadata
+            )
 
           # Standalone grounding_metadata event (when content is empty)
           if (
@@ -338,10 +350,14 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
               )
               self._output_transcription_text = ''
           if message.server_content.turn_complete:
+            g_metadata_to_yield = pending_grounding_metadata
             if text:
-              yield self.__build_full_text_response(text, is_thought)
+              yield self.__build_full_text_response(
+                  text, is_thought, g_metadata_to_yield
+              )
               text = ''
               is_thought = False
+              g_metadata_to_yield = None
             if tool_call_parts:
               logger.debug('Returning aggregated tool_call_parts')
               yield LlmResponse(
@@ -353,7 +369,8 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
             yield LlmResponse(
                 turn_complete=True,
                 interrupted=message.server_content.interrupted,
-                grounding_metadata=message.server_content.grounding_metadata,
+                grounding_metadata=message.server_content.grounding_metadata
+                or g_metadata_to_yield,
                 model_version=self._model_version,
                 live_session_id=live_session_id,
             )

tests/unittests/models/test_gemini_llm_connection.py

@@ -1343,3 +1343,70 @@ async def mock_receive_generator():
   content_response = next((r for r in responses if r.content), None)
   assert content_response is not None
   assert content_response.content == mock_content
+
+
+@pytest.mark.asyncio
+async def test_receive_grounding_metadata_pending(
+    gemini_connection, mock_gemini_session
+):
+  """Test that grounding metadata in partial chunks is pending and yielded on full text."""
+  grounding_metadata = types.GroundingMetadata(
+      web_search_queries=['stock price of google'],
+  )
+
+  def make_msg(text=None, g_meta=None, tc=False):
+    msg = mock.Mock(
+        usage_metadata=None,
+        tool_call=None,
+        session_resumption_update=None,
+        go_away=None,
+    )
+    msg.server_content = mock.Mock(
+        interrupted=False,
+        input_transcription=None,
+        output_transcription=None,
+        generation_complete=False,
+        turn_complete=tc,
+        grounding_metadata=g_meta,
+        model_turn=types.Content(
+            role='model', parts=[types.Part.from_text(text=text)]
+        )
+        if text
+        else None,
+    )
+    return msg
+
+  msg1 = make_msg(text='hello', g_meta=grounding_metadata)
+  msg2 = make_msg(text=' world')
+  msg3 = make_msg(tc=True)
+
+  async def gen():
+    yield msg1
+    yield msg2
+    yield msg3
+
+  mock_gemini_session.receive = mock.Mock(return_value=gen())
+
+  responses = [resp async for resp in gemini_connection.receive()]
+
+  # Expected responses:
+  # 1. Msg 1 partial (hello) with grounding_metadata
+  # 2. Msg 2 partial ( world) without grounding_metadata
+  # 3. Full text response (hello world) with PENDING grounding_metadata
+  # 4. Turn complete response without grounding_metadata (already cleared)
+  assert len(responses) == 4
+
+  assert responses[0].content.parts[0].text == 'hello'
+  assert responses[0].partial is True
+  assert responses[0].grounding_metadata == grounding_metadata
+
+  assert responses[1].content.parts[0].text == ' world'
+  assert responses[1].partial is True
+  assert responses[1].grounding_metadata is None
+
+  assert responses[2].content.parts[0].text == 'hello world'
+  assert responses[2].partial is False
+  assert responses[2].grounding_metadata == grounding_metadata
+
+  assert responses[3].turn_complete is True
+  assert responses[3].grounding_metadata is None