Address Gemini Code Assist review feedback

- Fix critical bug: Remove premature reset of last_grounding_metadata before turn_complete response to prevent data loss

- Simplify duplicate reset logic in interrupted handling

- Add grounding_metadata propagation to tool_call responses

- Add test for grounding_metadata with text content + turn_complete

- Add test for grounding_metadata with tool_call responses

All 27 tests pass.

Author: VedantMadane

src/google/adk/models/gemini_llm_connection.py

@@ -294,7 +294,6 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
             if text:
               yield self.__build_full_text_response(text, last_grounding_metadata)
               text = ''
-              last_grounding_metadata = None
             yield LlmResponse(
                 turn_complete=True,
                 interrupted=message.server_content.interrupted,
@@ -310,23 +309,25 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
             if text:
               yield self.__build_full_text_response(text, last_grounding_metadata)
               text = ''
-              last_grounding_metadata = None
             else:
               yield LlmResponse(
                   interrupted=message.server_content.interrupted,
                   grounding_metadata=last_grounding_metadata,
               )
-              last_grounding_metadata = None  # Reset after yielding
+            last_grounding_metadata = None  # Reset after yielding
         if message.tool_call:
           if text:
             yield self.__build_full_text_response(text, last_grounding_metadata)
             text = ''
-            last_grounding_metadata = None
           parts = [
               types.Part(function_call=function_call)
               for function_call in message.tool_call.function_calls
           ]
-          yield LlmResponse(content=types.Content(role='model', parts=parts))
+          yield LlmResponse(
+              content=types.Content(role='model', parts=parts),
+              grounding_metadata=last_grounding_metadata,
+          )
+          last_grounding_metadata = None  # Reset after yielding
         if message.session_resumption_update:
           logger.debug('Received session resumption message: %s', message)
           yield (

tests/unittests/models/test_gemini_llm_connection.py

@@ -887,3 +887,107 @@ async def mock_receive_generator():
   assert turn_complete_response is not None
   # The grounding_metadata should be carried over to turn_complete
   assert turn_complete_response.grounding_metadata == mock_grounding_metadata
+
+
+@pytest.mark.asyncio
+async def test_receive_grounding_metadata_with_text_and_turn_complete(
+    gemini_connection, mock_gemini_session
+):
+  """Test that grounding_metadata is preserved when text content is followed by turn_complete."""
+  mock_content = types.Content(
+      role='model', parts=[types.Part.from_text(text='response text')]
+  )
+  mock_grounding_metadata = types.GroundingMetadata(
+      retrieval_queries=['test query'],
+  )
+
+  # Message with both content and grounding, followed by turn_complete
+  mock_server_content = mock.Mock()
+  mock_server_content.model_turn = mock_content
+  mock_server_content.interrupted = False
+  mock_server_content.input_transcription = None
+  mock_server_content.output_transcription = None
+  mock_server_content.turn_complete = True
+  mock_server_content.generation_complete = False
+  mock_server_content.grounding_metadata = mock_grounding_metadata
+
+  mock_message = mock.Mock()
+  mock_message.usage_metadata = None
+  mock_message.server_content = mock_server_content
+  mock_message.tool_call = None
+  mock_message.session_resumption_update = None
+
+  async def mock_receive_generator():
+    yield mock_message
+
+  receive_mock = mock.Mock(return_value=mock_receive_generator())
+  mock_gemini_session.receive = receive_mock
+
+  responses = [resp async for resp in gemini_connection.receive()]
+
+  # Find content response with grounding
+  content_response = next((r for r in responses if r.content), None)
+  assert content_response is not None
+  assert content_response.grounding_metadata == mock_grounding_metadata
+
+  # Find turn_complete response - should also have grounding_metadata
+  turn_complete_response = next((r for r in responses if r.turn_complete), None)
+  assert turn_complete_response is not None
+  assert turn_complete_response.grounding_metadata == mock_grounding_metadata
+
+
+@pytest.mark.asyncio
+async def test_receive_grounding_metadata_with_tool_call(
+    gemini_connection, mock_gemini_session
+):
+  """Test that grounding_metadata is propagated with tool_call responses."""
+  mock_grounding_metadata = types.GroundingMetadata(
+      retrieval_queries=['test query'],
+  )
+
+  # First message with grounding metadata
+  mock_server_content1 = mock.Mock()
+  mock_server_content1.model_turn = None
+  mock_server_content1.interrupted = False
+  mock_server_content1.input_transcription = None
+  mock_server_content1.output_transcription = None
+  mock_server_content1.turn_complete = False
+  mock_server_content1.generation_complete = False
+  mock_server_content1.grounding_metadata = mock_grounding_metadata
+
+  message1 = mock.Mock()
+  message1.usage_metadata = None
+  message1.server_content = mock_server_content1
+  message1.tool_call = None
+  message1.session_resumption_update = None
+
+  # Second message with tool_call
+  mock_function_call = types.FunctionCall(
+      name='test_function', args={'param': 'value'}
+  )
+  mock_tool_call = mock.Mock()
+  mock_tool_call.function_calls = [mock_function_call]
+
+  message2 = mock.Mock()
+  message2.usage_metadata = None
+  message2.server_content = None
+  message2.tool_call = mock_tool_call
+  message2.session_resumption_update = None
+
+  async def mock_receive_generator():
+    yield message1
+    yield message2
+
+  receive_mock = mock.Mock(return_value=mock_receive_generator())
+  mock_gemini_session.receive = receive_mock
+
+  responses = [resp async for resp in gemini_connection.receive()]
+
+  # Find tool_call response
+  tool_call_response = next(
+      (r for r in responses if r.content and r.content.parts[0].function_call),
+      None,
+  )
+  assert tool_call_response is not None
+  # The grounding_metadata should be carried over to tool_call
+  assert tool_call_response.grounding_metadata == mock_grounding_metadata