Skip to content

fix(workflow): isolate and clean up single_turn LlmAgent node_input events - #7320

Open
abhayjoshi201 wants to merge 1 commit into
google:mainfrom
abhayjoshi201:fix/single-turn-node-input-leak
Open

abhayjoshi201 wants to merge 1 commit into
google:mainfrom
abhayjoshi201:fix/single-turn-node-input-leak

Conversation

@abhayjoshi201

Copy link
Copy Markdown

Fixes #7227

Summary

When a Workflow containing single_turn LlmAgent nodes is run from a root chat agent's tool (await tool_context.run_node(workflow, ...)), branch is None and isolation_scope is None. Previously, prepare_llm_agent_input appended a synthetic Event(author='user', message=agent_input) directly to the shared in-memory ctx.session.events list without scoping it to the target node or removing it after the node finished.

During sequential multi-tool turns (fc-a -> fr-a followed by fc-b -> fr-b in the same user turn), _rearrange_events_for_async_function_responses_in_history moved fr-a next to fc-a and left the unpersisted synthetic user_event (including any inline binary/PDF parts) in result_events, leaking it verbatim into the root chat agent's LLM request. Similarly, concurrent single_turn LlmAgent nodes running via asyncio.gather on branch=None could see each other's synthetic user_event while in flight.

Why isolation_scope Is Not Used for Unscoped single_turn Nodes

Assigning a synthetic isolation_scope to unscoped single_turn nodes causes two side effects:

  1. _get_contents (_contents.py) treats any non-None isolation_scope as a task-scoped agent and invokes _build_task_input_user_content, prepending the root workflow's user_content + _SINGLE_TURN_NUDGE.
  2. During inner tool round-trips (test_inner_llm_agent_node_input_survives_tool_round_trip), NodeRunner._enrich_event stamps the node's function_call and function_response events with ctx.isolation_scope (None), so a synthetic ic.isolation_scope causes _should_include_event_in_context to drop the single_turn agent's own tool call/response events on the post-tool LLM request.

Changes

  1. src/google/adk/workflow/_llm_agent_wrapper.py:
    • Stamp user_event.node_info.path = ctx.node_path in prepare_llm_agent_input when ctx.node_path is non-empty, and return user_event.
    • In run_llm_agent_as_node, remove the injected synthetic user_event from agent_ctx.session.events in a finally: block once agent.run_async(ic) completes.
  2. src/google/adk/flows/llm_flows/context/_contents.py:
    • Thread node_path=invocation_context.node_path into _get_contents, _get_current_turn_contents, and _should_include_event_in_context.
    • Exclude synthetic node-input user events (event.author == 'user' with a non-empty event.node_info.path) whenever event.node_info.path != (node_path or '').
  3. tests/unittests/workflow/test_llm_agent_as_node.py:
    • Add test_single_turn_node_input_does_not_leak_across_sequential_tools (verifying text + inline binary PDF parts reach the worker across sequential tool calls without leaking into any root agent LLM request).
    • Add test_parallel_single_turn_nodes_only_see_own_node_input (verifying concurrent single_turn nodes under branch=None only see their own node_input).

@nanhe17 nanhe17 left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verified locally on Windows (Python 3.12, current main 044a1ec vs this branch):

  • tests/unittests/workflow/test_workflow_dynamic_nodes.py::test_inner_llm_agent_node_input_survives_tool_round_trip passes on this branch. Worth flagging for reviewers: the earlier isolation_scope-assignment approach discussed in #7227 breaks exactly this test (a scoped node request no longer sees its own unscoped FC/FR parts after a tool round trip), so the node_info.path-based filtering + cleanup design here avoids that failure mode entirely.
  • The two new tests (test_single_turn_node_input_does_not_leak_across_sequential_tools, test_parallel_single_turn_nodes_only_see_own_node_input) pass.
  • Full tests/unittests/workflow/ suite: 694 passed, 1 skipped, 5 xfailed, no failures.

One non-blocking nit: the finally cleanup uses injected_input_event in agent_ctx.session.events / .remove(), which goes through pydantic value equality. Today Event.id (a uuid) keeps two equal-content synthetic events distinct, so this is safe as written — but an identity-based check (any(e is injected_input_event ...) would be robust even if a future Event change drops or reuses ids. Not a merge blocker.

The approach also keeps real user turns untouched, since Event.node_info defaults to NodeInfo() (empty path) via default_factory and the filter only fires when ev_node_path is non-empty.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

3 participants