Merge branch 'main' into fix/2942-jinja2-instruction-templating

Author: vaibhav-patel

src/google/adk/agents/invocation_context.py

@@ -31,6 +31,7 @@
 from ..artifacts.base_artifact_service import BaseArtifactService
 from ..auth.auth_credential import AuthCredential
 from ..auth.credential_service.base_credential_service import BaseCredentialService
+from ..events._branch_path import _BranchPath
 from ..events.event import Event
 from ..memory.base_memory_service import BaseMemoryService
 from ..plugins.plugin_manager import PluginManager
@@ -466,9 +467,28 @@ def should_pause_invocation(self, event: Event) -> bool:
     if not event.long_running_tool_ids or not event.get_function_calls():
       return False
 
+    events = self.session.events if self.session else []
     for fc in event.get_function_calls():
       if fc.id in event.long_running_tool_ids:
-        return True
+        # Check if there is a newer user event in the session that belongs to a sub-branch of this tool call.
+        # This indicates the tool call is resuming to process that nested input.
+        is_resolving_sub_branch = False
+        event_index = -1
+        # Search backwards since the checked event is typically near the end of history.
+        for i in range(len(events) - 1, -1, -1):
+          if events[i].id == event.id:
+            event_index = i
+            break
+        if event_index != -1:
+          is_resolving_sub_branch = any(
+              e.author == "user"
+              and e.branch
+              and fc.id in _BranchPath.from_string(e.branch).run_ids
+              for e in events[event_index + 1 :]
+          )
+
+        if not is_resolving_sub_branch:
+          return True
 
     return False
 
@@ -483,11 +503,23 @@ def _find_matching_function_call(
     if not function_responses:
       return None
 
-    # Search backwards from the event before the current response event.
+    events = self._get_events(current_invocation=True)
+    if events and events[-1].id == function_response_event.id:
+      search_space = events[:-1]
+    else:
+      search_space = events
+
     return find_event_by_function_call_id(
-        self._get_events(current_invocation=True)[:-1], function_responses[0].id
+        search_space, function_responses[0].id
     )
 
+  def stamp_event_branch_context(self, event: Event) -> None:
+    """Stamps the event with the branch and isolation scope of its matching function call."""
+    if function_call := self._find_matching_function_call(event):
+      event.branch = function_call.branch
+      if function_call.isolation_scope is not None:
+        event.isolation_scope = function_call.isolation_scope
+
 
 def new_invocation_context_id() -> str:
   return "e-" + cast(str, platform_uuid.new_uuid())

src/google/adk/models/gemini_llm_connection.py

@@ -86,33 +86,10 @@ async def send_history(self, history: list[types.Content]):
     ]
 
     if contents:
-      # Gemini Enterprise Agent Platform does not support history_config in the
-      # SDK. To initialize a live session with prior history without hitting a
-      # 1007 protocol error (invalid role mid-session), we consolidate previous
-      # multi-turn interactions into a unified contextual preamble on a single
-      # user role turn.
-      if (
-          self._is_gemini_3_1_flash_live
-          and self._api_backend != GoogleLLMVariant.GEMINI_API
-      ):
-        collapsed_text = 'Previous conversation history:\n'
-        for c in contents:
-          text_parts = ''.join(p.text for p in c.parts if p.text)
-          collapsed_text += f'[{c.role}]: {text_parts}\n'
-        contents = [
-            types.Content(
-                role='user', parts=[types.Part.from_text(text=collapsed_text)]
-            )
-        ]
-
       logger.debug('Sending history to live connection: %s', contents)
       await self._gemini_session.send_client_content(
           turns=contents,
-          turn_complete=(
-              True
-              if self._is_gemini_3_1_flash_live
-              else contents[-1].role == 'user'
-          ),
+          turn_complete=contents[-1].role == 'user',
       )
     else:
       logger.info('no content is sent')

src/google/adk/runners.py

