Skip to content

Commit 27edb47

Browse files
AtulJoshi1206copybara-github
authored andcommitted
fix: keep a node's branch when it relays sub-agent events
Merge #7940 A node that relays its sub-agents' events (such as a ParallelAgent or a SequentialAgent containing one, run as a Workflow node or transferred to from a root LlmAgent) no longer moves onto the branch of whichever sub-agent spoke last. Relayed events keep their own sub-branch; the node's later events and its successors stay on the node's branch, so an agent after a parallel step sees every sub-agent's output. Resuming a paused workflow recovers the same branch, so a node after the parallel step that paused for a tool confirmation resumes where it paused and runs the confirmed tool. PiperOrigin-RevId: 1013319278
1 parent 0894c03 commit 27edb47

6 files changed

Lines changed: 341 additions & 10 deletions

File tree

‎src/google/adk/workflow/_node_runner.py‎

Lines changed: 21 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -465,11 +465,7 @@ def _track_event_in_context(self, event: Event, ctx: Context) -> None:
465465
ctx.output = event.content
466466
if event.long_running_tool_ids is not None:
467467
ctx._interrupt_ids.update(event.long_running_tool_ids)
468-
# Only propagate decisions from native events (authored by this node or unspecified).
469-
# This prevents structured parent nodes (e.g. SequentialAgent) from intercepting
470-
# and bubbling up actions already handled internally by their nested sub-agents.
471-
is_native_node_event = not event.author or event.author == self._node.name
472-
if event.actions and is_native_node_event:
468+
if event.actions and self._is_native_event(event):
473469
if event.actions.route is not None:
474470
ctx.route = event.actions.route
475471
ctx._route_emitted = True
@@ -562,18 +558,34 @@ def _flush_deltas(self, event: Event, ctx: Context) -> None:
562558
event.actions.artifact_delta.update(artifact_delta)
563559
artifact_delta.clear()
564560

561+
def _is_native_event(self, event: Event) -> bool:
562+
"""Whether this node produced the event itself rather than relaying it.
563+
564+
A structured agent run as a node (e.g. a SequentialAgent or ParallelAgent)
565+
re-yields its sub-agents' events, still authored by those sub-agents. The
566+
sub-agents already acted on what such an event carries, so its route,
567+
transfer and branch apply to them, not to this node. Read the author before
568+
``_enrich_event`` overwrites it.
569+
"""
570+
return not event.author or event.author == self._node.name
571+
565572
def _enrich_event(self, event: Event, ctx: Context) -> None:
566573
"""Set author, node_info, invocation_id on the event."""
574+
is_native_event = self._is_native_event(event)
567575
event.author = ctx.event_author or self._node.name
568576
event.invocation_id = ctx._invocation_context.invocation_id
569577
event.node_info.path = ctx.node_path
570578
if event.branch is None:
571579
event.branch = ctx._invocation_context.branch
572-
elif event.branch == "":
573-
event.branch = None
574-
ctx._invocation_context.branch = None
575580
else:
576-
ctx._invocation_context.branch = event.branch
581+
if event.branch == "":
582+
event.branch = None
583+
# An event the node emits on another branch moves the node there. A
584+
# relayed event stays on its sub-agent's branch: following it would put
585+
# the node's later events and its successors on whichever sub-branch
586+
# spoke last, hiding the other branches from them.
587+
if is_native_event:
588+
ctx._invocation_context.branch = event.branch
577589
if event.output is not None:
578590
event.node_info.output_for = [ctx.node_path] + ctx._output_for_ancestors
579591
# An explicitly claimed output is the claimer's output: publish it on

