Skip to content

Commit 8d389e2

Browse files
he-yufengGWeale
authored andcommitted
fix: scope single-turn node inputs to workflow branch
Merge #6008 ## Summary Fixes #5989. `prepare_llm_agent_input()` appends the single-turn node input directly to `session.events`, so it bypasses `NodeRunner._enrich_event()`. In a dynamic workflow sub-branch, that left the private node input as a branchless user event. Branchless user events are visible to sibling branches, which can make another worker pick the wrong "current turn" boundary and slice out its own function call. This stamps the single-turn input event with the current invocation branch when one exists, while preserving the existing isolation-scope stamping. ## To verify - `PYTHONPATH=src python -m pytest tests\unittests\workflow\test_llm_agent_as_node.py -q -k "single_turn_input_event_inherits_branch_and_scope"` - `PYTHONPATH=src python -m pytest tests\unittests\workflow\test_llm_agent_as_node.py -q` - `python -m py_compile src\google\adk\workflow\_llm_agent_wrapper.py tests\unittests\workflow\test_llm_agent_as_node.py` - `python -m pyink --check src\google\adk\workflow\_llm_agent_wrapper.py tests\unittests\workflow\test_llm_agent_as_node.py` - `python -m isort --check-only src\google\adk\workflow\_llm_agent_wrapper.py tests\unittests\workflow\test_llm_agent_as_node.py` - `git diff --check` Co-authored-by: George Weale <[email protected]> COPYBARA_INTEGRATE_REVIEW=#6008 from he-yufeng:fix/node-input-branch-scope f178fb9 PiperOrigin-RevId: 938712082
1 parent febb250 commit 8d389e2

2 files changed

Lines changed: 27 additions & 2 deletions

File tree

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -235,8 +235,8 @@ def prepare_llm_agent_input(agent: Any, ctx: Context, node_input: Any) -> None:
235235
overrides ``ic.user_content`` so the content-builder can fall back
236236
to that as the first user turn.
237237
238-
No branch is set — task and single_turn agents scope via
239-
``isolation_scope`` rather than branch.
238+
For workflow nodes running in a sub-branch, stamp the input event with that
239+
branch. A private node input should not look like the shared root user turn.
240240
"""
241241
if node_input is None or agent.mode != 'single_turn':
242242
return
@@ -247,6 +247,9 @@ def prepare_llm_agent_input(agent: Any, ctx: Context, node_input: Any) -> None:
247247
iso = getattr(ctx, 'isolation_scope', None)
248248
if iso:
249249
user_event.isolation_scope = iso
250+
branch = ctx._invocation_context.branch
251+
if branch:
252+
user_event.branch = branch
250253
ctx.session.events.append(user_event)
251254

252255

‎tests/unittests/workflow/test_llm_agent_as_node.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,28 @@ def test_rerun_on_resume_defaults_true(self):
212212
assert wrapper.rerun_on_resume is True
213213

214214

215+
@pytest.mark.asyncio
216+
async def test_single_turn_input_event_inherits_branch_and_scope(
217+
request: pytest.FixtureRequest,
218+
):
219+
"""Private single-turn node input is scoped to the node branch."""
220+
from google.adk.workflow._llm_agent_wrapper import prepare_llm_agent_input
221+
222+
agent = _make_agent(mode='single_turn')
223+
ic = await create_parent_invocation_context(request.function.__name__, agent)
224+
ic.branch = 'parent.worker@1'
225+
ctx = Context(invocation_context=ic)
226+
ctx.isolation_scope = 'scope-1'
227+
228+
prepare_llm_agent_input(agent, ctx, 'hello')
229+
230+
event = ic.session.events[-1]
231+
assert event.author == 'user'
232+
assert event.content and event.content.role == 'user'
233+
assert event.branch == 'parent.worker@1'
234+
assert event.isolation_scope == 'scope-1'
235+
236+
215237
# --- build_node auto-wrapping ---
216238

217239

0 commit comments

Comments
 (0)