@@ -740,6 +740,7 @@ async def _append_user_event(
       if iso is not None:
         event.isolation_scope = iso
     _apply_run_config_custom_metadata(event, ic.run_config)
+    ic.stamp_event_branch_context(event)
     return await self.session_service.append_event(
         session=ic.session, event=event
     )
@@ -1482,10 +1483,7 @@ async def _append_new_message_to_session(
           content=new_message,
       )
     _apply_run_config_custom_metadata(event, invocation_context.run_config)
-    # If new_message is a function response, find the matching function call
-    # and use its branch as the new event's branch.
-    if function_call := invocation_context._find_matching_function_call(event):
-      event.branch = function_call.branch
+    invocation_context.stamp_event_branch_context(event)
 
     await self.session_service.append_event(
         session=invocation_context.session, event=event

tests/unittests/agents/test_invocation_context.py

@@ -152,14 +152,14 @@ def event_to_pause(self, long_running_function_call) -> Event:
     )
 
   def _create_test_invocation_context(
-      self, resumability_config
+      self, resumability_config: ResumabilityConfig | None = None
   ) -> InvocationContext:
     """Create a mock invocation context for testing."""
     ctx = InvocationContext(
         session_service=Mock(spec=BaseSessionService),
         agent=Mock(spec=BaseAgent),
         invocation_id='inv_1',
-        session=Mock(spec=Session),
+        session=Mock(spec=Session, events=[]),
         resumability_config=resumability_config,
     )
     return ctx
@@ -208,6 +208,69 @@ def test_should_not_pause_invocation_with_no_function_calls(
         nonpausable_event
     )
 
+  def test_should_not_pause_when_user_resumes_in_sub_branch(
+      self, event_to_pause, long_running_function_call
+  ):
+    """We do not pause the invocation if a subsequent user event belongs to a sub-branch."""
+    # Arrange
+    mock_invocation_context = self._create_test_invocation_context()
+    user_event = Event(
+        invocation_id='inv_1',
+        author='user',
+        branch=f'agent@{long_running_function_call.id}.child',
+    )
+    mock_invocation_context.session.events = [event_to_pause, user_event]
+
+    # Act
+    should_pause = mock_invocation_context.should_pause_invocation(
+        event_to_pause
+    )
+
+    # Assert
+    assert not should_pause
+
+  def test_should_not_pause_when_user_resumes_in_deeply_nested_sub_branch(
+      self, event_to_pause, long_running_function_call
+  ):
+    """We do not pause if the user resumes in a deeply nested sub-branch containing the tool call."""
+    # Arrange
+    mock_invocation_context = self._create_test_invocation_context()
+    user_event = Event(
+        invocation_id='inv_1',
+        author='user',
+        branch=f'parent@other.child@{long_running_function_call.id}.grandchild',
+    )
+    mock_invocation_context.session.events = [event_to_pause, user_event]
+
+    # Act
+    should_pause = mock_invocation_context.should_pause_invocation(
+        event_to_pause
+    )
+
+    # Assert
+    assert not should_pause
+
+  def test_should_pause_when_user_resumes_in_different_branch(
+      self, event_to_pause
+  ):
+    """We still pause the invocation if the subsequent user event belongs to a different branch."""
+    # Arrange
+    mock_invocation_context = self._create_test_invocation_context()
+    user_event = Event(
+        invocation_id='inv_1',
+        author='user',
+        branch='parent@different_id.child',
+    )
+    mock_invocation_context.session.events = [event_to_pause, user_event]
+
+    # Act
+    should_pause = mock_invocation_context.should_pause_invocation(
+        event_to_pause
+    )
+
+    # Assert
+    assert should_pause
+
   def test_is_resumable_true(self):
     """Tests that is_resumable is True when resumability is enabled."""
     invocation_context = self._create_test_invocation_context(
@@ -534,3 +597,32 @@ def test_find_matching_function_call_no_response_in_event(
     invocation_context = test_invocation_context([fc_event, fr_event])
     match = invocation_context._find_matching_function_call(fr_event_no_fr)
     assert match is None
+
+  def test_stamp_event_branch_context_preserves_isolation_scope(
+      self, test_invocation_context
+  ):
+    """Tests stamp_event_branch_context does not overwrite existing isolation_scope with None."""
+    fc = Part.from_function_call(name='some_tool', args={})
+    fc.function_call.id = 'test_function_call_id'
+    fc_event = Event(
+        invocation_id='inv_1',
+        author='agent',
+        branch='root@1',
+        isolation_scope=None,  # Coordinator FC has None scope
+        content=testing_utils.ModelContent([fc]),
+    )
+    fr = Part.from_function_response(
+        name='some_tool', response={'result': 'ok'}
+    )
+    fr.function_response.id = 'test_function_call_id'
+    fr_event = Event(
+        invocation_id='inv_1',
+        author='agent',
+        isolation_scope='task_123',  # Pre-populated active task scope
+        content=Content(role='user', parts=[fr]),
+    )
+    invocation_context = test_invocation_context([fc_event, fr_event])
+
+    invocation_context.stamp_event_branch_context(fr_event)
+    assert fr_event.branch == 'root@1'
+    assert fr_event.isolation_scope == 'task_123'

tests/unittests/models/test_gemini_llm_connection.py

@@ -949,54 +949,128 @@ async def test_send_history_filters_various_audio_mime_types(
 
 @pytest.mark.asyncio
 async def test_send_history_gemini_31_turn_complete(mock_gemini_session):
-  """Verify Gemini 3.1 Live history seeding explicitly appends turn_complete=True."""
+  """Verify Gemini 3.1 Live history seeding sets turn_complete based on history[-1].role == 'user'."""
   conn = GeminiLlmConnection(
       mock_gemini_session,
       api_backend=GoogleLLMVariant.GEMINI_API,
       model_version='gemini-3.1-flash-live-preview',
   )
   mock_gemini_session.send_client_content = mock.AsyncMock()
 
-  mock_contents = [
+  # Last turn is model -> turn_complete=False
+  mock_contents_model = [
       types.Content(role='user', parts=[types.Part.from_text(text='hi')]),
       types.Content(role='model', parts=[types.Part.from_text(text='hello')]),
   ]
-  await conn.send_history(mock_contents)
+  await conn.send_history(mock_contents_model)
 
   mock_gemini_session.send_client_content.assert_called_once_with(
-      turns=mock_contents,
+      turns=mock_contents_model,
+      turn_complete=False,
+  )
+
+  # Last turn is user -> turn_complete=True
+  mock_gemini_session.send_client_content.reset_mock()
+  mock_contents_user = [
+      types.Content(role='user', parts=[types.Part.from_text(text='hi')]),
+  ]
+  await conn.send_history(mock_contents_user)
+
+  mock_gemini_session.send_client_content.assert_called_once_with(
+      turns=mock_contents_user,
       turn_complete=True,
   )
 
 
 @pytest.mark.asyncio
-async def test_send_history_collapse_vertex_ai(mock_gemini_session):
-  """Verify history prompt collapse when seeding Gemini 3.1 Live on Vertex AI backend."""
+async def test_send_history_vertex_ai_no_collapse(mock_gemini_session):
+  """Verify history is sent without collapsing on Vertex AI backend."""
   conn = GeminiLlmConnection(
       mock_gemini_session,
       api_backend=GoogleLLMVariant.VERTEX_AI,
       model_version='gemini-3.1-flash-live-preview',
   )
   mock_gemini_session.send_client_content = mock.AsyncMock()
 
-  mock_contents = [
+  # Last turn is model -> turn_complete=False
+  mock_contents_model = [
       types.Content(role='user', parts=[types.Part.from_text(text='hi')]),
       types.Content(role='model', parts=[types.Part.from_text(text='hello')]),
   ]
-  await conn.send_history(mock_contents)
+  await conn.send_history(mock_contents_model)
 
-  assert mock_gemini_session.send_client_content.call_count == 1
-  called_turns = mock_gemini_session.send_client_content.call_args.kwargs[
-      'turns'
+  mock_gemini_session.send_client_content.assert_called_once_with(
+      turns=mock_contents_model,
+      turn_complete=False,
+  )
+
+  # Last turn is user -> turn_complete=True
+  mock_gemini_session.send_client_content.reset_mock()
+  mock_contents_user = [
+      types.Content(role='user', parts=[types.Part.from_text(text='hi')]),
+      types.Content(role='model', parts=[types.Part.from_text(text='hello')]),
+      types.Content(
+          role='user', parts=[types.Part.from_text(text='how are you?')]
+      ),
   ]
-  assert len(called_turns) == 1
-  assert called_turns[0].role == 'user'
-  assert 'Previous conversation history:' in called_turns[0].parts[0].text
-  assert '[user]: hi' in called_turns[0].parts[0].text
-  assert '[model]: hello' in called_turns[0].parts[0].text
-  assert (
-      mock_gemini_session.send_client_content.call_args.kwargs['turn_complete']
-      is True
+  await conn.send_history(mock_contents_user)
+
+  mock_gemini_session.send_client_content.assert_called_once_with(
+      turns=mock_contents_user,
+      turn_complete=True,
+  )
+
+
+@pytest.mark.asyncio
+async def test_send_history_turn_complete_determined_by_filtered_content(
+    mock_gemini_session,
+):
+  """Verify turn_complete is determined by the last element of filtered content instead of unfiltered history."""
+  conn = GeminiLlmConnection(
+      mock_gemini_session,
+      api_backend=GoogleLLMVariant.GEMINI_API,
+      model_version='gemini-3.1-flash-live-preview',
+  )
+  mock_gemini_session.send_client_content = mock.AsyncMock()
+
+  # Scenario: Last turn in history is a user audio turn (gets filtered out).
+  # The remaining last turn is model's turn -> turn_complete should be False.
+  audio_part = types.Part(
+      inline_data=types.Blob(data=b'\x00\xFF', mime_type='audio/pcm')
+  )
+  history_with_final_audio_user_turn = [
+      types.Content(role='user', parts=[types.Part.from_text(text='hi')]),
+      types.Content(role='model', parts=[types.Part.from_text(text='hello')]),
+      types.Content(role='user', parts=[audio_part]),
+  ]
+
+  await conn.send_history(history_with_final_audio_user_turn)
+
+  expected_contents = [
+      types.Content(role='user', parts=[types.Part.from_text(text='hi')]),
+      types.Content(role='model', parts=[types.Part.from_text(text='hello')]),
+  ]
+  mock_gemini_session.send_client_content.assert_called_once_with(
+      turns=expected_contents,
+      turn_complete=False,
+  )
+
+  # Scenario: Last turn in history is a model audio turn (gets filtered out).
+  # The remaining last turn is user's turn -> turn_complete should be True.
+  mock_gemini_session.send_client_content.reset_mock()
+  history_with_final_audio_model_turn = [
+      types.Content(role='user', parts=[types.Part.from_text(text='hi')]),
+      types.Content(role='model', parts=[audio_part]),
+  ]
+
+  await conn.send_history(history_with_final_audio_model_turn)
+
+  expected_contents = [
+      types.Content(role='user', parts=[types.Part.from_text(text='hi')]),
+  ]
+  mock_gemini_session.send_client_content.assert_called_once_with(
+      turns=expected_contents,
+      turn_complete=True,
   )