‎src/google/adk/workflow/utils/_rehydration_utils.py‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -414,8 +414,15 @@ def get_owner_key(event_path_builder: _NodePathBuilder) -> str | None:
414414
event.branch, owner_path_builder.node_name, tool_call_ids
415415
)
416416
)
417+
# Like NodeRunner._is_native_event: only the node's own events move its
418+
# branch. A stored native event carries the node's name or one it inherits
419+
# from an ancestor; a relayed one keeps its sub-agent's name.
420+
is_native = not event.author or event.author in (
421+
event_path_builder.static_path.split('/')
422+
)
417423
if (
418424
is_direct
425+
and is_native
419426
and event.branch
420427
and not has_claimed_branch
421428
and not is_tool_branch
@@ -445,7 +452,12 @@ def get_owner_key(event_path_builder: _NodePathBuilder) -> str | None:
445452
if is_direct or is_delegated:
446453
if event.output is not None:
447454
child.output = event.output
448-
if is_direct and not has_claimed_branch and not is_tool_branch:
455+
if (
456+
is_direct
457+
and is_native
458+
and not has_claimed_branch
459+
and not is_tool_branch
460+
):
449461
child.branch = event.branch
450462
elif use_message_as_output:
451463
child.output = event.content

‎tests/unittests/workflow/test_agent_transfer.py‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
from google.adk.agents.llm_agent import Agent
2222
from google.adk.agents.loop_agent import LoopAgent
2323
from google.adk.agents.loop_agent import LoopAgentState
24+
from google.adk.agents.parallel_agent import ParallelAgent
2425
from google.adk.agents.sequential_agent import SequentialAgent
2526
from google.adk.agents.sequential_agent import SequentialAgentState
2627
from google.adk.apps.app import App
@@ -323,6 +324,67 @@ def test_auto_to_sequential(is_resumable: bool):
323324
]
324325

325326

