fix(workflow): store a single-turn node's synthetic input in the session

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 974056680

Author: GWeale

src/google/adk/workflow/_llm_agent_wrapper.py

@@ -302,15 +302,18 @@ def prepare_llm_agent_context(agent: LlmAgent, ctx: Context) -> Context:
   return agent_ctx
 
 
-def prepare_llm_agent_input(
+async def prepare_llm_agent_input(
     agent: LlmAgent, ctx: Context, node_input: object
 ) -> None:
   """Prepares the input for running LlmAgent as a node.
 
-  For ``single_turn`` mode, append a user-role event with the input
-  directly to session.events (legacy behavior). When resuming with
-  ``resume_inputs``, skip appending to avoid injecting duplicate synthetic
-  user events that shadow user function responses.
+  For ``single_turn`` mode, write a user-role event carrying the input
+  through the session service. Appending to ``session.events`` by hand
+  instead reaches only the in-memory list, so the input never becomes part
+  of the stored session. Nothing is written when the input is the
+  invocation's own user message, because the runner has already recorded
+  that turn, nor when resuming with ``resume_inputs``, because a synthetic
+  user event would shadow the user's function responses.
 
   For ``task`` mode, the input is the parent's task-delegation FC
   args.  Those are NOT appended here — the content-builder
@@ -336,17 +339,32 @@ def prepare_llm_agent_input(
       or bool(ctx.resume_inputs)
   ):
     return
+  from ..runners import _apply_run_config_custom_metadata
+
   agent_input = to_user_content(node_input)
-  user_event = Event(author='user', message=agent_input)
+  ic = ctx._invocation_context
+  # A node wired directly to START is handed the invocation's own user
+  # message, which the runner already recorded. Matching on rendered content
+  # is approximate: a mid-workflow node whose input happens to equal the
+  # user's turn is skipped too.
+  if (
+      ic.user_content is not None
+      and to_user_content(ic.user_content) == agent_input
+  ):
+    return
+  user_event = Event(
+      author='user', message=agent_input, invocation_id=ic.invocation_id
+  )
   if user_event.content is not None:
     user_event.content.role = 'user'
+  _apply_run_config_custom_metadata(user_event, ic.run_config)
   iso = getattr(ctx, 'isolation_scope', None)
   if iso:
     user_event.isolation_scope = iso
-  branch = ctx._invocation_context.branch
+  branch = ic.branch
   if branch:
     user_event.branch = branch
-  ctx.session.events.append(user_event)
+  await ic.session_service.append_event(session=ic.session, event=user_event)
 
 
 def process_llm_agent_output(
@@ -404,7 +422,7 @@ async def run_llm_agent_as_node(
     agent.include_contents = 'none'
 
   agent_ctx = prepare_llm_agent_context(agent, ctx)
-  prepare_llm_agent_input(agent, agent_ctx, node_input)
+  await prepare_llm_agent_input(agent, agent_ctx, node_input)
 
   ic = agent_ctx.get_invocation_context()
   update: dict[str, object] = {'agent': agent}

tests/unittests/workflow/test_llm_agent_as_node.py

@@ -231,15 +231,34 @@ async def test_single_turn_input_event_inherits_branch_and_scope(
   ctx = Context(invocation_context=ic)
   ctx.isolation_scope = 'scope-1'
 
-  prepare_llm_agent_input(agent, ctx, 'hello')
+  await prepare_llm_agent_input(agent, ctx, 'hello')
 
   event = ic.session.events[-1]
   assert event.author == 'user'
+  assert event.invocation_id == ic.invocation_id
   assert event.content and event.content.role == 'user'
   assert event.branch == 'parent.worker@1'
   assert event.isolation_scope == 'scope-1'
 
 
+@pytest.mark.asyncio
+async def test_single_turn_input_event_carries_run_custom_metadata(
+    request: pytest.FixtureRequest,
+):
+  """Run-level custom metadata tags the synthetic input like any other event."""
+  from google.adk.agents.run_config import RunConfig
+  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.run_config = RunConfig(custom_metadata={'tenant': 'acme'})
+  ctx = Context(invocation_context=ic)
+
+  await prepare_llm_agent_input(agent, ctx, 'hello')
+
+  assert ic.session.events[-1].custom_metadata == {'tenant': 'acme'}
+
+
 @pytest.mark.asyncio
 async def test_single_turn_input_skipped_when_resuming(
     request: pytest.FixtureRequest,
@@ -265,14 +284,92 @@ async def test_single_turn_input_skipped_when_resuming(
   )
 
   initial_len = len(ic.session.events)
-  prepare_llm_agent_input(agent, ctx, 'hello')
+  await prepare_llm_agent_input(agent, ctx, 'hello')
 
   # Verify no duplicate user input was appended on resume
   assert len(ic.session.events) == initial_len
   # Verify the resumed node still sees the initial user input from turn 1
   assert ic.session.events[-1].content.parts[0].text == 'turn 1 initial input'
 
 
+@pytest.mark.asyncio
+async def test_single_turn_input_skipped_when_it_is_the_user_message(
+    request: pytest.FixtureRequest,
+):
+  """The runner already recorded the invocation's own user message."""
+  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.user_content = types.Content(
+      role='user', parts=[types.Part(text='who am i')]
+  )
+  ctx = Context(invocation_context=ic)
+
+  await prepare_llm_agent_input(agent, ctx, ic.user_content)
+
+  assert not ic.session.events
+
+
+@pytest.mark.asyncio
+async def test_single_turn_input_event_reaches_the_stored_session(
+    request: pytest.FixtureRequest,
+):
+  """A single-turn node's input becomes part of the stored session.
+
+  Appending to the in-memory events list reaches only the Session object the
+  node happens to hold, so the input is missing from the session anyone loads
+  afterwards and the model turn it prompted has no prompt.
+  """
+  from . import testing_utils
+
+  async def brief(node_input: Any) -> str:
+    return 'write the brief'
+
+  wrapper = build_node(_make_agent(mode='single_turn'))
+  wf = Workflow(name='wf', edges=[(START, brief), (brief, wrapper)])
+  runner = _new_workflow_runner(wf, request.function.__name__)
+
+  agent_clone = next(n for n in wf.graph.nodes if n.name == wrapper.name)
+  with _mock_leaf_run(agent_clone, content_text='Done.'):
+    await runner.run_async(testing_utils.get_user_content('start'))
+
+  stored = runner.session
+  assert any(
+      event.author == 'user'
+      and event.content
+      and event.content.parts
+      and event.content.parts[0].text == 'write the brief'
+      for event in stored.events
+  )
+
+
+@pytest.mark.asyncio
+async def test_first_node_does_not_store_the_user_message_twice(
+    request: pytest.FixtureRequest,
+):
+  """A node wired to START is fed the turn the runner already stored."""
+  from . import testing_utils
+
+  wrapper = build_node(_make_agent(mode='single_turn'))
+  wf = Workflow(name='wf', edges=[(START, wrapper)])
+  runner = _new_workflow_runner(wf, request.function.__name__)
+
+  agent_clone = next(n for n in wf.graph.nodes if n.name == wrapper.name)
+  with _mock_leaf_run(agent_clone, content_text='Done.'):
+    await runner.run_async(testing_utils.get_user_content('who am i'))
+
+  user_turns = [
+      event
+      for event in runner.session.events
+      if event.author == 'user'
+      and event.content
+      and event.content.parts
+      and event.content.parts[0].text == 'who am i'
+  ]
+  assert len(user_turns) == 1
+
+
 # --- build_node auto-wrapping ---
 
 
@@ -340,17 +437,16 @@ async def mock_run_async(*args, **kwargs):
           content=types.Content(parts=[types.Part(text='ok')]),
       )
 
