From a18fd640db5d8fc8f970f183382dd16e98dc1b1a Mon Sep 17 00:00:00 2001
From: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Date: Mon, 8 Jun 2026 07:37:44 +0800
Subject: [PATCH] fix: scope single-turn node inputs to workflow branch

---
 src/google/adk/workflow/_llm_agent_wrapper.py |  7 ++++--
 .../workflow/test_llm_agent_as_node.py        | 22 +++++++++++++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/src/google/adk/workflow/_llm_agent_wrapper.py b/src/google/adk/workflow/_llm_agent_wrapper.py
index 36a487b0e0..0d2f738595 100644
--- a/src/google/adk/workflow/_llm_agent_wrapper.py
+++ b/src/google/adk/workflow/_llm_agent_wrapper.py
@@ -233,8 +233,8 @@ def prepare_llm_agent_input(agent: Any, ctx: Context, node_input: Any) -> None:
   overrides ``ic.user_content`` so the content-builder can fall back
   to that as the first user turn.
 
-  No branch is set — task and single_turn agents scope via
-  ``isolation_scope`` rather than branch.
+  For workflow nodes running in a sub-branch, stamp the input event with that
+  branch. A private node input should not look like the shared root user turn.
   """
   if node_input is None or agent.mode != 'single_turn':
     return
@@ -245,6 +245,9 @@ def prepare_llm_agent_input(agent: Any, ctx: Context, node_input: Any) -> None:
   iso = getattr(ctx, 'isolation_scope', None)
   if iso:
     user_event.isolation_scope = iso
+  branch = getattr(ctx._invocation_context, 'branch', None)
+  if branch:
+    user_event.branch = branch
   ctx.session.events.append(user_event)
 
 
diff --git a/tests/unittests/workflow/test_llm_agent_as_node.py b/tests/unittests/workflow/test_llm_agent_as_node.py
index 8e4fcf116d..faeb5860ba 100644
--- a/tests/unittests/workflow/test_llm_agent_as_node.py
+++ b/tests/unittests/workflow/test_llm_agent_as_node.py
@@ -212,6 +212,28 @@ def test_rerun_on_resume_defaults_true(self):
     assert wrapper.rerun_on_resume is True
 
 
+@pytest.mark.asyncio
+async def test_single_turn_input_event_inherits_branch_and_scope(
+    request: pytest.FixtureRequest,
+):
+  """Private single-turn node input is scoped to the node branch."""
+  from google.adk.workflow._llm_agent_wrapper import prepare_llm_agent_input
+
+  agent = _make_agent(mode='single_turn')
+  ic = await create_parent_invocation_context(request.function.__name__, agent)
+  ic.branch = 'parent.worker@1'
+  ctx = Context(invocation_context=ic)
+  ctx.isolation_scope = 'scope-1'
+
+  prepare_llm_agent_input(agent, ctx, 'hello')
+
+  event = ic.session.events[-1]
+  assert event.author == 'user'
+  assert event.content and event.content.role == 'user'
+  assert event.branch == 'parent.worker@1'
+  assert event.isolation_scope == 'scope-1'
+
+
 # --- build_node auto-wrapping ---