327+
@pytest.mark.parametrize('is_resumable', [True, False])
328+
def test_auto_to_sequential_with_parallel_keeps_merger_branch(
329+
is_resumable: bool,
330+
):
331+
"""A merger after a parallel step stays on its sequence's branch.
332+
333+
After the transfer the pipeline runs as one node that relays the events of
334+
every agent inside it. The merger's reply belongs on the pipeline's branch,
335+
as when the same pipeline is the root agent, not on whichever parallel
336+
sub-branch spoke last; otherwise only that one worker sees the merger's
337+
reply on later turns.
338+
"""
339+
# root (auto) - pipeline (seq) - fan_out (parallel) - worker_a (single)
340+
# \ worker_b (single)
341+
# \ merger (single)
342+
worker_a = Agent(
343+
name='worker_a',
344+
model=testing_utils.MockModel.create(responses=['from a']),
345+
)
346+
worker_b = Agent(
347+
name='worker_b',
348+
model=testing_utils.MockModel.create(responses=['from b']),
349+
)
350+
merger = Agent(
351+
name='merger',
352+
model=testing_utils.MockModel.create(responses=['merged']),
353+
)
354+
pipeline = SequentialAgent(
355+
name='pipeline',
356+
sub_agents=[
357+
ParallelAgent(name='fan_out', sub_agents=[worker_a, worker_b]),
358+
merger,
359+
],
360+
)
361+
root_agent = Agent(
362+
name='root_agent',
363+
model=testing_utils.MockModel.create(
364+
responses=[transfer_call_part('pipeline')]
365+
),
366+
sub_agents=[pipeline],
367+
)
368+
app = App(
369+
name='test_app',
370+
root_agent=root_agent,
371+
resumability_config=ResumabilityConfig(is_resumable=is_resumable),
372+
)
373+
runner = testing_utils.InMemoryRunner(app=app)
374+
375+
events = runner.run('test1')
376+
377+
assert {
378+
event.author: event.branch
379+
for event in events
380+
if event.content and event.author != 'root_agent'
381+
} == {
382+
'worker_a': 'fan_out.worker_a',
383+
'worker_b': 'fan_out.worker_b',
384+
'merger': None,
385+
}
386+
387+
326388
@pytest.mark.parametrize('is_resumable', [True, False])
327389
def test_auto_to_sequential_to_auto(is_resumable: bool):
328390
response = [

‎tests/unittests/workflow/test_node_runner_ctx.py‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -968,6 +968,56 @@ async def _run_impl(self, *, ctx, node_input):
968968
assert parent_ctx._invocation_context.branch == 'parent_branch'
969969

970970

971+
@pytest.mark.asyncio
972+
async def test_own_event_branch_moves_node_branch():
973+
"""An event the node authors on another branch moves the node there."""
974+
975+
class _Node(BaseNode):
976+
977+
async def _run_impl(self, *, ctx, node_input):
978+
yield Event(branch='other_branch')
979+
yield Event(output='result')
980+
981+
parent_ctx, events = _make_ctx()
982+
parent_ctx._invocation_context.branch = 'parent_branch'
983+
child_ctx = await NodeRunner(
984+
node=_Node(name='n'), parent_ctx=parent_ctx
985+
).run()
986+
987+
assert [e.branch for e in events] == ['other_branch', 'other_branch']
988+
assert child_ctx._invocation_context.branch == 'other_branch'
989+
990+
991+
@pytest.mark.asyncio
992+
async def test_relayed_event_branch_does_not_move_node_branch():
993+
"""A sub-agent's event relayed by the node does not move the node's branch.
994+
995+
A structured agent run as a node, such as a ParallelAgent, re-yields events
996+
its sub-agents emitted on their own sub-branches. Each relayed event keeps
997+
its sub-branch, but the node stays on its own branch, so its later events
998+
and the branch its successors inherit do not depend on which sub-agent
999+
spoke last.
1000+
"""
1001+
1002+
class _Node(BaseNode):
1003+
1004+
async def _run_impl(self, *, ctx, node_input):
1005+
yield Event(author='sub_agent', branch='parent_branch.n.sub_agent')
1006+
yield Event(output='result')
1007+
1008+
parent_ctx, events = _make_ctx()
1009+
parent_ctx._invocation_context.branch = 'parent_branch'
1010+
child_ctx = await NodeRunner(
1011+
node=_Node(name='n'), parent_ctx=parent_ctx
1012+
).run()
1013+
1014+
assert [e.branch for e in events] == [
1015+
'parent_branch.n.sub_agent',
1016+
'parent_branch',
1017+
]
1018+
assert child_ctx._invocation_context.branch == 'parent_branch'
1019+
1020+
9711021
@pytest.mark.asyncio
9721022
async def test_override_isolation_scope_used_in_node_runner():
9731023
"""NodeRunner sets isolation_scope on child context and enriches emitted events."""

‎tests/unittests/workflow/test_workflow_agent_as_node.py‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,9 @@
1919
from google.adk.agents.base_agent import BaseAgent
2020
from google.adk.agents.context import Context
2121
from google.adk.agents.invocation_context import InvocationContext as BaseInvocationContext
22+
from google.adk.agents.llm_agent import LlmAgent
23+
from google.adk.agents.parallel_agent import ParallelAgent
24+
from google.adk.agents.sequential_agent import SequentialAgent
2225
from google.adk.events.event import Event
2326
from google.adk.sessions.in_memory_session_service import InMemorySessionService
2427
from google.adk.sessions.session import Session
@@ -27,6 +30,8 @@
2730
from google.genai import types
2831
import pytest
2932

33+
from tests.unittests import testing_utils
34+
3035
from .workflow_testing_utils import create_runner
3136
from .workflow_testing_utils import InputCapturingNode
3237
from .workflow_testing_utils import simplify_events_with_node
@@ -165,3 +170,116 @@ async def test_run_async_with_agent_node_piping_data(
165170
# AgentNode does not record content as node output, so the next node
166171
# receives None as input.
167172
assert node_b.received_inputs == [None]
173+
174+
175+
async def _run_workflow(wf: Workflow) -> list[Event]:
176+
ss = InMemorySessionService()
177+
runner = create_runner(app_name='test', node=wf, session_service=ss)
178+
session = await ss.create_session(app_name='test', user_id='u')
179+
msg = types.Content(parts=[types.Part(text='start')], role='user')
180+
return [
181+
event
182+
async for event in runner.run_async(
183+
user_id='u', session_id=session.id, new_message=msg
184+
)
185+
]
186+
187+
188+
@pytest.mark.asyncio
189+
async def test_node_after_parallel_agent_node_sees_every_branch():
190+
"""A node after a ParallelAgent node runs on the workflow's own branch.
191+
192+
The ParallelAgent relays each sub-agent's events on that sub-agent's
193+
branch. Its successor must not inherit whichever of those branches spoke
194+
last: an agent there reads history filtered by branch, so it would see only
195+
that one sub-agent's answer, and its own replies would land on that
196+
sub-agent's branch.
197+
"""
198+
researcher_a = LlmAgent(
199+
name='researcher_a',
200+
model=testing_utils.MockModel.create(responses=['finding from a']),
201+
)
202+
researcher_b = LlmAgent(
203+
name='researcher_b',
204+
model=testing_utils.MockModel.create(responses=['finding from b']),
205+
)
206+
research = ParallelAgent(
207+
name='research', sub_agents=[researcher_a, researcher_b]
208+
)
209+
summarizer_model = testing_utils.MockModel.create(responses=['summary'])
210+
report = SequentialAgent(
211+
name='report',
212+
sub_agents=[LlmAgent(name='summarizer', model=summarizer_model)],
213+
)
214+
215+
events = await _run_workflow(
216+
Workflow(name='wf', edges=[(START, research, report)])
217+
)
218+
219+
summarizer_request = ' '.join(
220+
part.text
221+
for content in summarizer_model.requests[0].contents
222+
for part in content.parts or []
223+
if part.text
224+
)
225+
assert 'finding from a' in summarizer_request
226+
assert 'finding from b' in summarizer_request
227+
assert {event.author: event.branch for event in events if event.content} == {
228+
'researcher_a': 'research.researcher_a',
229+
'researcher_b': 'research.researcher_b',
230+
'summarizer': None,
231+
}
232+
233+
234+
class _Dispatcher(BaseAgent):
235+
"""Runs its sub-agent on a sub-branch and hands back the final answer."""
236+
237+
async def _run_async_impl(
238+
self, ctx: BaseInvocationContext
239+
) -> AsyncGenerator[Event, None]:
240+
target = self.sub_agents[0]
241+
parent_branch = ctx.branch
242+
ctx.branch = f'{self.name}.{target.name}'
243+
async for event in target.run_async(ctx):
244+
if event.is_final_response():
245+
event.branch = parent_branch
246+
yield event
247+
248+
249+
@pytest.mark.asyncio
250+
async def test_answer_relayed_onto_parent_branch_stays_there():
251+
"""A relayed answer re-stamped onto the parent branch is not pulled back.
252+
253+
The dispatcher keeps its sub-agent's tool traffic on a sub-branch and hands
254+
the final answer back on its own branch, so later agents see the answer
255+
but not the tool calls. The tool events must not move the dispatcher node
256+
onto the sub-branch, or the branch-less answer gets stamped with it.
257+
"""
258+
259+
def lookup_arrival() -> dict[str, str]:
260+
"""Looks up the arrival time."""
261+
return {'arrival': 'Tuesday 14:05'}
262+
263+
bookings = LlmAgent(
264+
name='bookings',
265+
model=testing_utils.MockModel.create(
266+
responses=[
267+
types.Part.from_function_call(name='lookup_arrival', args={}),
268+
'You land on Tuesday at 14:05.',
269+
]
270+
),
271+
tools=[lookup_arrival],
272+
)
273+
dispatcher = _Dispatcher(name='dispatcher', sub_agents=[bookings])
274+
275+
events = await _run_workflow(Workflow(name='wf', edges=[(START, dispatcher)]))
276+
277+
bookings_events = [e for e in events if e.author == 'bookings' and e.content]
278+
assert [e.branch for e in bookings_events] == [
279+
'dispatcher.bookings',
280+
'dispatcher.bookings',
281+
None,
282+
]
283+
assert bookings_events[-1].content.parts[0].text == (
284+
'You land on Tuesday at 14:05.'
285+
)

0 commit comments

Comments
 (0)