fix(live): keep streaming tool yields from completing turns

Fixes #5947.

Co-authored-by: George Weale <gweale@google.com>
COPYBARA_INTEGRATE_REVIEW=#6009 from he-yufeng:fix/live-stream-turn-complete f72e0b5
PiperOrigin-RevId: 943418613

Author: 2 people

src/google/adk/agents/live_request_queue.py

@@ -55,6 +55,9 @@ class LiveRequest(BaseModel):
   close: bool = False
   """If set, close the queue. queue.shutdown() is only supported in Python 3.13+."""
 
+  partial: bool = False
+  """If set, the content is a partial turn update that does not complete the current model turn."""
+
 
 class LiveRequestQueue:
   """Queue used to send LiveRequest in a live(bidirectional streaming) way."""
@@ -65,8 +68,8 @@ def __init__(self):
   def close(self):
     self._queue.put_nowait(LiveRequest(close=True))
 
-  def send_content(self, content: types.Content):
-    self._queue.put_nowait(LiveRequest(content=content))
+  def send_content(self, content: types.Content, partial: bool = False):
+    self._queue.put_nowait(LiveRequest(content=content, partial=partial))
 
   def send_realtime(self, blob: types.Blob):
     self._queue.put_nowait(LiveRequest(blob=blob))

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -811,9 +811,9 @@ async def _send_to_model(
         is_function_response = content.parts and any(
             part.function_response for part in content.parts
         )
-        if not is_function_response:
-          if not content.role:
-            content.role = 'user'
+        if not is_function_response and not content.role:
+          content.role = 'user'
+        if not is_function_response and not live_request.partial:
           user_content_event = Event(
               id=Event.new_id(),
               invocation_id=invocation_context.invocation_id,
@@ -824,7 +824,9 @@ async def _send_to_model(
               session=invocation_context.session,
               event=user_content_event,
           )
-        await llm_connection.send_content(live_request.content)
+        await llm_connection._send_content(
+            live_request.content, partial=live_request.partial
+        )
 
   async def _receive_from_model(
       self,

src/google/adk/flows/llm_flows/functions.py

@@ -995,7 +995,9 @@ async def run_tool_and_update_queue(tool, function_args, tool_context):
             updated_content = _build_function_response_content(
                 tool, result, tool_context.function_call_id
             )
-            invocation_context.live_request_queue.send_content(updated_content)
+            invocation_context.live_request_queue.send_content(
+                updated_content, partial=True
+            )
       except asyncio.CancelledError:
         raise  # Re-raise to properly propagate the cancellation
 

src/google/adk/models/base_llm_connection.py

@@ -51,6 +51,21 @@ async def send_content(self, content: types.Content):
     """
     pass
 
+  async def _send_content(
+      self, content: types.Content, *, partial: bool = False
+  ) -> None:
+    """Sends content, optionally as a partial (non-turn-completing) update.
+
+    The default implementation ignores ``partial`` and completes the turn.
+    Connections that support turn-based partial updates override this.
+
+    Args:
+      content: The content to send to the model.
+      partial: Whether this content is a partial turn update that does not
+        complete the model turn.
+    """
+    await self.send_content(content)
+
   @abstractmethod
   async def send_realtime(self, blob: types.Blob):
     """Sends a chunk of audio or a frame of video to the model in realtime.

src/google/adk/models/gemini_llm_connection.py

@@ -104,6 +104,18 @@ async def send_content(self, content: types.Content):
     Args:
       content: The content to send to the model.
     """
+    await self._send_content(content)
+
+  async def _send_content(
+      self, content: types.Content, *, partial: bool = False
+  ) -> None:
+    """Sends content, optionally as a partial (non-turn-completing) update.
+
+    Args:
+      content: The content to send to the model.
+      partial: Whether this content is a partial turn update that does not
+        complete the model turn.
+    """
     assert content.parts
     if content.parts[0].function_response:
       # All parts have to be function responses.
@@ -115,7 +127,8 @@ async def send_content(self, content: types.Content):
     else:
       logger.debug('Sending LLM new content %s', content)
       if (
-          self._is_gemini_3_x_live
+          not partial
+          and self._is_gemini_3_x_live
           and len(content.parts) == 1
           and content.parts[0].text
       ):
@@ -127,7 +140,7 @@ async def send_content(self, content: types.Content):
         await self._gemini_session.send(
             input=types.LiveClientContent(
                 turns=[content],
-                turn_complete=True,
+                turn_complete=not partial,
             )
         )
 

tests/unittests/agents/test_live_request_queue.py

@@ -40,6 +40,17 @@ def test_send_content():
     mock_put_nowait.assert_called_once_with(LiveRequest(content=content))
 
 
+def test_send_content_sets_partial():
+  queue = LiveRequestQueue()
+  content = MagicMock(spec=types.Content)
+
+  with patch.object(queue._queue, "put_nowait") as mock_put_nowait:
+    queue.send_content(content, partial=True)
+    mock_put_nowait.assert_called_once_with(
+        LiveRequest(content=content, partial=True)
+    )
+
+
 def test_send_realtime():
   queue = LiveRequestQueue()
   blob = MagicMock(spec=types.Blob)

tests/unittests/flows/llm_flows/test_base_llm_flow_realtime.py

@@ -197,5 +197,36 @@ async def test_send_to_model_with_text_content(mock_llm_connection):
   await flow._send_to_model(mock_llm_connection, invocation_context)
 
   # Verify send_content was called instead of send_realtime
-  mock_llm_connection.send_content.assert_called_once_with(content)
+  mock_llm_connection._send_content.assert_called_once_with(
+      content, partial=False
+  )
   mock_llm_connection.send_realtime.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_send_to_model_with_intermediate_text_content(
+    mock_llm_connection,
+):
+  agent = Agent(name='test_agent', model='mock')
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content=''
+  )
+  invocation_context.live_request_queue = LiveRequestQueue()
+  invocation_context.session_service.append_event = mock.AsyncMock()
+
+  flow = TestBaseLlmFlow()
+
+  content = types.Content(
+      role='user', parts=[types.Part.from_text(text='progress')]
+  )
+  invocation_context.live_request_queue.send(
+      LiveRequest(content=content, partial=True)
+  )
+  invocation_context.live_request_queue.close()
+
+  await flow._send_to_model(mock_llm_connection, invocation_context)
+
+  mock_llm_connection._send_content.assert_called_once_with(
+      content, partial=True
+  )
+  invocation_context.session_service.append_event.assert_not_called()

tests/unittests/models/test_gemini_llm_connection.py

@@ -124,6 +124,22 @@ async def test_send_content_text(gemini_connection, mock_gemini_session):
   assert call_args['input'].turn_complete is True
 
 
+@pytest.mark.asyncio
+async def test_send_content_text_can_keep_turn_open(
+    gemini_connection, mock_gemini_session
+):
+  content = types.Content(
+      role='user', parts=[types.Part.from_text(text='progress')]
+  )
+
+  await gemini_connection._send_content(content, partial=True)
+
+  mock_gemini_session.send.assert_called_once()
+  call_args = mock_gemini_session.send.call_args[1]
+  assert call_args['input'].turns == [content]
+  assert call_args['input'].turn_complete is False
+
+
 @pytest.mark.asyncio
 async def test_send_content_function_response(
     gemini_connection, mock_gemini_session