+    async def skip_input(agent, ctx, node_input):
+      return None
+
     object.__setattr__(wrapper, 'run_async', mock_run_async)
     monkeypatch.setattr(
         agent_wrapper,
         'prepare_llm_agent_context',
         lambda agent, ctx: ctx,
     )
-    monkeypatch.setattr(
-        agent_wrapper,
-        'prepare_llm_agent_input',
-        lambda agent, ctx, node_input: None,
-    )
+    monkeypatch.setattr(agent_wrapper, 'prepare_llm_agent_input', skip_input)
     ctx = MagicMock(spec=Context)
     ic = MagicMock()
     ctx.get_invocation_context.return_value = ic

tests/unittests/workflow/test_workflow_llm_agent_interruptions.py

@@ -701,14 +701,15 @@ def get_invocation_context(self):
 
   from google.adk.workflow import _llm_agent_wrapper
 
+  async def skip_input(a, c, i):
+    return None
+
   monkeypatch.setattr(
       _llm_agent_wrapper,
       'prepare_llm_agent_context',
       lambda a, c: DummyAgentCtx(ic),
   )
-  monkeypatch.setattr(
-      _llm_agent_wrapper, 'prepare_llm_agent_input', lambda a, c, i: None
-  )
+  monkeypatch.setattr(_llm_agent_wrapper, 'prepare_llm_agent_input', skip_input)
 
   # Simulate Runner adding the event to the correct branch!
   _append_function_response(
@@ -823,14 +824,15 @@ def get_invocation_context(self):
 
   from google.adk.workflow import _llm_agent_wrapper
 
+  async def skip_input(a, c, i):
+    return None
+
   monkeypatch.setattr(
       _llm_agent_wrapper,
       'prepare_llm_agent_context',
       lambda a, c: DummyAgentCtx(ic),
   )
-  monkeypatch.setattr(
-      _llm_agent_wrapper, 'prepare_llm_agent_input', lambda a, c, i: None
-  )
+  monkeypatch.setattr(_llm_agent_wrapper, 'prepare_llm_agent_input', skip_input)
 
   # Simulate Runner adding the events to the correct branches!
   _append_function_response(