fix: Fix branch parsing and unify branch stamping

- Fix branch parsing bug in should_pause_invocation to correctly handle dot-separated nested branch paths and avoid substring collisions.
- Unify branch stamping in runners.py by replacing manual branch matching with InvocationContext.stamp_event_branch_context().

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 938841304

Author: DeanChensj

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/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'