Skip to content

Commit 5c46937

Browse files
DeanChensjcopybara-github
authored andcommitted
fix(workflow): Prevent Shared InvocationContext branch mutation
Previously, the fallback path failed to make a copy, leading to child events mutating the parent and siblings' branches. Co-authored-by: Shangjie Chen <deanchen@google.com> PiperOrigin-RevId: 930738022
1 parent c67c7af commit 5c46937

2 files changed

Lines changed: 24 additions & 0 deletions

File tree

src/google/adk/workflow/_node_runner.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,8 @@ def _create_child_context(
207207
ic = ic.model_copy(update={"branch": branch})
208208
elif self._override_branch is not None:
209209
ic = ic.model_copy(update={"branch": self._override_branch})
210+
else:
211+
ic = ic.model_copy()
210212

211213
ctx = Context(
212214
ic,

tests/unittests/workflow/test_node_runner_ctx.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -588,6 +588,28 @@ async def _run_impl(self, *, ctx, node_input):
588588
assert events[0].branch == 'parent_branch'
589589

590590

591+
@pytest.mark.asyncio
592+
async def test_child_event_branch_does_not_mutate_parent_ic():
593+
"""A child node altering its branch does not mutate the parent's shared InvocationContext branch."""
594+
595+
class _Node(BaseNode):
596+
597+
async def _run_impl(self, *, ctx, node_input):
598+
yield Event(output='result', branch='new_child_branch')
599+
600+
parent_ctx, events = _make_ctx()
601+
parent_ctx._invocation_context.branch = 'parent_branch'
602+
await NodeRunner(
603+
node=_Node(name='n'),
604+
parent_ctx=parent_ctx,
605+
use_sub_branch=False,
606+
).run()
607+
608+
assert events[0].branch == 'new_child_branch'
609+
# The parent's branch must remain unchanged.
610+
assert parent_ctx._invocation_context.branch == 'parent_branch'
611+
612+
591613
@pytest.mark.asyncio
592614
async def test_override_isolation_scope_used_in_node_runner():
593615
"""NodeRunner sets isolation_scope on child context and enriches emitted events."""

0 commit comments

Comments
 (0)