feat: Add sub-branch creation and segment appending to _BranchPath

Introduce `append` method and `create_sub_branch` classmethod on `_BranchPath`
to standardize dynamic branch path creation across workflow node runners, agent tools,
and parallel agent execution.

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

Author: DeanChensj

src/google/adk/agents/parallel_agent.py

@@ -25,6 +25,7 @@
 from typing_extensions import deprecated
 from typing_extensions import override
 
+from ..events._branch_path import _BranchPath
 from ..events.event import Event
 from ..utils.context_utils import Aclosing
 from .base_agent import BaseAgent
@@ -44,10 +45,8 @@ def _create_branch_ctx_for_sub_agent(
   """Create isolated branch for every sub-agent."""
   invocation_context = invocation_context.model_copy()
   branch_suffix = f'{agent.name}.{sub_agent.name}'
-  invocation_context.branch = (
-      f'{invocation_context.branch}.{branch_suffix}'
-      if invocation_context.branch
-      else branch_suffix
+  invocation_context.branch = _BranchPath.create_sub_branch(
+      invocation_context.branch, name=branch_suffix
   )
   return invocation_context
 

src/google/adk/events/_branch_path.py

@@ -37,11 +37,11 @@ def from_string(cls, path_str: str | None) -> _BranchPath:
     """Parses a _BranchPath from a dot-separated string representation."""
     if not path_str:
       return cls([])
-    return cls(path_str.split('.'))
+    return cls(path_str.split("."))
 
   def __str__(self) -> str:
     """Returns the dot-separated string representation of the path."""
-    return '.'.join(self._segments)
+    return ".".join(self._segments)
 
   def __eq__(self, other: object) -> bool:
     """Returns True if segments are equal."""
@@ -64,7 +64,7 @@ def run_ids(self) -> set[str]:
     """
     ids = set()
     for segment in self._segments:
-      parts = segment.rsplit('@', 1)
+      parts = segment.rsplit("@", 1)
       if len(parts) > 1 and parts[1]:
         ids.add(parts[1])
     return ids
@@ -99,3 +99,53 @@ def common_prefix(paths: list[_BranchPath]) -> _BranchPath:
       else:
         break
     return _BranchPath(common_segments)
+
+  def append(
+      self,
+      segment_or_path: str | _BranchPath,
+      run_id: str | None = None,
+  ) -> _BranchPath:
+    """Returns a new _BranchPath with segment(s) appended.
+
+    Args:
+      segment_or_path: A segment name (str), dot-separated path (str), or another
+        _BranchPath instance to append.
+      run_id: Optional run ID (or function_call_id) to format segment as
+        'name@run_id'.
+    """
+    if isinstance(segment_or_path, _BranchPath):
+      if run_id is not None:
+        raise ValueError(
+            "run_id cannot be provided when segment_or_path is a _BranchPath"
+            " instance."
+        )
+      return _BranchPath(self._segments + segment_or_path.segments)
+
+    if run_id is not None:
+      if "." in segment_or_path:
+        raise ValueError(
+            "run_id cannot be provided when segment_or_path is a dot-separated"
+            " path."
+        )
+      segment = f"{segment_or_path}@{run_id}"
+      return _BranchPath(self._segments + [segment])
+
+    new_segments = [s for s in segment_or_path.split(".") if s]
+    return _BranchPath(self._segments + new_segments)
+
+  @classmethod
+  def create_sub_branch(
+      cls,
+      base_branch: str | None,
+      *,
+      name: str,
+      run_id: str | None = None,
+  ) -> str:
+    """Creates a new dot-separated branch path string by appending a segment.
+
+    Example:
+      _BranchPath.create_sub_branch('parent', name='child', run_id='1') ->
+      'parent.child@1'
+      _BranchPath.create_sub_branch(None, name='agent') -> 'agent'
+    """
+    return str(cls.from_string(base_branch).append(name, run_id=run_id))

src/google/adk/tools/agent_tool.py

@@ -27,6 +27,7 @@
 
 from . import _automatic_function_calling_util
 from ..agents.common_configs import AgentRefConfig
+from ..events._branch_path import _BranchPath
 from ..features import FeatureName
 from ..features import is_feature_enabled
 from ..memory.in_memory_memory_service import InMemoryMemoryService
@@ -364,8 +365,9 @@ async def run_async(
     # Align subagent branch scoping with node execution (Node as Tool) using function_call_id.
     fc_id = tool_context.function_call_id
     base_branch = tool_context.get_invocation_context().branch
-    segment = f'{self.agent.name}@{fc_id}'
-    tool_branch = f'{base_branch}.{segment}' if base_branch else segment
+    tool_branch = _BranchPath.create_sub_branch(
+        base_branch, name=self.agent.name, run_id=fc_id
+    )
 
     try:
       return await tool_context.run_node(

src/google/adk/workflow/_node_runner.py

@@ -28,6 +28,7 @@
 from typing import Any
 from typing import TYPE_CHECKING
 
+from ..events._branch_path import _BranchPath
 from ..telemetry import node_tracing
 
 if TYPE_CHECKING:
@@ -206,8 +207,9 @@ def _create_child_context(
     )
 
     if self._use_sub_branch:
-      segment = f"{self._node.name}@{self._run_id}"
-      branch = f"{base_branch}.{segment}" if base_branch else segment
+      branch = _BranchPath.create_sub_branch(
+          base_branch, name=self._node.name, run_id=self._run_id
+      )
       ic = ic.model_copy(update={"branch": branch})
     elif self._override_branch is not None:
       ic = ic.model_copy(update={"branch": self._override_branch})

tests/unittests/events/test_branch_path.py

@@ -153,6 +153,64 @@ def test_constructor_copies_segments_list():
   assert path.segments == ["parent", "child"]
 
 
+def test_append_single_segment_returns_new_path():
+  """append adds a single segment to an existing path."""
+  path = _BranchPath.from_string("parent")
+  new_path = path.append("child")
+
+  assert new_path == _BranchPath.from_string("parent.child")
+  assert path == _BranchPath.from_string("parent")  # Immutability check
+
+
+def test_append_with_run_id_formats_segment():
+  """append formats the segment as 'name@run_id' when run_id is provided."""
+  path = _BranchPath.from_string("parent")
+  new_path = path.append("child", run_id="call_123")
+
+  assert new_path == _BranchPath.from_string("parent.child@call_123")
+
+
+def test_append_another_branch_path():
+  """append combines segments from another _BranchPath instance."""
+  path1 = _BranchPath.from_string("parent")
+  path2 = _BranchPath.from_string("child.grandchild")
+  new_path = path1.append(path2)
+
+  assert new_path == _BranchPath.from_string("parent.child.grandchild")
+
+
+def test_create_sub_branch_formats_string_correctly():
+  """create_sub_branch constructs sub-branch strings safely."""
+  # With base branch and run ID
+  res1 = _BranchPath.create_sub_branch(
+      "parent.sub", name="child", run_id="run_1"
+  )
+  assert res1 == "parent.sub.child@run_1"
+
+  # Without base branch (None or empty)
+  res2 = _BranchPath.create_sub_branch(None, name="agent", run_id="fc_456")
+  assert res2 == "agent@fc_456"
+
+  # Dot-separated sub-path without run ID
+  res3 = _BranchPath.create_sub_branch("parent", name="agent.sub_agent")
+  assert res3 == "parent.agent.sub_agent"
+
+
+def test_append_with_run_id_and_branch_path_raises_value_error():
+  """append raises ValueError when run_id is provided with a _BranchPath."""
+  path1 = _BranchPath.from_string("parent")
+  path2 = _BranchPath.from_string("child")
+  with pytest.raises(ValueError, match="run_id cannot be provided"):
+    path1.append(path2, run_id="123")
+
+
+def test_append_with_run_id_and_dot_separated_path_raises_value_error():
+  """append raises ValueError when run_id is provided with a dot-separated path."""
+  path = _BranchPath.from_string("parent")
+  with pytest.raises(ValueError, match="run_id cannot be provided"):
+    path.append("child.sub", run_id="123")
+
+
 def test_run_ids_filters_out_empty_run_ids():
   """run_ids filters out segments with empty run IDs (e.g. ending with '@')."""
   path = _BranchPath.from_string("parent@.child@2.node@")