chore: import upstream snapshot with attribution

tests/unittests/workflow/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/unittests/workflow/test_agent_node.py

@@ -0,0 +1,71 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for BaseAgent acting as a workflow node."""
+
+from __future__ import annotations
+
+from typing import AsyncGenerator
+
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.context import Context
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.events.event import Event
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.sessions.session import Session
+import pytest
+
+
+class MockAgent(BaseAgent):
+  """A mock agent that yields predefined events."""
+
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    yield Event(author=self.name)
+    yield Event(author="sub_agent")
+
+
+@pytest.mark.asyncio
+async def test_base_agent_as_node_run():
+  """Tests that BaseAgent runs as a node and preserves event authors."""
+  agent = MockAgent(name="mock_agent")
+
+  # Setup minimal context
+  session = Session(app_name="test", user_id="user", id="session")
+  session_service = InMemorySessionService()
+  ic = InvocationContext(
+      invocation_id="inv",
+      session=session,
+      session_service=session_service,
+  )
+  ctx = Context(ic, node_path="wf")
+
+  events = []
+  async for event in agent.run(ctx=ctx, node_input=None):
+    events.append(event)
+
+  assert len(events) == 2
+
+  # First event from mock_agent
+  assert events[0].author == "mock_agent"
+  assert events[0].node_info.path == "wf"
+
+  # Second event from sub_agent
+  assert events[1].author == "sub_agent"
+  # Path should not be set by BaseAgent for sub_agent if author doesn't match agent name
+  assert not events[1].node_info.path
+
+  # Also check if ctx.event_author was updated to preserve author for NodeRunner
+  assert ctx.event_author == "sub_agent"

tests/unittests/workflow/test_dynamic_node_scheduler.py

@@ -0,0 +1,798 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for DynamicNodeScheduler.
+
+Verifies the three scheduling cases (fresh, dedup, resume) and the
+lazy event scan that reconstructs dynamic node state.
+"""
+
+from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
+
+from google.adk.agents.context import Context
+from google.adk.events.event import Event
+from google.adk.events.event import NodeInfo
+from google.adk.workflow._base_node import BaseNode
+from google.adk.workflow._dynamic_node_scheduler import DynamicNodeRun
+from google.adk.workflow._dynamic_node_scheduler import DynamicNodeScheduler
+from google.adk.workflow._dynamic_node_scheduler import DynamicNodeState
+from google.adk.workflow._node_state import NodeState
+from google.adk.workflow._node_status import NodeStatus
+from google.adk.workflow._workflow import _LoopState
+from pydantic import BaseModel
+from pydantic import ValidationError
+import pytest
+
+# --- Fixtures ---
+
+
+def _make_parent_ctx(events=None):
+  """Create a minimal parent Context with mock IC."""
+  ic = MagicMock()
+  ic.invocation_id = 'inv-1'
+  ic.session = MagicMock()
+  ic.session.state = {}
+  ic.session.events = events or []
+  ic.run_config = None
+
+  collected = []
+
+  async def _enqueue(event):
+    collected.append(event)
+
+  ic._enqueue_event = AsyncMock(side_effect=_enqueue)
+
+  ctx = MagicMock(spec=Context)
+  ctx._invocation_context = ic
+  ctx.node_path = 'wf/parent'
+  ctx.run_id = 'run-parent'
+  ctx.event_author = 'wf'
+  ctx._workflow_scheduler = None
+  ctx._output_for_ancestors = []
+  ctx._output_delegated = False
+  ctx._child_run_counters = {}
+
+  return ctx, collected
+
+
+def _make_event(
+    path='',
+    output=None,
+    interrupt_ids=None,
+    run_id=None,
+    author='node',
+    invocation_id='inv-1',
+    output_for=None,
+):
+  """Create a minimal Event for session event lists."""
+  event = MagicMock(spec=Event)
+  event.invocation_id = invocation_id
+  event.author = author
+  event.output = output
+  event.partial = False
+  event.node_info = MagicMock(spec=NodeInfo)
+  event.node_info.path = path
+  event.node_info.output_for = output_for
+  event.node_info.message_as_output = None
+  event.branch = None
+  event.isolation_scope = None
+  event.long_running_tool_ids = set(interrupt_ids) if interrupt_ids else None
+  event.content = None
+  event.actions = None
+  return event
+
+
+def _make_fr_event(fc_id, response, invocation_id='inv-1'):
+  """Create a user FR event."""
+  event = MagicMock(spec=Event)
+  event.invocation_id = invocation_id
+  event.author = 'user'
+  event.output = None
+  event.node_info = MagicMock(spec=NodeInfo)
+  event.node_info.path = ''
+  event.node_info.message_as_output = None
+  event.branch = None
+  event.isolation_scope = None
+  event.long_running_tool_ids = None
+
+  fr = MagicMock()
+  fr.id = fc_id
+  fr.response = response
+
+  part = MagicMock()
+  part.function_response = fr
+
+  content = MagicMock()
+  content.parts = [part]
+  event.content = content
+  return event
+
+
+# =========================================================================
+# _rehydrate_from_events — lazy scan
+# =========================================================================
+
+
+@pytest.mark.asyncio
+async def test_rehydrate_finds_completed_node():
+  """Scan finds output event → node marked COMPLETED."""
+  events = [
+      _make_event(
+          path='wf/parent/child@r-1',
+          output='result',
+      ),
+  ]
+  ctx, _ = _make_parent_ctx(events=events)
+  ls = _LoopState()
+  scheduler = DynamicNodeScheduler(state=ls)
+
+  scheduler._rehydrate_from_events(ctx, 'wf/parent/child@r-1')
+
+  assert 'wf/parent/child@r-1' in ls.runs
+  run = ls.runs['wf/parent/child@r-1']
+  assert run.recovered_state is not None
+  assert run.recovered_state.output == 'result'
+
+
+@pytest.mark.asyncio
+async def test_rehydrate_ignores_events_from_different_invocation():
+  """Scan ignores events with a different invocation_id."""
+  events = [
+      _make_event(
+          path='wf/parent/child@r-1',
+          output='result',
+          invocation_id='inv-different',
+      ),
+  ]
+  ctx, _ = _make_parent_ctx(events=events)
+  ctx._invocation_context.invocation_id = 'inv-current'
+  ls = _LoopState()
+  scheduler = DynamicNodeScheduler(state=ls)
+
+  scheduler._rehydrate_from_events(ctx, 'wf/parent/child@r-1')
+
+  assert 'wf/parent/child@r-1' not in ls.runs
+
+
+@pytest.mark.asyncio
+async def test_rehydrate_finds_interrupted_node():
+  """Scan finds interrupt event → node marked WAITING."""
+  events = [
+      _make_event(
+          path='wf/parent/child@r-1',
+          interrupt_ids=['fc-1'],
+      ),
+  ]
+  ctx, _ = _make_parent_ctx(events=events)
+  ls = _LoopState()
+  scheduler = DynamicNodeScheduler(state=ls)
+
+  scheduler._rehydrate_from_events(ctx, 'wf/parent/child@r-1')
+
+  assert 'wf/parent/child@r-1' in ls.runs
+  run = ls.runs['wf/parent/child@r-1']
+  assert run.recovered_state is not None
+  assert 'fc-1' in run.recovered_state.interrupt_ids
+
+
+@pytest.mark.asyncio
+async def test_rehydrate_with_target_run_id_skips_others():
+  """Scan with unique path only rehydrates that specific run."""
+  events = [
+      _make_event(
+          path='wf/parent/child@r-1',
+          output='result-1',
+      ),
+      _make_event(
+          path='wf/parent/child@r-2',
+          output='result-2',
+      ),
+  ]
+  ctx, _ = _make_parent_ctx(events=events)
+  ls = _LoopState()
+  scheduler = DynamicNodeScheduler(state=ls)
+
+  # When targeting r-2
+  scheduler._rehydrate_from_events(ctx, 'wf/parent/child@r-2')
+
+  # Then only r-2 is in state
+  assert 'wf/parent/child@r-2' in ls.runs
+  assert 'wf/parent/child@r-1' not in ls.runs
+  run = ls.runs['wf/parent/child@r-2']
+  assert run.recovered_state is not None
+  assert run.recovered_state.output == 'result-2'
+
+
+@pytest.mark.asyncio
+async def test_rehydrate_includes_delegated():
+  """Scan includes events delegated to that run."""
+  events = [
+      _make_event(
+          path='wf/parent/child@r-target/inner@r-inner',
+          output='delegated-val',
+          output_for=['wf/parent/child@r-target'],
+      ),
+  ]
+  ctx, _ = _make_parent_ctx(events=events)
+  ls = _LoopState()
+  scheduler = DynamicNodeScheduler(state=ls)
+
+  scheduler._rehydrate_from_events(ctx, 'wf/parent/child@r-target')
+
+  assert 'wf/parent/child@r-target' in ls.runs
+  run = ls.runs['wf/parent/child@r-target']
+  assert run.recovered_state is not None
+  assert run.recovered_state.output == 'delegated-val'
+
+
+@pytest.mark.asyncio
+async def test_rehydrate_resolves_interrupt_with_fr():
+  """Scan finds interrupt + FR → all resolved, ready to re-run."""
+  events = [
+      _make_event(
+          path='wf/parent/child@r-1',
+          interrupt_ids=['fc-1'],
+      ),
+      _make_fr_event('fc-1', {'approved': True}),
+  ]
+  ctx, _ = _make_parent_ctx(events=events)
+  ls = _LoopState()
+  scheduler = DynamicNodeScheduler(state=ls)
+
+  scheduler._rehydrate_from_events(ctx, 'wf/parent/child@r-1')
+
+  run = ls.runs['wf/parent/child@r-1']
+  assert run.recovered_state is not None
+  assert 'fc-1' in run.recovered_state.resolved_ids
+
+
+@pytest.mark.asyncio
+async def test_rehydrate_no_events_does_nothing():
+  """Scan with no matching events does not populate dynamic_nodes."""
+  events = [
+      _make_event(path='wf/other/node', output='x'),
+  ]
+  ctx, _ = _make_parent_ctx(events=events)
+  ls = _LoopState()
+  scheduler = DynamicNodeScheduler(state=ls)
+
+  scheduler._rehydrate_from_events(ctx, 'wf/parent/child@r-1')
+
+  assert 'wf/parent/child@r-1' not in ls.runs
+
+
+@pytest.mark.asyncio
+async def test_rehydrate_subtree_interrupt():
+  """Interrupts from nested descendants are collected."""
+  events = [
+      _make_event(
+          path='wf/parent/child@r-1/inner@r-inner',
+          interrupt_ids=['fc-deep'],
+      ),
+  ]
+  ctx, _ = _make_parent_ctx(events=events)
+  ls = _LoopState()
+  scheduler = DynamicNodeScheduler(state=ls)
+
+  scheduler._rehydrate_from_events(ctx, 'wf/parent/child@r-1')
+
+  assert 'wf/parent/child@r-1' in ls.runs
+  run = ls.runs['wf/parent/child@r-1']
+  assert run.recovered_state is not None
+  assert 'fc-deep' in run.recovered_state.interrupt_ids
+
+
+@pytest.mark.asyncio
+async def test_rehydrate_parallel_worker_interrupts():
+  """Interrupts from parallel child nodes sharing the parent's path."""
+  events = [
+      _make_event(
+          # Child has exact same path as parent
+          path='wf/parent/parallel',
+          interrupt_ids=['fc-1'],
+          run_id='r-child-1',
+      ),
+      _make_event(
+          path='wf/parent/parallel',
+          interrupt_ids=['fc-2'],
+          run_id='r-child-2',
+      ),
+  ]
+  ctx, _ = _make_parent_ctx(events=events)
+  ls = _LoopState()
+  scheduler = DynamicNodeScheduler(state=ls)
+
+  # Rehydrate the parent which has run_id 'r-parent'
+  scheduler._rehydrate_from_events(ctx, 'wf/parent/parallel')
+
+  assert 'wf/parent/parallel' in ls.runs
+  run = ls.runs['wf/parent/parallel']
+  assert run.recovered_state is not None
+  assert 'fc-1' in run.recovered_state.interrupt_ids
+  assert 'fc-2' in run.recovered_state.interrupt_ids
+
+
+@pytest.mark.asyncio
+async def test_rehydrate_output_for_delegation():
+  """Output via output_for delegation is recognized."""
+  events = [
+      _make_event(
+          path='wf/parent/child@r-1/inner@r-inner',
+          output='delegated',
+          output_for=['wf/parent/child@r-1'],
+      ),
+  ]
+  ctx, _ = _make_parent_ctx(events=events)
+  ls = _LoopState()
+  scheduler = DynamicNodeScheduler(state=ls)
+
+  scheduler._rehydrate_from_events(ctx, 'wf/parent/child@r-1')
+
+  run = ls.runs['wf/parent/child@r-1']
+  assert run.recovered_state is not None
+  assert run.recovered_state.output == 'delegated'
+
+
+# =========================================================================
+# __call__ — dispatch logic
+# =========================================================================
+
+
+# =========================================================================
+# DefaultNodeScheduler — standalone scheduler
+# =========================================================================
+
+
+@pytest.mark.asyncio
+async def test_fresh_execution_runs_node():
+  """DefaultNodeScheduler runs a fresh node just like DynamicNodeScheduler."""
+
+  class _Child(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      yield f'ct: {node_input}'
+
+  ctx, _ = _make_parent_ctx()
+  tracker = DynamicNodeScheduler(state=DynamicNodeState())
+
+  mock_child_ctx = MagicMock(spec=Context)
+  mock_child_ctx.error = None
+  mock_child_ctx.interrupt_ids = set()
+  mock_child_ctx.output = 'ct: data'
+  mock_child_ctx.actions = MagicMock()
+  mock_child_ctx.actions.transfer_to_agent = None
+  ctx._run_node_standalone = AsyncMock(return_value=mock_child_ctx)
+
+  child_ctx = await tracker(
+      ctx,
+      _Child(name='child'),
+      'data',
+      node_name='child',
+      run_id='1',
+  )
+
+  assert child_ctx.output == 'ct: data'
+
+
+@pytest.mark.asyncio
+async def test_completed_dedup_returns_cached():
+  """DefaultNodeScheduler returns cached output for completed nodes."""
+  ctx, _ = _make_parent_ctx()
+  tracker = DynamicNodeScheduler(state=DynamicNodeState())
+
+  # Pre-populate state as if node already completed.
+  from google.adk.workflow.utils._rehydration_utils import _ChildScanState
+
+  tracker._state.runs['wf/parent/child@r-1'] = DynamicNodeRun(
+      state=NodeState(run_id='r-1'),
+      recovered_state=_ChildScanState(
+          run_id='r-1',
+          output='cached',
+      ),
+  )
+
+  child_ctx = await tracker(
+      ctx,
+      BaseNode(name='child'),
+      'input',
+      node_name='child',
+      run_id='r-1',
+  )
+
+  assert child_ctx.output == 'cached'
+
+
+@pytest.mark.asyncio
+async def test_concurrent_dedup_returns_running_task():
+  """Scheduler deduplicates concurrent executions of the same running task."""
+  import asyncio
+
+  ctx, _ = _make_parent_ctx()
+  tracker = DynamicNodeScheduler(state=DynamicNodeState())
+
+  # Mock an active running task (not done yet!)
+  running_task = asyncio.Future()
+
+  tracker._state.runs['wf/parent/child@r-1'] = DynamicNodeRun(
+      state=NodeState(run_id='r-1'),
+      task=running_task,
+  )
+
+  # Dispatch the scheduler in the background
+  scheduler_task = asyncio.create_task(
+      tracker(
+          ctx,
+          BaseNode(name='child'),
+          'input',
+          node_name='child',
+          run_id='r-1',
+      )
+  )
+
+  # Let the event loop run one tick to execute the scheduler interception
+  await asyncio.sleep(0)
+
+  # Resolve the running task dynamically
+  mock_context = MagicMock(spec=Context)
+  running_task.set_result(mock_context)
+
+  res_ctx = await scheduler_task
+  assert res_ctx is mock_context
+
+
+@pytest.mark.asyncio
+async def test_waiting_resolved_resumes_node():
+  """DefaultNodeScheduler re-runs nodes with resolved interrupts."""
+
+  class _Resumable(BaseNode):
+    rerun_on_resume: bool = True
+
+    async def _run_impl(self, *, ctx, node_input):
+      if ctx.resume_inputs and 'fc-1' in ctx.resume_inputs:
+        yield f'resumed: {ctx.resume_inputs["fc-1"]}'
+        return
+      yield 'should not reach here'
+
+  ctx, _ = _make_parent_ctx()
+  tracker = DynamicNodeScheduler(state=DynamicNodeState())
+
+  # Pre-populate state as if node interrupted and was resolved.
+  from google.adk.workflow.utils._rehydration_utils import _ChildScanState
+
+  tracker._state.runs['wf/parent/child@r-1'] = DynamicNodeRun(
+      state=NodeState(run_id='r-1'),
+      recovered_state=_ChildScanState(
+          run_id='r-1',
+          interrupt_ids={'fc-1'},
+          resolved_ids={'fc-1'},
+          resolved_responses={'fc-1': 'approved'},
+      ),
+  )
+
+  mock_child_ctx = MagicMock(spec=Context)
+  mock_child_ctx.error = None
+  mock_child_ctx.interrupt_ids = set()
+  mock_child_ctx.output = 'resumed: approved'
+  mock_child_ctx.actions = MagicMock()
+  mock_child_ctx.actions.transfer_to_agent = None
+  ctx._run_node_standalone = AsyncMock(return_value=mock_child_ctx)
+
+  child_ctx = await tracker(
+      ctx,
+      _Resumable(name='child'),
+      'input',
+      node_name='child',
+      run_id='r-1',
+  )
+
+  assert child_ctx.output == 'resumed: approved'
+
+
+@pytest.mark.asyncio
+async def test_waiting_unresolved_propagates_interrupts():
+  """DefaultNodeScheduler propagates unresolved interrupts."""
+  ctx, _ = _make_parent_ctx()
+  tracker = DynamicNodeScheduler(state=DynamicNodeState())
+
+  from google.adk.workflow.utils._rehydration_utils import _ChildScanState
+
+  tracker._state.runs['wf/parent/child@r-1'] = DynamicNodeRun(
+      state=NodeState(run_id='r-1'),
+      recovered_state=_ChildScanState(
+          run_id='r-1',
+          interrupt_ids={'fc-1'},
+      ),
+  )
+
+  child_ctx = await tracker(
+      ctx,
+      BaseNode(name='child'),
+      'input',
+      node_name='child',
+      run_id='r-1',
+  )
+
+  assert child_ctx.interrupt_ids == {'fc-1'}
+  assert 'fc-1' in tracker._state.interrupt_ids
+
+
+@pytest.mark.asyncio
+async def test_calling_waiting_node_without_rerun_raises_value_error():
+  """Calling a dynamic node that is waiting for output with rerun_on_resume=False raises ValueError."""
+
+  # Given a dynamic node waiting for output with rerun_on_resume=False
+  class _WaitingNode(BaseNode):
+    wait_for_output: bool = True
+
+    async def _run_impl(self, *, ctx, node_input):
+      yield 'should not reach here'
+
+  ctx, _ = _make_parent_ctx()
+  ls = _LoopState()
+  from google.adk.workflow.utils._rehydration_utils import _ChildScanState
+
+  ls.runs['wf/parent/child@r-1'] = DynamicNodeRun(
+      state=NodeState(run_id='r-1'),
+      recovered_state=_ChildScanState(
+          run_id='r-1',
+          interrupt_ids={'pause_req'},
+          resolved_ids={'pause_req'},
+      ),
+  )
+  scheduler = DynamicNodeScheduler(state=ls)
+
+  # When it is called again
+  # Then it raises ValueError
+  with pytest.raises(
+      ValueError, match='is waiting for output but was called again'
+  ):
+    await scheduler(
+        ctx,
+        _WaitingNode(name='child'),
+        'input',
+        node_name='child',
+        run_id='r-1',
+    )
+
+
+def test_get_dynamic_tasks_excludes_done_tasks():
+  """get_dynamic_tasks should not return completed tasks (regression for #6082)."""
+  import asyncio
+
+  loop = asyncio.new_event_loop()
+  try:
+
+    async def _done():
+      return None
+
+    done_task = loop.run_until_complete(
+        asyncio.ensure_future(_done(), loop=loop)
+    )
+    running_coro = asyncio.sleep(9999)
+    running_task = loop.create_task(running_coro)
+
+    state = DynamicNodeState()
+    state.runs['path/done@r-1'] = DynamicNodeRun(
+        state=NodeState(run_id='r-1'),
+        task=done_task,
+    )
+    state.runs['path/running@r-2'] = DynamicNodeRun(
+        state=NodeState(run_id='r-2'),
+        task=running_task,
+    )
+    state.runs['path/no-task@r-3'] = DynamicNodeRun(
+        state=NodeState(run_id='r-3'),
+        task=None,
+    )
+
+    tasks = state.get_dynamic_tasks()
+
+    assert tasks == [running_task]
+    running_task.cancel()
+  finally:
+    loop.close()
+
+
+class _ModelA(BaseModel):
+  x: int
+
+
+@pytest.mark.asyncio
+async def test_runtime_schema_validation_passes():
+  """Tests that runtime schema validation passes when input matches schema."""
+  ctx, _ = _make_parent_ctx()
+  ls = _LoopState()
+  scheduler = DynamicNodeScheduler(state=ls)
+
+  node = BaseNode(name='child', input_schema=_ModelA)
+
+  # We mock _run_node_internal to avoid full execution, we only care about validation in __call__
+  scheduler._run_node_internal = AsyncMock(return_value=MagicMock(spec=Context))
+
+  await scheduler(
+      ctx,
+      node,
+      {'x': 1},
+      node_name='child',
+      run_id='1',
+  )
+  # Should not raise
+
+
+@pytest.mark.asyncio
+async def test_runtime_schema_validation_raises():
+  """Tests that runtime schema validation raises when input mismatches schema."""
+  ctx, _ = _make_parent_ctx()
+  ls = _LoopState()
+  scheduler = DynamicNodeScheduler(state=ls)
+
+  node = BaseNode(name='child', input_schema=_ModelA)
+
+  with pytest.raises(ValidationError):
+    await scheduler(
+        ctx,
+        node,
+        {'x': 'string'},  # Invalid type for x
+        node_name='child',
+        run_id='1',
+    )
+
+
+@pytest.mark.asyncio
+async def test_runtime_schema_validation_missing_schema_passes():
+  """Tests that runtime schema validation passes when no schema is defined."""
+  ctx, _ = _make_parent_ctx()
+  ls = _LoopState()
+  scheduler = DynamicNodeScheduler(state=ls)
+
+  node = BaseNode(name='child')  # No input schema
+
+  scheduler._run_node_internal = AsyncMock(return_value=MagicMock(spec=Context))
+
+  await scheduler(
+      ctx,
+      node,
+      {'x': 1},
+      node_name='child',
+      run_id='1',
+  )
+  # Should not raise
+
+
+@pytest.mark.asyncio
+async def test_runtime_schema_validation_content_fallback():
+  """Tests that runtime schema validation handles Content objects by extraction."""
+  ctx, _ = _make_parent_ctx()
+  ls = _LoopState()
+  scheduler = DynamicNodeScheduler(state=ls)
+
+  node = BaseNode(name='child', input_schema=_ModelA)
+
+  scheduler._run_node_internal = AsyncMock(return_value=MagicMock(spec=Context))
+
+  from google.genai import types
+
+  msg = types.Content(parts=[types.Part(text='{"x": 1}')], role='user')
+
+  await scheduler(
+      ctx,
+      node,
+      msg,
+      node_name='child',
+      run_id='1',
+  )
+  # Should not raise
+
+
+# =========================================================================
+# Replay Sequence Ordering preservation for Dynamic Nodes
+# =========================================================================
+
+
+@pytest.mark.asyncio
+async def test_dynamic_node_replay_ordering_preserved(
+    request: pytest.FixtureRequest,
+):
+  """Test that parallel dynamic nodes maintain their chronological completion order during replay."""
+  import asyncio
+
+  from google.adk.events.request_input import RequestInput
+  from google.adk.workflow import node
+  from google.adk.workflow import START
+  from google.adk.workflow._workflow import Workflow
+  from google.genai import types
+
+  from .. import testing_utils
+
+  execution_order = []
+  recorded_winner_vals = []
+
+  @node
+  async def source_a(*, ctx, node_input):
+    await asyncio.sleep(0.1)
+    execution_order.append('source_a_executed')
+    yield 'result_a'
+
+  @node
+  async def source_b(*, ctx, node_input):
+    # No sleep, completes immediately in Run 1
+    execution_order.append('source_b_executed')
+    yield 'result_b'
+
+  @node(rerun_on_resume=True)
+  async def hitl_node(*, ctx, node_input):
+    if 'req_h' not in ctx.resume_inputs:
+      yield RequestInput(interrupt_id='req_h', message='input h')
+      return
+    execution_order.append(f'hitl_resumed_with_{node_input}')
+    yield f'h_{node_input}'
+
+  @node(rerun_on_resume=True)
+  async def parent(*, ctx, node_input):
+    completed_order = []
+
+    async def run_and_record(node_func, run_id):
+      res = await ctx.run_node(node_func, run_id=run_id)
+      completed_order.append(res)
+      return res
+
+    task_a = asyncio.create_task(run_and_record(source_a, 'a'))
+    task_b = asyncio.create_task(run_and_record(source_b, 'b'))
+
+    await asyncio.wait([task_a, task_b], return_when=asyncio.ALL_COMPLETED)
+
+    winner_val = completed_order[0]
+    recorded_winner_vals.append(winner_val)
+
+    await ctx.run_node(hitl_node, node_input=winner_val, run_id='h')
+
+  wf_name = request.node.name.replace('[', '_').replace(']', '')
+  agent = Workflow(name=wf_name, edges=[(START, parent)])
+  runner = testing_utils.InMemoryRunner(node=agent)
+
+  # Run 1: source_b finishes first, source_a finishes second. hitl_node interrupts.
+  events1 = await runner.run_async(testing_utils.get_user_content('start'))
+
+  req_events = [e for e in events1 if e.long_running_tool_ids]
+  assert len(req_events) == 1
+  assert execution_order == ['source_b_executed', 'source_a_executed']
+
+  invocation_id = events1[0].invocation_id
+
+  # Clear execution order to track replay/resume behavior accurately
+  execution_order.clear()
+  recorded_winner_vals.clear()
+
+  # Run 2: Resume with response
+  resume_payload = types.Content(
+      role='user',
+      parts=[
+          types.Part(
+              function_response=types.FunctionResponse(  # type: ignore[call-arg]  # Third-party SDK signature
+                  id='req_h',
+                  name='user_input',
+                  response={'text': 'response_h'},
+              )
+          ),
+      ],
+  )
+
+  await runner.run_async(
+      new_message=resume_payload, invocation_id=invocation_id
+  )
+
+  # Assert source_a and source_b were replayed from cache in exact historical order,
+  # ensuring winner_val correctly resolves to 'result_b' without re-execution.
+  assert recorded_winner_vals == ['result_b']

tests/unittests/workflow/test_dynamic_use_as_output.py

@@ -0,0 +1,588 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for ctx.run_node(use_as_output=True) dynamic terminal paths."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+from typing import AsyncGenerator
+
+from google.adk.agents.context import Context
+from google.adk.apps.app import App
+from google.adk.apps.app import ResumabilityConfig
+from google.adk.events.event import Event
+from google.adk.events.request_input import RequestInput
+from google.adk.workflow import BaseNode
+from google.adk.workflow import FunctionNode
+from google.adk.workflow import node
+from google.adk.workflow import START
+from google.adk.workflow._workflow import Workflow
+from google.genai import types
+from pydantic import ConfigDict
+import pytest
+from typing_extensions import override
+
+from .. import testing_utils
+from .workflow_testing_utils import get_outputs as _get_outputs
+
+
+def _make_app(name: str, agent: Workflow, resumable: bool) -> App:
+  return App(
+      name=name,
+      root_agent=agent,
+      resumability_config=(
+          ResumabilityConfig(is_resumable=True) if resumable else None
+      ),
+  )
+
+
+# ---------------------------------------------------------------------------
+
+# FunctionNode delegates to FunctionNode
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize('resumable', [True, False])
+@pytest.mark.asyncio
+async def test_use_as_output_function_to_function(
+    request: pytest.FixtureRequest, resumable: bool
+):
+  """Node A delegates output to dynamic child B via use_as_output=True.
+
+  Only B's output event should appear; A should not emit a duplicate.
+  """
+
+  def func_b() -> str:
+    return 'from_b'
+
+  node_b = FunctionNode(func=func_b)
+
+  async def func_a(ctx: Context) -> str:
+    return await ctx.run_node(node_b, use_as_output=True)
+
+  node_a = FunctionNode(func=func_a, rerun_on_resume=True)
+
+  wf_name = request.node.name.replace('[', '_').replace(']', '')
+  agent = Workflow(name=wf_name, edges=[(START, node_a)])
+  app = _make_app(wf_name, agent, resumable)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  # Only one output event (from B). A's output is suppressed.
+  outputs = _get_outputs(events)
+  assert outputs == ['from_b']
+
+
+# ---------------------------------------------------------------------------
+# FunctionNode delegates to Workflow
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize('resumable', [True, False])
+@pytest.mark.asyncio
+async def test_use_as_output_function_to_workflow(
+    request: pytest.FixtureRequest, resumable: bool
+):
+  """Node A delegates output to a Workflow B via use_as_output=True.
+
+  The Workflow's terminal node output should be used as A's output.
+  """
+
+  def step_1() -> str:
+    return 'step_1_done'
+
+  def step_2(node_input: str) -> str:
+    return f'final:{node_input}'
+
+  inner_wf = Workflow(
+      name='inner_wf',
+      edges=[
+          (START, step_1),
+          (step_1, step_2),
+      ],
+  )
+
+  async def func_a(ctx: Context):
+    return await ctx.run_node(inner_wf, use_as_output=True)
+
+  node_a = FunctionNode(func=func_a, rerun_on_resume=True)
+
+  wf_name = request.node.name.replace('[', '_').replace(']', '')
+  agent = Workflow(name=wf_name, edges=[(START, node_a)])
+  app = _make_app(wf_name, agent, resumable)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  # step_2 is the terminal node; its output is func_a's output.
+  # step_1's output is intermediate (not terminal), so only step_2 appears.
+  outputs = _get_outputs(events)
+  assert 'final:step_1_done' in outputs
+  # func_a should NOT have a duplicate output event.
+  assert outputs.count('final:step_1_done') == 1
+
+
+# ---------------------------------------------------------------------------
+# Custom BaseNode delegates to FunctionNode
+# ---------------------------------------------------------------------------
+
+
+class _DelegatingNode(BaseNode):
+  """A custom BaseNode that delegates output via use_as_output."""
+
+  model_config = ConfigDict(arbitrary_types_allowed=True)
+  delegate: BaseNode
+
+  @override
+  async def _run_impl(
+      self, *, ctx: Context, node_input: Any
+  ) -> AsyncGenerator[Any, None]:
+    await ctx.run_node(self.delegate, use_as_output=True)
+    # Intentionally yield nothing — output comes from delegate.
+    return
+    yield  # Make this an async generator
+
+
+@pytest.mark.parametrize('resumable', [True, False])
+@pytest.mark.asyncio
+async def test_use_as_output_custom_node_to_function(
+    request: pytest.FixtureRequest, resumable: bool
+):
+  """Custom BaseNode delegates output to dynamic child via use_as_output."""
+
+  def func_b() -> str:
+    return 'delegated_output'
+
+  node_b = FunctionNode(func=func_b)
+  node_a = _DelegatingNode(
+      name='delegator', delegate=node_b, rerun_on_resume=True
+  )
+
+  wf_name = request.node.name.replace('[', '_').replace(']', '')
+  agent = Workflow(name=wf_name, edges=[(START, node_a)])
+  app = _make_app(wf_name, agent, resumable)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  outputs = _get_outputs(events)
+  assert outputs == ['delegated_output']
+
+
+# ---------------------------------------------------------------------------
+# Multiple use_as_output=True calls (fan-out delegation)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize('resumable', [True, False])
+@pytest.mark.asyncio
+async def test_use_as_output_multiple_disallowed(
+    request: pytest.FixtureRequest, resumable: bool
+):
+  """V2 engine forbids calling use_as_output=True multiple times from the same node."""
+
+  def func_b() -> str:
+    return 'from_b'
+
+  def func_c() -> str:
+    return 'from_c'
+
+  node_b = FunctionNode(func=func_b)
+  node_c = FunctionNode(func=func_c)
+
+  async def func_a(ctx: Context):
+    await ctx.run_node(node_b, use_as_output=True)
+    # V2 engine should throw ValueError on second call
+    with pytest.raises(
+        ValueError, match='already has a use_as_output delegate'
+    ):
+      await ctx.run_node(node_c, use_as_output=True)
+    return 'failure_as_expected'
+
+  node_a = FunctionNode(func=func_a, rerun_on_resume=True)
+
+  wf_name = request.node.name.replace('[', '_').replace(']', '')
+  agent = Workflow(name=wf_name, edges=[(START, node_a)])
+  app = _make_app(wf_name, agent, resumable)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  outputs = _get_outputs(events)
+  # func_a's output is suppressed because it successfully delegated to func_b
+  # before failing on func_c. So only func_b's output appears.
+  assert outputs == ['from_b']
+
+
+# ---------------------------------------------------------------------------
+# Nested dynamic delegation (A → B → C, all use_as_output=True)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize('resumable', [True, False])
+@pytest.mark.asyncio
+async def test_use_as_output_nested_delegation(
+    request: pytest.FixtureRequest, resumable: bool
+):
+  """Chained delegation: A delegates to B, B delegates to C.
+
+  Only C's output should appear; both A's and B's outputs are suppressed.
+  """
+
+  def func_c() -> str:
+    return 'from_c'
+
+  node_c = FunctionNode(func=func_c)
+
+  async def func_b(ctx: Context) -> str:
+    return await ctx.run_node(node_c, use_as_output=True)
+
+  node_b = FunctionNode(func=func_b, rerun_on_resume=True)
+
+  async def func_a(ctx: Context) -> str:
+    return await ctx.run_node(node_b, use_as_output=True)
+
+  node_a = FunctionNode(func=func_a, rerun_on_resume=True)
+
+  wf_name = request.node.name.replace('[', '_').replace(']', '')
+  agent = Workflow(name=wf_name, edges=[(START, node_a)])
+  app = _make_app(wf_name, agent, resumable)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  outputs = _get_outputs(events)
+  assert outputs == ['from_c']
+
+
+@pytest.mark.parametrize('resumable', [True, False])
+@pytest.mark.asyncio
+async def test_use_as_output_nested_delegation_with_downstream(
+    request: pytest.FixtureRequest, resumable: bool
+):
+  """Chained delegation with a downstream node that consumes the output.
+
+  A delegates to B, B delegates to C. D is downstream of A and should
+  receive C's output as its node_input.
+  """
+
+  def func_c() -> str:
+    return 'from_c'
+
+  node_c = FunctionNode(func=func_c)
+
+  async def func_b(ctx: Context) -> str:
+    return await ctx.run_node(node_c, use_as_output=True)
+
+  node_b = FunctionNode(func=func_b, rerun_on_resume=True)
+
+  async def func_a(ctx: Context) -> str:
+    return await ctx.run_node(node_b, use_as_output=True)
+
+  node_a = FunctionNode(func=func_a, rerun_on_resume=True)
+
+  def func_d(node_input: str) -> str:
+    return f'received:{node_input}'
+
+  wf_name = request.node.name.replace('[', '_').replace(']', '')
+  agent = Workflow(
+      name=wf_name,
+      edges=[(START, node_a), (node_a, func_d)],
+  )
+  app = _make_app(wf_name, agent, resumable)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  outputs = _get_outputs(events)
+  assert 'received:from_c' in outputs
+
+
+# ---------------------------------------------------------------------------
+# use_as_output=False (default) still emits both events
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize('resumable', [True, False])
+@pytest.mark.asyncio
+async def test_without_use_as_output_emits_both(
+    request: pytest.FixtureRequest, resumable: bool
+):
+  """Without use_as_output, both parent and child emit output events."""
+
+  def func_b() -> str:
+    return 'from_b'
+
+  node_b = FunctionNode(func=func_b)
+
+  async def func_a(ctx: Context) -> str:
+    return await ctx.run_node(node_b)
+
+  node_a = FunctionNode(func=func_a, rerun_on_resume=True)
+
+  wf_name = request.node.name.replace('[', '_').replace(']', '')
+  agent = Workflow(name=wf_name, edges=[(START, node_a)])
+  app = _make_app(wf_name, agent, resumable)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  # Both B and A emit output events (no dedup).
+  outputs = _get_outputs(events)
+  assert outputs == ['from_b', 'from_b']
+
+
+# ---------------------------------------------------------------------------
+# use_as_output with Workflow containing HITL
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_use_as_output_workflow_with_hitl(
+    request: pytest.FixtureRequest,
+):
+  """Dynamic Workflow child with HITL node pauses and resumes correctly.
+
+  Requires resumable=True because the inner Workflow's internal node state
+  (hitl_step in WAITING) must be persisted across invocations. Without
+  resumability, re-spawning the dynamic Workflow starts its graph fresh.
+
+  Flow:
+    1. func_a calls ctx.run_node(inner_wf, use_as_output=True)
+    2. inner_wf's hitl_step yields RequestInput → workflow pauses
+    3. User responds → workflow resumes
+    4. inner_wf completes, func_a re-runs with cached result
+    5. Only inner_wf's terminal node output appears
+  """
+  resumable = True
+
+  async def hitl_step():
+    yield RequestInput(
+        interrupt_id='wf_req1',
+        message='enter value',
+    )
+
+  def final_step(node_input: Any) -> str:
+    text = node_input['text'] if isinstance(node_input, dict) else node_input
+    return f'final:{text}'
+
+  inner_wf = Workflow(
+      name='inner_wf',
+      edges=[
+          (START, hitl_step),
+          (hitl_step, final_step),
+      ],
+  )
+
+  async def func_a(ctx: Context):
+    return await ctx.run_node(inner_wf, use_as_output=True)
+
+  node_a = FunctionNode(func=func_a, rerun_on_resume=True)
+
+  wf_name = request.node.name.replace('[', '_').replace(']', '')
+  agent = Workflow(name=wf_name, edges=[(START, node_a)])
+  app = _make_app(wf_name, agent, resumable)
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  # Run 1: Should pause at hitl_step.
+  events1 = await runner.run_async(testing_utils.get_user_content('start'))
+  outputs1 = _get_outputs(events1)
+  assert outputs1 == []
+
+  # Resume with user response.
+  invocation_id = events1[0].invocation_id
+  resume_payload = testing_utils.UserContent(
+      types.Part(
+          function_response=types.FunctionResponse(
+              id='wf_req1',
+              name='user_input',
+              response={'text': 'hello'},
+          )
+      )
+  )
+  events2 = await runner.run_async(
+      new_message=resume_payload, invocation_id=invocation_id
+  )
+
+  outputs2 = _get_outputs(events2)
+  assert 'final:hello' in outputs2
+  assert outputs2.count('final:hello') == 1
+
+
+@pytest.mark.asyncio
+async def test_use_as_output_static_node_not_rerun_on_resume(
+    request: pytest.FixtureRequest,
+):
+  """Static node delegating output to dynamic child is not re-run on resume.
+
+  Setup:
+    - wf with node_a (FunctionNode) and hitl_step.
+    - node_a calls ctx.run_node(node_b, use_as_output=True).
+    - hitl_step yields RequestInput to pause.
+  Act:
+    - Run 1: Run workflow, node_a executes, hitl_step pauses.
+    - Run 2: Resume with user response for hitl_step.
+  Assert:
+    - Run 1: node_b output is emitted as node_a's output.
+    - Run 2: node_a is NOT re-run (run_count remains 1).
+  """
+
+  run_count = 0
+
+  @node
+  def node_b() -> str:
+    return 'from_b'
+
+  @node(rerun_on_resume=True)
+  async def node_a(ctx: Context):
+    nonlocal run_count
+    run_count += 1
+    return await ctx.run_node(node_b, use_as_output=True)
+
+  node_a.wait_for_output = True
+
+  @node
+  async def hitl_step():
+    yield RequestInput(
+        interrupt_id='pause_req',
+        message='pause',
+    )
+
+  @node
+  def final_step(node_input: Any) -> None:
+    return None
+
+  wf_name = request.node.name.replace('[', '_').replace(']', '')
+  agent = Workflow(
+      name=wf_name,
+      edges=[
+          (START, node_a),
+          (START, hitl_step),
+          (hitl_step, final_step),
+      ],
+  )
+
+  runner = testing_utils.InMemoryRunner(root_agent=agent)
+
+  events1 = await runner.run_async(testing_utils.get_user_content('start'))
+
+  assert run_count == 1
+  outputs1 = [e.output for e in events1 if e.output]
+  assert 'from_b' in outputs1
+
+  invocation_id = events1[0].invocation_id
+  resume_payload = testing_utils.UserContent(
+      types.Part(
+          function_response=types.FunctionResponse(
+              id='pause_req',
+              name='user_input',
+              response={'text': 'continue'},
+          )
+      )
+  )
+
+  events2 = await runner.run_async(
+      new_message=resume_payload, invocation_id=invocation_id
+  )
+
+  assert run_count == 1
+
+
+@pytest.mark.asyncio
+async def test_use_as_output_instance_isolation(request: pytest.FixtureRequest):
+  """Output delegation is isolated to specific dynamic instances.
+
+  Setup:
+    - wf with node_a calling node_b twice in parallel with different run_ids ('b1', 'b2').
+    - node_b yields RequestInput to pause.
+  Act:
+    - Run 1: Run workflow, both node_b instances pause.
+    - Run 2: Resume only the 'b1' instance.
+  Assert:
+    - Run 1: Both instances run (run_count_b is 2).
+    - Run 2: Only instance 'b1' completes and delegates output.
+      Instance 'b2' remains paused and does not emit output.
+  """
+  run_count_b = 0
+
+  @node
+  def node_c() -> str:
+    return 'from_c'
+
+  @node(rerun_on_resume=True)
+  async def node_b(ctx: Context):
+    nonlocal run_count_b
+
+    interrupt_id = f'pause_{ctx.node_path}'
+    if ctx.resume_inputs and interrupt_id in ctx.resume_inputs:
+      response = ctx.resume_inputs[interrupt_id]
+      if (
+          response
+          and isinstance(response, dict)
+          and response.get('text') == 'continue'
+      ):
+        await ctx.run_node(node_c, use_as_output=True)
+        return
+
+    run_count_b += 1
+    yield RequestInput(
+        interrupt_id=interrupt_id,
+        message='pause',
+    )
+
+  @node(rerun_on_resume=True)
+  async def node_a(ctx: Context):
+    task1 = ctx.run_node(node_b, run_id='b1')
+    task2 = ctx.run_node(node_b, run_id='b2')
+    await asyncio.gather(task1, task2)
+
+  wf_name = request.node.name.replace('[', '_').replace(']', '')
+  agent = Workflow(name=wf_name, edges=[(START, node_a)])
+
+  runner = testing_utils.InMemoryRunner(root_agent=agent)
+
+  # Given the workflow is started with parallel instances of node_b
+  events1 = await runner.run_async(testing_utils.get_user_content('start'))
+
+  # Then both instances execute and pause
+  assert run_count_b == 2
+
+  # Find the interrupt ID for instance b1
+  interrupt_id_b1 = None
+  for event in events1:
+    if event.long_running_tool_ids:
+      for interrupt_id in event.long_running_tool_ids:
+        if 'b1' in interrupt_id:
+          interrupt_id_b1 = interrupt_id
+          break
+
+  assert interrupt_id_b1 is not None
+
+  # When resuming only instance b1
+  invocation_id = events1[0].invocation_id
+  resume_payload = testing_utils.UserContent(
+      types.Part(
+          function_response=types.FunctionResponse(
+              id=interrupt_id_b1,
+              name='user_input',
+              response={'text': 'continue'},
+          )
+      )
+  )
+
+  events2 = await runner.run_async(
+      new_message=resume_payload, invocation_id=invocation_id
+  )
+
+  # Then node_b is not re-run, and only one output is emitted from node_c
+  assert run_count_b == 2
+
+  outputs2 = [e.output for e in events2 if e.output]
+  assert outputs2.count('from_c') == 1

tests/unittests/workflow/test_graph.py

@@ -0,0 +1,96 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for Graph validation and routing."""
+
+import logging
+
+from google.adk.workflow import Edge
+from google.adk.workflow import START
+from google.adk.workflow._graph import DEFAULT_ROUTE
+from google.adk.workflow._graph import Graph
+
+from .workflow_testing_utils import TestingNode
+
+
+def test_valid_graph() -> None:
+  """Tests that a valid graph passes validation."""
+  node_a = TestingNode(name='NodeA')
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+      ],
+  )
+  graph.validate_graph()  # Should not raise
+
+
+def test_get_next_pending_nodes() -> None:
+  """Tests that get_next_pending_nodes returns correct nodes based on routes."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  node_c = TestingNode(name='NodeC')
+  node_d = TestingNode(name='NodeD')
+
+  graph = Graph(
+      edges=[
+          Edge(from_node=node_a, to_node=node_b),  # Unconditional
+          Edge(from_node=node_a, to_node=node_c, route='route1'),  # Conditional
+          Edge(
+              from_node=node_a, to_node=node_d, route=DEFAULT_ROUTE
+          ),  # Default
+      ],
+  )
+
+  # Test unconditional edge triggered
+  next_nodes = graph.get_next_pending_nodes('NodeA', routes_to_match=None)
+  assert set(next_nodes) == {'NodeB', 'NodeD'}
+
+  # Test specific route matched
+  next_nodes = graph.get_next_pending_nodes('NodeA', routes_to_match='route1')
+  assert set(next_nodes) == {'NodeB', 'NodeC'}
+
+  # Test unmatched route falls back to default
+  next_nodes = graph.get_next_pending_nodes(
+      'NodeA', routes_to_match='unknown_route'
+  )
+  assert set(next_nodes) == {'NodeB', 'NodeD'}
+
+  # Test list of routes to match
+  next_nodes = graph.get_next_pending_nodes(
+      'NodeA', routes_to_match=['route1', 'unknown_route']
+  )
+  assert set(next_nodes) == {'NodeB', 'NodeC'}
+
+
+def test_get_next_pending_nodes_unmatched_route_warning(caplog) -> None:
+  """Tests that a warning is logged when a route is unmatched and there's no DEFAULT_ROUTE."""
+  node_a = TestingNode(name='NodeA')
+  node_c = TestingNode(name='NodeC')
+
+  graph = Graph(
+      edges=[
+          Edge(from_node=node_a, to_node=node_c, route='route1'),
+      ],
+  )
+
+  with caplog.at_level(logging.WARNING):
+    next_nodes = graph.get_next_pending_nodes(
+        'NodeA', routes_to_match='unknown_route'
+    )
+
+  assert not next_nodes
+  assert any(
+      'has conditional/DEFAULT edges but none were matched' in record.message
+      for record in caplog.records
+  )

tests/unittests/workflow/test_join_node.py

@@ -0,0 +1,277 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Testings for the JoinNode."""
+
+from google.adk import workflow
+from google.adk.apps import app
+from google.adk.workflow import _base_node as base_node
+from google.adk.workflow import _graph as workflow_graph
+from google.adk.workflow import _join_node as join_node
+from google.adk.workflow import START
+from google.adk.workflow._workflow import Workflow
+from pydantic import BaseModel
+import pytest
+
+from . import workflow_testing_utils
+from .. import testing_utils
+
+
+def _build_join_node_workflow(
+    request: pytest.FixtureRequest,
+) -> tuple[
+    workflow_testing_utils.InputCapturingNode, testing_utils.InMemoryRunner
+]:
+  """Builds a workflow with a JoinNode."""
+  node_a = workflow_testing_utils.TestingNode(
+      name='NodeA', output={'a': 1, 'b': 1}
+  )
+  node_b = workflow_testing_utils.TestingNode(name='NodeB', output={'b': 2})
+  node_join = join_node.JoinNode(name='NodeJoin')
+  node_capture = workflow_testing_utils.InputCapturingNode(name='NodeCapture')
+  agent = workflow.Workflow(
+      name='test_join_node',
+      edges=[
+          workflow_graph.Edge(from_node=base_node.START, to_node=node_a),
+          workflow_graph.Edge(from_node=base_node.START, to_node=node_b),
+          workflow_graph.Edge(from_node=node_a, to_node=node_join),
+          workflow_graph.Edge(from_node=node_b, to_node=node_join),
+          workflow_graph.Edge(from_node=node_join, to_node=node_capture),
+      ],
+  )
+  app_instance = app.App(
+      name=request.function.__name__,
+      root_agent=agent,
+  )
+  return node_capture, testing_utils.InMemoryRunner(app=app_instance)
+
+
+def test_get_common_branch_prefix():
+  """Tests _get_common_branch_prefix."""
+  assert join_node._get_common_branch_prefix(['A@1', 'A@2']) == ''
+  assert join_node._get_common_branch_prefix(['A@1.B@1', 'A@1.B@2']) == 'A@1'
+  assert join_node._get_common_branch_prefix(['A@1', 'A@1']) == 'A@1'
+  assert join_node._get_common_branch_prefix(['A@1', '']) == ''
+  assert join_node._get_common_branch_prefix(['', '']) == ''
+  assert join_node._get_common_branch_prefix([]) == ''
+
+
+@pytest.mark.asyncio
+async def test_join_node_waits_for_all_inputs(request: pytest.FixtureRequest):
+  """Tests JoinNode with fan-in."""
+  node_capture, runner = _build_join_node_workflow(request)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  assert node_capture.received_inputs == [{
+      'NodeA': {'a': 1, 'b': 1},
+      'NodeB': {'b': 2},
+  }]
+
+
+@pytest.mark.asyncio
+async def test_join_node_with_none_state(request: pytest.FixtureRequest):
+  """Tests JoinNode with fan-in when node state is None."""
+  node_capture, runner = _build_join_node_workflow(request)
+  # Run once to set state to None
+  await runner.run_async(testing_utils.get_user_content('start'))
+  # Run again to trigger join_node with state=None
+  await runner.run_async(testing_utils.get_user_content('start'))
+
+  assert node_capture.received_inputs == [
+      {'NodeA': {'a': 1, 'b': 1}, 'NodeB': {'b': 2}},
+      {'NodeA': {'a': 1, 'b': 1}, 'NodeB': {'b': 2}},
+  ]
+
+
+@pytest.mark.asyncio
+async def test_join_node_with_none_inputs(request: pytest.FixtureRequest):
+  """Tests JoinNode with fan-in when incoming edges have None output."""
+  node_a = workflow_testing_utils.TestingNode(
+      name='NodeA', output=None, route='NodeJoin'
+  )
+  node_b = workflow_testing_utils.TestingNode(
+      name='NodeB', output=None, route='NodeJoin'
+  )
+  node_join = join_node.JoinNode(name='NodeJoin')
+  node_capture = workflow_testing_utils.InputCapturingNode(name='NodeCapture')
+  agent = workflow.Workflow(
+      name='test_join_node_none_inputs',
+      edges=[
+          workflow_graph.Edge(from_node=base_node.START, to_node=node_a),
+          workflow_graph.Edge(from_node=base_node.START, to_node=node_b),
+          workflow_graph.Edge(from_node=node_a, to_node=node_join),
+          workflow_graph.Edge(from_node=node_b, to_node=node_join),
+          workflow_graph.Edge(from_node=node_join, to_node=node_capture),
+      ],
+  )
+  app_instance = app.App(
+      name=request.function.__name__,
+      root_agent=agent,
+  )
+  runner = testing_utils.InMemoryRunner(app=app_instance)
+
+  await runner.run_async(testing_utils.get_user_content('start'))
+
+  assert node_capture.received_inputs == [
+      {'NodeA': None, 'NodeB': None},
+  ]
+
+
+# ── JoinNode input_schema ──────────────────────────────────────
+# input_schema on JoinNode validates each trigger input individually
+# (each predecessor's output), not the joined dict.
+
+
+class _TriggerInput(BaseModel):
+  key: str
+  value: int
+
+
+@pytest.mark.asyncio
+async def test_join_node_input_schema_validates_per_trigger(
+    request: pytest.FixtureRequest,
+):
+  """JoinNode input_schema validates each trigger input individually."""
+
+  def node_a() -> dict:
+    return {'key': 'a', 'value': 1}
+
+  def node_b() -> dict:
+    return {'key': 'b', 'value': 2}
+
+  join = join_node.JoinNode(name='join', input_schema=_TriggerInput)
+  capture = workflow_testing_utils.InputCapturingNode(name='capture')
+
+  agent = Workflow(
+      name='wf',
+      edges=[
+          (START, node_a),
+          (START, node_b),
+          (node_a, join),
+          (node_b, join),
+          (join, capture),
+      ],
+  )
+  app_instance = app.App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app_instance)
+  await runner.run_async(testing_utils.get_user_content('start'))
+
+  assert capture.received_inputs == [{
+      'node_a': {'key': 'a', 'value': 1},
+      'node_b': {'key': 'b', 'value': 2},
+  }]
+
+
+@pytest.mark.asyncio
+async def test_join_node_input_schema_rejects_invalid_trigger(
+    request: pytest.FixtureRequest,
+):
+  """JoinNode input_schema rejects invalid trigger input early."""
+
+  def node_a() -> dict:
+    return {'key': 'a', 'value': 1}
+
+  def node_b() -> dict:
+    return {'wrong': 'shape'}  # missing required fields
+
+  join = join_node.JoinNode(name='join', input_schema=_TriggerInput)
+  capture = workflow_testing_utils.InputCapturingNode(name='capture')
+
+  agent = Workflow(
+      name='wf',
+      edges=[
+          (START, node_a),
+          (START, node_b),
+          (node_a, join),
+          (node_b, join),
+          (join, capture),
+      ],
+  )
+  app_instance = app.App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app_instance)
+  with pytest.raises(Exception):
+    await runner.run_async(testing_utils.get_user_content('start'))
+
+
+@pytest.mark.asyncio
+async def test_join_node_input_schema_none_trigger_passes(
+    request: pytest.FixtureRequest,
+):
+  """JoinNode input_schema skips validation for None trigger input."""
+  # Given
+  node_a_fn = workflow_testing_utils.TestingNode(
+      name='NodeA', output=None, route='join'
+  )
+  node_b_fn = workflow_testing_utils.TestingNode(
+      name='NodeB', output={'key': 'b', 'value': 2}
+  )
+  join = join_node.JoinNode(name='join', input_schema=_TriggerInput)
+  capture = workflow_testing_utils.InputCapturingNode(name='capture')
+
+  agent = Workflow(
+      name='wf',
+      edges=[
+          (START, node_a_fn),
+          (START, node_b_fn),
+          (node_a_fn, join),
+          (node_b_fn, join),
+          (join, capture),
+      ],
+  )
+  app_instance = app.App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app_instance)
+
+  # When
+  await runner.run_async(testing_utils.get_user_content('start'))
+
+  # Then
+  assert capture.received_inputs == [{
+      'NodeA': None,
+      'NodeB': {'key': 'b', 'value': 2},
+  }]
+
+
+@pytest.mark.asyncio
+async def test_join_node_computes_common_branch_prefix(
+    request: pytest.FixtureRequest,
+):
+  """Tests JoinNode computes common branch prefix for final output."""
+  node_capture, runner = _build_join_node_workflow(request)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  # Find the final output event from JoinNode
+  join_events = [
+      e
+      for e in events
+      if 'NodeJoin' in e.node_info.path and e.output is not None
+  ]
+  assert len(join_events) == 1
+  join_event = join_events[0]
+
+  # NodeA and NodeB run in parallel, so they should have branches like 'NodeA@1' and 'NodeB@1'.
+  a_events = [e for e in events if 'NodeA' in e.node_info.path]
+  b_events = [e for e in events if 'NodeB' in e.node_info.path]
+
+  assert any('NodeA@' in e.branch for e in a_events if e.branch)
+  assert any('NodeB@' in e.branch for e in b_events if e.branch)
+
+  # The common prefix of 'NodeA@1' and 'NodeB@1' is empty string.
+  # So JoinNode should set branch to empty string (which is converted to None).
+  assert join_event.branch is None
+
+  # The node after JoinNode (NodeCapture) should also have branch=None
+  capture_events = [e for e in events if 'NodeCapture' in e.node_info.path]
+  assert len(capture_events) > 0
+  for e in capture_events:
+    assert e.branch is None

tests/unittests/workflow/test_node_runner_ctx.py

@@ -0,0 +1,658 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for NodeRunner → Context as result channel.
+
+Verifies that NodeRunner correctly populates ctx.output, ctx.route,
+and ctx.interrupt_ids from yielded events and direct assignment,
+and that resume state (prior_output, prior_interrupt_ids) is carried
+forward correctly.
+"""
+
+from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
+
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.context import Context
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.events.event import Event
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.sessions.session import Session
+from google.adk.workflow._base_node import BaseNode
+from google.adk.workflow._node_runner import NodeRunner
+from google.genai import types
+import pytest
+
+# --- Helpers ---
+
+
+def _make_ctx(invocation_id='inv-test', enqueue_events=None, node_path=''):
+  """Create a minimal Context mock with IC."""
+  mock_agent = MagicMock(spec=BaseAgent)
+  real_session = Session(
+      id='test_session', app_name='test_app', user_id='test_user'
+  )
+  real_session_service = InMemorySessionService()
+
+  ic = InvocationContext(
+      invocation_id=invocation_id,
+      agent=mock_agent,
+      session=real_session,
+      session_service=real_session_service,
+  )
+
+  collected = enqueue_events if enqueue_events is not None else []
+
+  async def _enqueue(event):
+    collected.append(event)
+
+  object.__setattr__(ic, '_enqueue_event', AsyncMock(side_effect=_enqueue))
+
+  ctx = Context(
+      invocation_context=ic,
+      node_path=node_path,
+  )
+  return ctx, collected
+
+
+# =========================================================================
+# Context as RESULT — fields populated by NodeRunner after execution
+# =========================================================================
+
+
+# --- ctx.output from yielded events ---
+
+
+@pytest.mark.asyncio
+async def test_yield_value_sets_ctx_output():
+  """Yielding a value sets ctx.output on the returned context."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      yield 'hello'
+
+  parent_ctx, _ = _make_ctx()
+  child_ctx = await NodeRunner(
+      node=_Node(name='n'), parent_ctx=parent_ctx
+  ).run()
+
+  assert child_ctx.output == 'hello'
+
+
+@pytest.mark.asyncio
+async def test_yield_event_output_sets_ctx_output():
+  """Yielding Event(output=X) sets ctx.output."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      yield Event(output='from_event')
+
+  parent_ctx, _ = _make_ctx()
+  child_ctx = await NodeRunner(
+      node=_Node(name='n'), parent_ctx=parent_ctx
+  ).run()
+
+  assert child_ctx.output == 'from_event'
+
+
+@pytest.mark.asyncio
+async def test_no_yield_leaves_ctx_output_none():
+  """A node that yields nothing leaves ctx.output as None."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      return
+      yield  # noqa: unreachable
+
+  parent_ctx, _ = _make_ctx()
+  child_ctx = await NodeRunner(
+      node=_Node(name='n'), parent_ctx=parent_ctx
+  ).run()
+
+  assert child_ctx.output is None
+
+
+# --- ctx.output set directly ---
+
+
+@pytest.mark.asyncio
+async def test_ctx_output_set_directly():
+  """Setting ctx.output directly produces a deferred output event."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      ctx.output = 'direct'
+      yield  # noqa: unreachable
+
+  parent_ctx, events = _make_ctx()
+  child_ctx = await NodeRunner(
+      node=_Node(name='n'), parent_ctx=parent_ctx
+  ).run()
+
+  assert child_ctx.output == 'direct'
+  output_events = [e for e in events if e.output is not None]
+  assert len(output_events) == 1
+  assert output_events[0].output == 'direct'
+
+
+@pytest.mark.asyncio
+async def test_ctx_output_direct_with_state_delta():
+  """Deferred output bundles pending state deltas onto the same event."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      ctx.state['key'] = 'val'
+      ctx.output = 'result'
+      yield  # noqa: unreachable
+
+  parent_ctx, events = _make_ctx()
+  child_ctx = await NodeRunner(
+      node=_Node(name='n'), parent_ctx=parent_ctx
+  ).run()
+
+  assert child_ctx.output == 'result'
+  output_events = [e for e in events if e.output is not None]
+  assert len(output_events) == 1
+  assert output_events[0].actions.state_delta['key'] == 'val'
+
+
+@pytest.mark.asyncio
+async def test_deferred_output_emitted_after_intermediate():
+  """ctx.output set directly emits after intermediate content events."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      ctx.output = 'deferred'
+      yield Event(content=types.Content(parts=[types.Part(text='working')]))
+
+  parent_ctx, events = _make_ctx()
+  child_ctx = await NodeRunner(
+      node=_Node(name='n'), parent_ctx=parent_ctx
+  ).run()
+
+  assert child_ctx.output == 'deferred'
+  assert len(events) == 2
+  assert events[0].content.parts[0].text == 'working'
+  assert events[1].output == 'deferred'
+
+
+# --- ctx.output validation ---
+
+
+@pytest.mark.asyncio
+async def test_double_output_raises():
+  """Setting ctx.output twice raises ValueError."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      ctx.output = 'first'
+      ctx.output = 'second'
+      yield  # noqa: unreachable
+
+  parent_ctx, events = _make_ctx()
+  await NodeRunner(node=_Node(name='n'), parent_ctx=parent_ctx).run()
+  error_events = [e for e in events if e.error_code]
+  assert len(error_events) == 1
+  assert error_events[0].error_code == 'ValueError'
+  assert 'already set' in error_events[0].error_message
+
+
+@pytest.mark.asyncio
+async def test_yield_then_ctx_output_raises():
+  """Yielding output then setting ctx.output raises ValueError."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      yield 'first'
+      ctx.output = 'second'
+
+  parent_ctx, events = _make_ctx()
+  await NodeRunner(node=_Node(name='n'), parent_ctx=parent_ctx).run()
+  error_events = [e for e in events if e.error_code]
+  assert len(error_events) == 1
+  assert error_events[0].error_code == 'ValueError'
+  assert 'already set' in error_events[0].error_message
+
+
+# --- ctx.route ---
+
+
+@pytest.mark.asyncio
+async def test_yield_route_sets_ctx_route():
+  """Yielding Event(route=R) sets ctx.route."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      yield Event(output='out', route='next')
+
+  parent_ctx, _ = _make_ctx()
+  child_ctx = await NodeRunner(
+      node=_Node(name='n'), parent_ctx=parent_ctx
+  ).run()
+
+  assert child_ctx.output == 'out'
+  assert child_ctx.route == 'next'
+
+
+@pytest.mark.asyncio
+async def test_ctx_route_set_directly():
+  """Setting ctx.route directly is readable after run."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      ctx.route = 'branch_a'
+      yield 'out'
+
+  parent_ctx, _ = _make_ctx()
+  child_ctx = await NodeRunner(
+      node=_Node(name='n'), parent_ctx=parent_ctx
+  ).run()
+
+  assert child_ctx.route == 'branch_a'
+
+
+# --- ctx.interrupt_ids ---
+
+
+@pytest.mark.asyncio
+async def test_interrupt_sets_ctx_interrupt_ids():
+  """Yielding an interrupt event populates ctx.interrupt_ids."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      yield Event(
+          content=types.Content(
+              parts=[
+                  types.Part(
+                      function_call=types.FunctionCall(
+                          name='tool', args={}, id='fc-1'
+                      )
+                  )
+              ]
+          ),
+          long_running_tool_ids={'fc-1'},
+      )
+
+  parent_ctx, _ = _make_ctx()
+  child_ctx = await NodeRunner(
+      node=_Node(name='n'), parent_ctx=parent_ctx
+  ).run()
+
+  assert child_ctx.interrupt_ids == {'fc-1'}
+  assert child_ctx.output is None
+
+
+@pytest.mark.asyncio
+async def test_output_and_interrupt_coexist():
+  """Output and interrupt can coexist across separate events."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      yield 'result'
+      yield Event(
+          content=types.Content(
+              parts=[
+                  types.Part(
+                      function_call=types.FunctionCall(
+                          name='tool', args={}, id='fc-1'
+                      )
+                  )
+              ]
+          ),
+          long_running_tool_ids={'fc-1'},
+      )
+
+  parent_ctx, _ = _make_ctx()
+  child_ctx = await NodeRunner(
+      node=_Node(name='n'), parent_ctx=parent_ctx
+  ).run()
+
+  assert child_ctx.output == 'result'
+  assert child_ctx.interrupt_ids == {'fc-1'}
+
+
+@pytest.mark.asyncio
+async def test_duplicate_interrupt_ids_deduplicated():
+  """Duplicate interrupt IDs are deduplicated (set semantics)."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      yield Event(long_running_tool_ids={'fc-1', 'fc-2'})
+      yield Event(long_running_tool_ids={'fc-2', 'fc-3'})
+
+  parent_ctx, _ = _make_ctx()
+  child_ctx = await NodeRunner(
+      node=_Node(name='n'), parent_ctx=parent_ctx
+  ).run()
+
+  assert child_ctx.interrupt_ids == {'fc-1', 'fc-2', 'fc-3'}
+
+
+# --- Output delegation (use_as_output) ---
+
+
+@pytest.mark.asyncio
+async def test_delegated_output_not_enqueued():
+  """When output is delegated, the output event is not enqueued."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      ctx._output_delegated = True
+      yield 'delegated_value'
+
+  parent_ctx, events = _make_ctx()
+  child_ctx = await NodeRunner(
+      node=_Node(name='n'), parent_ctx=parent_ctx
+  ).run()
+
+  assert child_ctx.output == 'delegated_value'
+  output_events = [e for e in events if e.output is not None]
+  assert len(output_events) == 0
+
+
+@pytest.mark.asyncio
+async def test_delegated_ctx_output_not_emitted():
+  """When output is delegated and set via ctx.output, no event emitted."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      ctx._output_delegated = True
+      ctx.output = 'delegated_direct'
+      yield  # noqa: unreachable
+
+  parent_ctx, events = _make_ctx()
+  child_ctx = await NodeRunner(
+      node=_Node(name='n'), parent_ctx=parent_ctx
+  ).run()
+
+  assert child_ctx.output == 'delegated_direct'
+  output_events = [e for e in events if e.output is not None]
+  assert len(output_events) == 0
+
+
+@pytest.mark.asyncio
+async def test_delegated_output_preserves_event_details():
+  """When output is delegated, the event is enqueued but output is suppressed."""
+  from google.adk.events.event_actions import EventActions
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      ctx._output_delegated = True
+      yield Event(
+          output='delegated_value',
+          actions=EventActions(state_delta={'foo': 'bar'}),
+          content=types.Content(role='model', parts=[types.Part(text='hello')]),
+      )
+
+  parent_ctx, events = _make_ctx()
+  child_ctx = await NodeRunner(
+      node=_Node(name='n'), parent_ctx=parent_ctx
+  ).run()
+
+  assert child_ctx.output == 'delegated_value'
+  assert len(events) == 1
+  event = events[0]
+  assert event.output is None
+  assert event.actions.state_delta == {'foo': 'bar'}
+  assert event.content.parts[0].text == 'hello'
+
+
+# =========================================================================
+# Context as INPUT — resume state provided to NodeRunner at construction
+# =========================================================================
+
+
+@pytest.mark.asyncio
+async def test_prior_output_carried_forward():
+  """Prior output from a previous run is available on ctx.output."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      return
+      yield  # noqa: unreachable
+
+  parent_ctx, _ = _make_ctx()
+  child_ctx = await NodeRunner(
+      node=_Node(name='n'),
+      parent_ctx=parent_ctx,
+      prior_output='cached_result',
+  ).run()
+
+  assert child_ctx.output == 'cached_result'
+
+
+@pytest.mark.asyncio
+async def test_prior_interrupt_ids_carried_forward():
+  """Prior interrupt IDs from a previous run are on ctx."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      return
+      yield  # noqa: unreachable
+
+  parent_ctx, _ = _make_ctx()
+  child_ctx = await NodeRunner(
+      node=_Node(name='n'),
+      parent_ctx=parent_ctx,
+      prior_interrupt_ids={'fc-old'},
+  ).run()
+
+  assert 'fc-old' in child_ctx.interrupt_ids
+
+
+@pytest.mark.asyncio
+async def test_prior_and_new_interrupt_ids_merged():
+  """New interrupt IDs are merged with prior ones."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      yield Event(
+          content=types.Content(
+              parts=[
+                  types.Part(
+                      function_call=types.FunctionCall(
+                          name='tool', args={}, id='fc-new'
+                      )
+                  )
+              ]
+          ),
+          long_running_tool_ids={'fc-new'},
+      )
+
+  parent_ctx, _ = _make_ctx()
+  child_ctx = await NodeRunner(
+      node=_Node(name='n'),
+      parent_ctx=parent_ctx,
+      prior_interrupt_ids={'fc-old'},
+  ).run()
+
+  assert child_ctx.interrupt_ids == {'fc-old', 'fc-new'}
+
+
+# =========================================================================
+# event_author — parent orchestrator overrides event author
+# =========================================================================
+
+
+@pytest.mark.asyncio
+async def test_event_author_defaults_to_node_name():
+  """Without event_author, events use the node's own name."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      yield 'result'
+
+  parent_ctx, events = _make_ctx()
+  await NodeRunner(node=_Node(name='my_node'), parent_ctx=parent_ctx).run()
+
+  assert events[0].author == 'my_node'
+
+
+@pytest.mark.asyncio
+async def test_event_author_overrides_node_name():
+  """When parent sets event_author, events use that instead."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      yield 'result'
+
+  parent_ctx, events = _make_ctx()
+  parent_ctx.event_author = 'my_workflow'
+  await NodeRunner(node=_Node(name='my_node'), parent_ctx=parent_ctx).run()
+
+  assert events[0].author == 'my_workflow'
+
+
+@pytest.mark.asyncio
+async def test_event_author_overrides_preset_author():
+  """event_author always wins, even over a pre-set event author."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      yield Event(author='custom_author', output='result')
+
+  parent_ctx, events = _make_ctx()
+  parent_ctx.event_author = 'my_workflow'
+  await NodeRunner(node=_Node(name='my_node'), parent_ctx=parent_ctx).run()
+
+  assert events[0].author == 'my_workflow'
+
+
+# =========================================================================
+# Branch propagation tests
+# =========================================================================
+
+
+@pytest.mark.asyncio
+async def test_override_branch_used_in_node_runner():
+  """NodeRunner uses override_branch if provided."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      yield Event(output='result')
+
+  parent_ctx, events = _make_ctx()
+  await NodeRunner(
+      node=_Node(name='n'),
+      parent_ctx=parent_ctx,
+      override_branch='custom_branch',
+  ).run()
+
+  assert events[0].branch == 'custom_branch'
+
+
+@pytest.mark.asyncio
+async def test_use_sub_branch_appends_segment_to_branch():
+  """NodeRunner appends node_name@run_id to branch when use_sub_branch is True."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      yield Event(output='result')
+
+  parent_ctx, events = _make_ctx()
+  parent_ctx._invocation_context.branch = 'parent_branch'
+  await NodeRunner(
+      node=_Node(name='n'),
+      parent_ctx=parent_ctx,
+      use_sub_branch=True,
+      run_id='1',
+  ).run()
+
+  assert events[0].branch == 'parent_branch.n@1'
+
+
+@pytest.mark.asyncio
+async def test_sequential_branch_propagation():
+  """NodeRunner inherits parent branch when use_sub_branch is False."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      yield Event(output='result')
+
+  parent_ctx, events = _make_ctx()
+  parent_ctx._invocation_context.branch = 'parent_branch'
+  await NodeRunner(
+      node=_Node(name='n'),
+      parent_ctx=parent_ctx,
+      use_sub_branch=False,
+  ).run()
+
+  assert events[0].branch == 'parent_branch'
+
+
+@pytest.mark.asyncio
+async def test_child_event_branch_does_not_mutate_parent_ic():
+  """A child node altering its branch does not mutate the parent's shared InvocationContext branch."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      yield Event(output='result', branch='new_child_branch')
+
+  parent_ctx, events = _make_ctx()
+  parent_ctx._invocation_context.branch = 'parent_branch'
+  await NodeRunner(
+      node=_Node(name='n'),
+      parent_ctx=parent_ctx,
+      use_sub_branch=False,
+  ).run()
+
+  assert events[0].branch == 'new_child_branch'
+  # The parent's branch must remain unchanged.
+  assert parent_ctx._invocation_context.branch == 'parent_branch'
+
+
+@pytest.mark.asyncio
+async def test_override_isolation_scope_used_in_node_runner():
+  """NodeRunner sets isolation_scope on child context and enriches emitted events."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(self, *, ctx, node_input):
+      assert ctx.isolation_scope == 'task:fc-999'
+      yield Event(output='result')
+
+  parent_ctx, events = _make_ctx()
+  await NodeRunner(
+      node=_Node(name='n'),
+      parent_ctx=parent_ctx,
+      override_isolation_scope='task:fc-999',
+  ).run()
+
+  assert events[0].isolation_scope == 'task:fc-999'

tests/unittests/workflow/test_node_runner_integration.py

@@ -0,0 +1,539 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for NodeRunner ↔ node integration.
+
+Verifies that NodeRunner correctly drives BaseNode.run(), enriches
+events, flushes state/artifact deltas, and delivers events to the
+session.
+"""
+
+from typing import Any
+from typing import AsyncGenerator
+from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
+
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.context import Context
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.events.event import Event
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.sessions.session import Session
+from google.adk.workflow._base_node import BaseNode
+from google.adk.workflow._node_runner import NodeRunner
+from google.genai import types
+import pytest
+
+# --- Test helper nodes ---
+
+
+class _EchoNode(BaseNode):
+
+  async def _run_impl(
+      self, *, ctx: Context, node_input: Any
+  ) -> AsyncGenerator[Any, None]:
+    yield node_input
+
+
+class _EmptyNode(BaseNode):
+
+  async def _run_impl(
+      self, *, ctx: Context, node_input: Any
+  ) -> AsyncGenerator[Any, None]:
+    return
+    yield
+
+
+class _MultiEventNode(BaseNode):
+
+  async def _run_impl(
+      self, *, ctx: Context, node_input: Any
+  ) -> AsyncGenerator[Any, None]:
+    yield Event(author='step1')
+    yield Event(author='step2')
+    yield Event(author='step3')
+
+
+class _InterruptNode(BaseNode):
+
+  async def _run_impl(
+      self, *, ctx: Context, node_input: Any
+  ) -> AsyncGenerator[Any, None]:
+    yield Event(
+        content=types.Content(
+            parts=[
+                types.Part(
+                    function_call=types.FunctionCall(
+                        name='long_tool', args={}, id='fc-1'
+                    )
+                )
+            ]
+        ),
+        long_running_tool_ids={'fc-1'},
+    )
+
+
+class _InterruptThenMoreNode(BaseNode):
+
+  async def _run_impl(
+      self, *, ctx: Context, node_input: Any
+  ) -> AsyncGenerator[Any, None]:
+    yield Event(
+        content=types.Content(
+            parts=[
+                types.Part(
+                    function_call=types.FunctionCall(
+                        name='long_tool', args={}, id='fc-2'
+                    )
+                )
+            ]
+        ),
+        long_running_tool_ids={'fc-2'},
+    )
+    yield Event(author='after_interrupt_1')
+    yield Event(author='after_interrupt_2')
+
+
+class _ErrorNode(BaseNode):
+
+  async def _run_impl(
+      self, *, ctx: Context, node_input: Any
+  ) -> AsyncGenerator[Any, None]:
+    raise RuntimeError('node failure')
+    yield  # pylint: disable=unreachable
+
+
+class _OutputWithRouteNode(BaseNode):
+
+  async def _run_impl(
+      self, *, ctx: Context, node_input: Any
+  ) -> AsyncGenerator[Any, None]:
+    yield Event(output='routed_output', route='next')
+
+
+class _StateMutatingNode(BaseNode):
+
+  async def _run_impl(
+      self, *, ctx: Context, node_input: Any
+  ) -> AsyncGenerator[Any, None]:
+    ctx.state['key1'] = 'value1'
+    ctx.state['key2'] = 42
+    yield 'done'
+
+
+class _ResumeInputReadingNode(BaseNode):
+  captured: list[Any] = []
+
+  async def _run_impl(
+      self, *, ctx: Context, node_input: Any
+  ) -> AsyncGenerator[Any, None]:
+    self.captured.append(ctx.resume_inputs)
+    yield 'resumed'
+
+
+class _ArtifactSavingNode(BaseNode):
+
+  async def _run_impl(
+      self, *, ctx: Context, node_input: Any
+  ) -> AsyncGenerator[Any, None]:
+    ctx.actions.artifact_delta['doc.txt'] = 1
+    yield 'saved'
+
+
+# --- Helpers ---
+
+
+def _make_ctx(invocation_id='inv-test', enqueue_events=None, node_path=''):
+  """Create a minimal Context mock with IC."""
+
+  mock_agent = MagicMock(spec=BaseAgent)
+  real_session = Session(
+      id='test_session', app_name='test_app', user_id='test_user'
+  )
+  real_session_service = InMemorySessionService()
+
+  ic = InvocationContext(
+      invocation_id=invocation_id,
+      agent=mock_agent,
+      session=real_session,
+      session_service=real_session_service,
+  )
+
+  collected = enqueue_events if enqueue_events is not None else []
+
+  async def _enqueue(event):
+    collected.append(event)
+
+  object.__setattr__(ic, '_enqueue_event', AsyncMock(side_effect=_enqueue))
+
+  ctx = Context(
+      invocation_context=ic,
+      node_path=node_path,
+  )
+  return ctx, collected
+
+
+# --- Tests ---
+
+
+@pytest.mark.asyncio
+async def test_node_output_returned_in_result():
+  """Running a node that produces output returns it in the result."""
+  node = _EchoNode(name='echo')
+  ctx, _ = _make_ctx()
+  result = await NodeRunner(node=node, parent_ctx=ctx).run(node_input='hello')
+  assert result.output == 'hello'
+  assert result.interrupt_ids == set()
+
+
+@pytest.mark.asyncio
+async def test_no_output_returns_none():
+  """Running a node that produces no output returns None."""
+  node = _EmptyNode(name='empty')
+  ctx, _ = _make_ctx()
+  result = await NodeRunner(node=node, parent_ctx=ctx).run()
+  assert result.output is None
+  assert result.interrupt_ids == set()
+
+
+@pytest.mark.asyncio
+async def test_event_author_is_node_name():
+  """Events are authored by the node's name."""
+  node = _EchoNode(name='my_node')
+  ctx, events = _make_ctx()
+  await NodeRunner(node=node, parent_ctx=ctx).run(node_input='data')
+
+  output_events = [e for e in events if e.output is not None]
+  assert output_events[0].author == 'my_node'
+
+
+@pytest.mark.asyncio
+async def test_event_path_contains_node_name():
+  """Event node_info.path includes the node name and execution context."""
+  node = _EchoNode(name='path_test')
+  ctx, events = _make_ctx(invocation_id='inv-123')
+  runner = NodeRunner(node=node, parent_ctx=ctx, run_id='exec-456')
+  await runner.run(node_input='data')
+
+  output_events = [e for e in events if e.output is not None]
+  event = output_events[0]
+  assert event.node_info.path == 'path_test@exec-456'
+  assert event.invocation_id == 'inv-123'
+
+
+@pytest.mark.asyncio
+async def test_interrupt_captured_in_result():
+  """A node that signals an interrupt reports it in the result."""
+  node = _InterruptNode(name='interrupt_node')
+  ctx, _ = _make_ctx()
+  result = await NodeRunner(node=node, parent_ctx=ctx).run()
+  assert 'fc-1' in result.interrupt_ids
+
+
+@pytest.mark.asyncio
+async def test_node_continues_after_interrupt():
+  """A node that interrupts can still produce more events before finishing."""
+  node = _InterruptThenMoreNode(name='flag_finish')
+  ctx, events = _make_ctx()
+  result = await NodeRunner(node=node, parent_ctx=ctx).run()
+  assert 'fc-2' in result.interrupt_ids
+  assert len(events) >= 3
+
+
+@pytest.mark.asyncio
+async def test_state_mutations_emitted_as_delta():
+  """State changes made by a node are delivered as a separate event."""
+  node = _StateMutatingNode(name='state_node')
+  ctx, events = _make_ctx()
+  await NodeRunner(node=node, parent_ctx=ctx).run()
+
+  all_deltas = {}
+  for e in events:
+    if e.actions.state_delta:
+      all_deltas.update(e.actions.state_delta)
+  assert all_deltas.get('key1') == 'value1'
+  assert all_deltas.get('key2') == 42
+
+
+@pytest.mark.asyncio
+async def test_artifact_delta_emitted():
+  """Artifact saves made by a node are delivered as a delta event."""
+  node = _ArtifactSavingNode(name='artifact_node')
+  ctx, events = _make_ctx()
+  await NodeRunner(node=node, parent_ctx=ctx).run()
+
+  artifact_deltas = {}
+  for e in events:
+    if e.actions.artifact_delta:
+      artifact_deltas.update(e.actions.artifact_delta)
+  assert 'doc.txt' in artifact_deltas
+
+
+@pytest.mark.asyncio
+async def test_events_enqueued_in_yield_order():
+  """Multiple events from a node arrive in the order they were produced."""
+  node = _MultiEventNode(name='multi')
+  ctx, events = _make_ctx()
+  await NodeRunner(node=node, parent_ctx=ctx).run()
+
+  # All 3 events enqueued, authored by node name (framework overrides).
+  assert len(events) == 3
+  assert all(e.author == 'multi' for e in events)
+
+
+@pytest.mark.asyncio
+async def test_node_exception_propagates():
+  """A node that raises an error surfaces it to the caller."""
+  node = _ErrorNode(name='error_node')
+  ctx, events = _make_ctx()
+  child_ctx = await NodeRunner(node=node, parent_ctx=ctx).run()
+
+  # Verify error is recorded on returned context
+  assert child_ctx.error is not None
+  assert isinstance(child_ctx.error, RuntimeError)
+  assert str(child_ctx.error) == 'node failure'
+
+  # Verify error event was enqueued
+  error_events = [e for e in events if e.error_code]
+  assert len(error_events) == 1
+  assert error_events[0].error_code == 'RuntimeError'
+  assert 'node failure' in error_events[0].error_message
+
+
+@pytest.mark.asyncio
+async def test_resume_inputs_available_on_context():
+  """Resume inputs are accessible to the node during execution."""
+  node = _ResumeInputReadingNode(name='resume_node')
+  node.captured = []
+  ctx, _ = _make_ctx()
+  resume = {'int-1': 'user_response'}
+  await NodeRunner(node=node, parent_ctx=ctx).run(resume_inputs=resume)
+  assert node.captured[0] == resume
+
+
+@pytest.mark.asyncio
+async def test_node_path_includes_parent():
+  """A child node's node_path is parent_node_path/child_name."""
+  node = _EchoNode(name='child')
+  ctx, events = _make_ctx(node_path='parent_path')
+  runner = NodeRunner(node=node, parent_ctx=ctx)
+  await runner.run(node_input='x')
+
+  output_events = [e for e in events if e.output is not None]
+  assert output_events[0].node_info.path == 'parent_path/child@1'
+
+
+@pytest.mark.asyncio
+async def test_run_id_generated_when_omitted():
+  """Each node run gets a unique execution ID by default."""
+  node = _EchoNode(name='auto_id')
+  ctx, _ = _make_ctx()
+
+  runner = NodeRunner(node=node, parent_ctx=ctx)
+
+  assert runner.run_id
+  assert isinstance(runner.run_id, str)
+
+
+@pytest.mark.asyncio
+async def test_explicit_run_id_used():
+  """A caller-specified execution ID is used on the runner and events."""
+  node = _EchoNode(name='explicit_id')
+  ctx, events = _make_ctx()
+
+  runner = NodeRunner(node=node, parent_ctx=ctx, run_id='my-exec-id')
+
+  assert runner.run_id == 'my-exec-id'
+  await runner.run(node_input='data')
+  assert events[0].node_info.path == 'explicit_id@my-exec-id'
+
+
+@pytest.mark.asyncio
+async def test_route_captured_in_result():
+  """A node's routing decision is available in the result."""
+  node = _OutputWithRouteNode(name='route_node')
+  ctx, _ = _make_ctx()
+  result = await NodeRunner(node=node, parent_ctx=ctx).run()
+  assert result.output == 'routed_output'
+  assert result.route == 'next'
+
+
+@pytest.mark.asyncio
+async def test_preset_author_overridden_by_framework():
+  """Framework always sets author — preset author is overridden."""
+  node = _MultiEventNode(name='multi')
+  ctx, events = _make_ctx()
+  await NodeRunner(node=node, parent_ctx=ctx).run()
+
+  # All events get node name, not the preset 'step1'/'step2'/'step3'.
+  assert all(e.author == 'multi' for e in events)
+
+
+class _MultiOutputNode(BaseNode):
+
+  async def _run_impl(
+      self, *, ctx: Context, node_input: Any
+  ) -> AsyncGenerator[Any, None]:
+    yield Event(output='first')
+    yield Event(output='second')
+
+
+@pytest.mark.asyncio
+async def test_multiple_outputs_raises():
+  """A node that produces more than one output is rejected."""
+  node = _MultiOutputNode(name='multi_out')
+  ctx, events = _make_ctx()
+  await NodeRunner(node=node, parent_ctx=ctx).run()
+  error_events = [e for e in events if e.error_code]
+  assert len(error_events) == 1
+  assert error_events[0].error_code == 'ValueError'
+  assert 'at most one output' in error_events[0].error_message
+
+
+@pytest.mark.asyncio
+async def test_all_events_delivered():
+  """All events from a node are delivered to the session."""
+  node = _EchoNode(name='enqueue_test')
+  ctx, events = _make_ctx()
+  await NodeRunner(node=node, parent_ctx=ctx).run(node_input='data')
+  assert len(events) >= 1
+
+
+# --- Delta flushing tests ---
+
+
+@pytest.mark.asyncio
+async def test_state_delta_bundled_with_output_event():
+  """State deltas set before yield are flushed onto the output event."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Any, None]:
+      ctx.state['color'] = 'blue'
+      ctx.state['count'] = 7
+      yield 'result'
+
+  ctx, events = _make_ctx()
+
+  await NodeRunner(node=_Node(name='bundled'), parent_ctx=ctx).run()
+
+  assert len(events) == 1
+  assert events[0].output == 'result'
+  assert events[0].actions.state_delta.get('color') == 'blue'
+  assert events[0].actions.state_delta.get('count') == 7
+
+
+@pytest.mark.asyncio
+async def test_state_after_last_yield_emitted_separately():
+  """State set after the last yield is emitted as a separate event."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Any, None]:
+      yield 'early'
+      ctx.state['late_key'] = 'late_value'
+
+  ctx, events = _make_ctx()
+
+  await NodeRunner(node=_Node(name='late_state'), parent_ctx=ctx).run()
+
+  assert events[0].output == 'early'
+  assert events[1].actions.state_delta.get('late_key') == 'late_value'
+
+
+@pytest.mark.asyncio
+async def test_deltas_skip_partial_events():
+  """Partial events carry no deltas — deltas flush to next non-partial."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Any, None]:
+      ctx.state['before_partial'] = True
+      yield Event(
+          content=types.Content(parts=[types.Part(text='streaming...')]),
+          partial=True,
+      )
+      ctx.state['after_partial'] = True
+      yield 'final'
+
+  ctx, events = _make_ctx()
+
+  await NodeRunner(node=_Node(name='partial_skip'), parent_ctx=ctx).run()
+
+  assert events[0].partial is True
+  assert not events[0].actions or not events[0].actions.state_delta
+  assert events[1].output == 'final'
+  assert events[1].actions.state_delta.get('before_partial') is True
+  assert events[1].actions.state_delta.get('after_partial') is True
+
+
+@pytest.mark.asyncio
+async def test_artifact_and_state_bundled_together():
+  """Both state and artifact deltas are flushed onto the same event."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Any, None]:
+      ctx.state['s1'] = 'v1'
+      ctx.actions.artifact_delta['file.txt'] = 1
+      yield 'done'
+
+  ctx, events = _make_ctx()
+
+  await NodeRunner(node=_Node(name='both_deltas'), parent_ctx=ctx).run()
+
+  assert len(events) == 1
+  assert events[0].output == 'done'
+  assert events[0].actions.state_delta.get('s1') == 'v1'
+  assert events[0].actions.artifact_delta.get('file.txt') == 1
+
+
+@pytest.mark.asyncio
+async def test_node_input_schema_validation():
+  """NodeRunner fails if node input does not match input_schema."""
+  from pydantic import BaseModel
+  from pydantic import ValidationError
+
+  class _MyInput(BaseModel):
+    name: str
+    age: int
+
+  node = _EchoNode(name='schema_node', input_schema=_MyInput)
+  ctx, _ = _make_ctx()
+
+  # Valid input (dict that matches schema)
+  result = await NodeRunner(node=node, parent_ctx=ctx).run(
+      node_input={'name': 'Alice', 'age': 30}
+  )
+
+  # _validate_schema converts model instances to dicts!
+  assert isinstance(result.output, dict)
+  assert result.output['name'] == 'Alice'
+  assert result.output['age'] == 30
+
+  # Invalid input (missing field)
+  result = await NodeRunner(node=node, parent_ctx=ctx).run(
+      node_input={'name': 'Alice'}
+  )
+
+  assert result.error is not None
+  assert isinstance(result.error, ValidationError)

tests/unittests/workflow/test_state_schema.py

@@ -0,0 +1,424 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for Workflow state_schema runtime enforcement."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from google.adk.agents.context import Context
+from google.adk.apps.app import App
+from google.adk.events.event import Event
+from google.adk.sessions.state import State
+from google.adk.sessions.state import StateSchemaError
+from google.adk.workflow import FunctionNode
+from google.adk.workflow import START
+from google.adk.workflow._workflow import Workflow
+from pydantic import BaseModel
+import pytest
+
+from .. import testing_utils
+from .workflow_testing_utils import create_parent_invocation_context
+
+# ── Schema models for testing ────────────────────────────────────────
+
+
+class _PipelineSchema(BaseModel):
+  counter: int
+  name: str
+  optional_field: Optional[str] = None
+
+
+class _NodeSchema(BaseModel):
+  x: int
+  y: int
+
+
+# ── Unit tests: State validation ─────────────────────────────────────
+
+
+def test_state_rejects_unknown_key() -> None:
+  """State with schema raises on unknown key."""
+  state = State(value={}, delta={}, schema=_PipelineSchema)
+  with pytest.raises(StateSchemaError, match='bad_key'):
+    state['bad_key'] = 'value'
+
+
+def test_state_accepts_declared_key() -> None:
+  """State with schema accepts keys that exist in the schema."""
+  state = State(value={}, delta={}, schema=_PipelineSchema)
+  state['counter'] = 5
+  state['name'] = 'hello'
+  assert state['counter'] == 5
+  assert state['name'] == 'hello'
+
+
+def test_state_rejects_wrong_type() -> None:
+  """State with schema raises when value type doesn't match annotation."""
+  state = State(value={}, delta={}, schema=_PipelineSchema)
+  with pytest.raises(StateSchemaError, match='counter'):
+    state['counter'] = 'not_an_int'
+
+
+def test_state_accepts_optional_none() -> None:
+  """Optional fields accept None."""
+  state = State(value={}, delta={}, schema=_PipelineSchema)
+  state['optional_field'] = None
+  assert state['optional_field'] is None
+
+
+def test_state_allows_prefixed_keys() -> None:
+  """Prefixed keys (app:, user:, temp:) bypass schema validation."""
+  state = State(value={}, delta={}, schema=_PipelineSchema)
+  state['app:anything'] = 'value'
+  state['user:pref'] = 42
+  state['temp:cache'] = [1, 2, 3]
+  assert state['app:anything'] == 'value'
+
+
+def test_state_update_validates_all_keys() -> None:
+  """State.update validates each key-value pair."""
+  state = State(value={}, delta={}, schema=_PipelineSchema)
+  with pytest.raises(StateSchemaError, match='unknown'):
+    state.update({'counter': 1, 'unknown': 'x'})
+
+
+def test_state_no_schema_allows_all() -> None:
+  """Without schema, any key/value is accepted (backward compat)."""
+  state = State(value={}, delta={})
+  state['anything'] = 'goes'
+  state['whatever'] = 42
+  assert state['anything'] == 'goes'
+
+
+# ── Startup validation tests ─────────────────────────────────────────
+
+
+def test_startup_rejects_mismatched_param() -> None:
+  """FunctionNode param not in state_schema raises at construction."""
+
+  def node_with_bad_param(ctx: Context, unknown_param: str) -> str:
+    return 'done'
+
+  with pytest.raises(StateSchemaError, match='unknown_param'):
+    Workflow(
+        name='wf',
+        edges=[(START, node_with_bad_param)],
+        state_schema=_PipelineSchema,
+    )
+
+
+def test_startup_accepts_matching_params() -> None:
+  """FunctionNode params matching schema fields pass construction."""
+
+  def node_with_good_params(ctx: Context, counter: int, name: str) -> str:
+    return 'done'
+
+  wf = Workflow(
+      name='wf',
+      edges=[(START, node_with_good_params)],
+      state_schema=_PipelineSchema,
+  )
+  assert wf.state_schema is _PipelineSchema
+
+
+def test_startup_skips_ctx_and_node_input() -> None:
+  """Framework params (ctx, node_input) are not checked against schema."""
+
+  def node_with_framework_params(ctx: Context, node_input: str) -> str:
+    return 'done'
+
+  Workflow(
+      name='wf',
+      edges=[(START, node_with_framework_params)],
+      state_schema=_PipelineSchema,
+  )
+
+
+def test_startup_no_validation_when_schema_none() -> None:
+  """No startup validation when state_schema is not set."""
+
+  def node_with_any_param(ctx: Context, anything: str) -> str:
+    return 'done'
+
+  Workflow(
+      name='wf',
+      edges=[(START, node_with_any_param)],
+  )
+
+
+def test_workflow_state_schema_field_exists() -> None:
+  """Workflow accepts a state_schema parameter."""
+
+  def produce_done():
+    return Event(output='done')
+
+  wf = Workflow(
+      name='wf',
+      edges=[(START, produce_done)],
+      state_schema=_PipelineSchema,
+  )
+  assert wf.state_schema is _PipelineSchema
+
+
+def test_workflow_state_schema_defaults_to_none() -> None:
+  """state_schema defaults to None when not provided."""
+
+  def produce_done():
+    return Event(output='done')
+
+  wf = Workflow(
+      name='wf',
+      edges=[(START, produce_done)],
+  )
+  assert wf.state_schema is None
+
+
+# ── Runtime enforcement tests (workflow execution) ───────────────────
+
+
+@pytest.mark.asyncio
+async def test_workflow_valid_state_writes_succeed(
+    request: pytest.FixtureRequest,
+) -> None:
+  """A workflow with valid state writes runs without errors."""
+
+  def write_state(ctx: Context) -> str:
+    ctx.state['counter'] = 5
+    ctx.state['name'] = 'hello'
+    return 'done'
+
+  wf = Workflow(
+      name='wf',
+      edges=[(START, write_state)],
+      state_schema=_PipelineSchema,
+  )
+  app = App(name=request.function.__name__, root_agent=wf)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+  data_events = [e for e in events if isinstance(e, Event) and e.output]
+  assert any(e.output == 'done' for e in data_events)
+
+
+@pytest.mark.asyncio
+async def test_workflow_rejects_unknown_key_via_ctx_state(
+    request: pytest.FixtureRequest,
+) -> None:
+  """ctx.state write with unknown key raises StateSchemaError."""
+
+  def write_bad_key(ctx: Context) -> str:
+    ctx.state['unknown_key'] = 'value'
+    return 'done'
+
+  wf = Workflow(
+      name='wf',
+      edges=[(START, write_bad_key)],
+      state_schema=_PipelineSchema,
+  )
+  app = App(name=request.function.__name__, root_agent=wf)
+  runner = testing_utils.InMemoryRunner(app=app)
+  with pytest.raises(StateSchemaError, match='unknown_key'):
+    await runner.run_async(testing_utils.get_user_content('start'))
+
+
+@pytest.mark.asyncio
+async def test_workflow_rejects_unknown_key_via_event_state(
+    request: pytest.FixtureRequest,
+) -> None:
+  """Event(state={...}) with unknown key raises StateSchemaError."""
+
+  def emit_bad_state() -> Event:
+    return Event(state={'arbitrary_key': 'value'}, output='done')
+
+  wf = Workflow(
+      name='wf',
+      edges=[(START, emit_bad_state)],
+      state_schema=_PipelineSchema,
+  )
+  app = App(name=request.function.__name__, root_agent=wf)
+  runner = testing_utils.InMemoryRunner(app=app)
+  with pytest.raises(StateSchemaError, match='arbitrary_key'):
+    await runner.run_async(testing_utils.get_user_content('start'))
+
+
+@pytest.mark.asyncio
+async def test_workflow_accepts_valid_event_state(
+    request: pytest.FixtureRequest,
+) -> None:
+  """Event(state={...}) with valid keys succeeds."""
+
+  def emit_valid_state() -> Event:
+    return Event(state={'counter': 10, 'name': 'ok'}, output='done')
+
+  wf = Workflow(
+      name='wf',
+      edges=[(START, emit_valid_state)],
+      state_schema=_PipelineSchema,
+  )
+  app = App(name=request.function.__name__, root_agent=wf)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+  data_events = [e for e in events if isinstance(e, Event) and e.output]
+  assert any(e.output == 'done' for e in data_events)
+
+
+@pytest.mark.asyncio
+async def test_workflow_allows_prefixed_keys_at_runtime(
+    request: pytest.FixtureRequest,
+) -> None:
+  """Prefixed keys bypass schema validation during workflow execution."""
+
+  def write_prefixed(ctx: Context) -> str:
+    ctx.state['temp:debug'] = True
+    ctx.state['app:config'] = 'val'
+    return 'done'
+
+  wf = Workflow(
+      name='wf',
+      edges=[(START, write_prefixed)],
+      state_schema=_PipelineSchema,
+  )
+  app = App(name=request.function.__name__, root_agent=wf)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+  data_events = [e for e in events if isinstance(e, Event) and e.output]
+  assert any(e.output == 'done' for e in data_events)
+
+
+@pytest.mark.asyncio
+async def test_workflow_without_schema_allows_anything(
+    request: pytest.FixtureRequest,
+) -> None:
+  """When state_schema=None, any key/value is accepted (backward compat)."""
+
+  def write_anything(ctx: Context) -> str:
+    ctx.state['any_key'] = 'any_value'
+    ctx.state['another'] = 42
+    return 'done'
+
+  wf = Workflow(
+      name='wf',
+      edges=[(START, write_anything)],
+  )
+  app = App(name=request.function.__name__, root_agent=wf)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+  data_events = [e for e in events if isinstance(e, Event) and e.output]
+  assert any(e.output == 'done' for e in data_events)
+
+
+# ── Per-node state_schema tests ────────────────────────────────────
+
+
+@pytest.mark.asyncio
+async def test_node_level_schema_validates_writes(
+    request: pytest.FixtureRequest,
+) -> None:
+  """A FunctionNode with its own state_schema validates state writes."""
+
+  def write_bad_key(ctx: Context) -> str:
+    ctx.state['bad_key'] = 'value'
+    return 'done'
+
+  node = FunctionNode(
+      name='guarded',
+      func=write_bad_key,
+      state_schema=_NodeSchema,
+  )
+  wf = Workflow(
+      name='wf',
+      edges=[(START, node)],
+  )
+  app = App(name=request.function.__name__, root_agent=wf)
+  runner = testing_utils.InMemoryRunner(app=app)
+  with pytest.raises(StateSchemaError, match='bad_key'):
+    await runner.run_async(testing_utils.get_user_content('start'))
+
+
+@pytest.mark.asyncio
+async def test_node_level_schema_accepts_valid_writes(
+    request: pytest.FixtureRequest,
+) -> None:
+  """A FunctionNode with its own state_schema accepts declared keys."""
+
+  def write_good_keys(ctx: Context) -> str:
+    ctx.state['x'] = 1
+    ctx.state['y'] = 2
+    return 'done'
+
+  node = FunctionNode(
+      name='guarded',
+      func=write_good_keys,
+      state_schema=_NodeSchema,
+  )
+  wf = Workflow(
+      name='wf',
+      edges=[(START, node)],
+  )
+  app = App(name=request.function.__name__, root_agent=wf)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+  data_events = [e for e in events if isinstance(e, Event) and e.output]
+  assert any(e.output == 'done' for e in data_events)
+
+
+@pytest.mark.asyncio
+async def test_node_schema_overrides_workflow_schema(
+    request: pytest.FixtureRequest,
+) -> None:
+  """Node-level state_schema takes precedence over workflow-level schema."""
+
+  def write_node_key(ctx: Context) -> str:
+    ctx.state['x'] = 10
+    return 'done'
+
+  node = FunctionNode(
+      name='guarded',
+      func=write_node_key,
+      state_schema=_NodeSchema,
+  )
+  wf = Workflow(
+      name='wf',
+      edges=[(START, node)],
+      state_schema=_PipelineSchema,
+  )
+  app = App(name=request.function.__name__, root_agent=wf)
+  runner = testing_utils.InMemoryRunner(app=app)
+  # 'x' is in _NodeSchema but NOT in _PipelineSchema — should succeed
+  # because node schema overrides workflow schema
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+  data_events = [e for e in events if isinstance(e, Event) and e.output]
+  assert any(e.output == 'done' for e in data_events)
+
+
+@pytest.mark.asyncio
+async def test_node_without_schema_inherits_workflow_schema(
+    request: pytest.FixtureRequest,
+) -> None:
+  """Node without state_schema inherits validation from parent workflow."""
+
+  def write_bad_key(ctx: Context) -> str:
+    ctx.state['unknown'] = 'value'
+    return 'done'
+
+  wf = Workflow(
+      name='wf',
+      edges=[(START, write_bad_key)],
+      state_schema=_PipelineSchema,
+  )
+  app = App(name=request.function.__name__, root_agent=wf)
+  runner = testing_utils.InMemoryRunner(app=app)
+  with pytest.raises(StateSchemaError, match='unknown'):
+    await runner.run_async(testing_utils.get_user_content('start'))

tests/unittests/workflow/test_task_api_e2e.py

@@ -0,0 +1,517 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""End-to-end tests for the Task Delegation API matrix.
+
+Covers the complete cross-product of dispatch-shape × hierarchy-depth so
+the chat-coordinator wrapper, the workflow-node task path, and the
+nested-delegation path are all exercised:
+
+* LlmAgent root → single task sub-agent (basic FC delegation).
+* LlmAgent root → multiple task sub-agents (sequential delegation).
+* LlmAgent root → task sub-agent → nested task sub-agent (chained).
+* Workflow with a task-mode node (no FC delegation).
+* Workflow with a task-mode node that itself has a task sub-agent.
+* Dynamic node case (task agent dispatched via ``ctx.run_node``).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+from typing import AsyncGenerator
+
+from google.adk.agents.context import Context
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.apps.app import App
+from google.adk.events.event import Event
+from google.adk.workflow import node
+from google.adk.workflow import START
+from google.adk.workflow._base_node import BaseNode
+from google.adk.workflow._workflow import Workflow
+from google.genai import types
+from pydantic import BaseModel
+import pytest
+
+from tests.unittests import testing_utils
+
+# ---------------------------------------------------------------------------
+# Fixture helpers
+# ---------------------------------------------------------------------------
+
+
+def _delegate_part(target_name: str, request_text: str) -> types.Part:
+  """LLM response calling a task sub-agent (the _TaskAgentTool FC)."""
+  return types.Part.from_function_call(
+      name=target_name, args={'request': request_text}
+  )
+
+
+def _finish_part(args: dict[str, Any]) -> types.Part:
+  """LLM response calling finish_task with the given args."""
+  return types.Part.from_function_call(name='finish_task', args=args)
+
+
+def _text_part(text: str) -> types.Part:
+  return types.Part.from_text(text=text)
+
+
+def _make_task_agent(
+    name: str,
+    responses: list,
+    *,
+    sub_agents: list[LlmAgent] | None = None,
+) -> LlmAgent:
+  return LlmAgent(
+      name=name,
+      model=testing_utils.MockModel.create(responses=responses),
+      mode='task',
+      sub_agents=sub_agents or [],
+  )
+
+
+def _collect_finish_outputs(events: list[Event]) -> list[Any]:
+  """Pull out finish_task FC arg dicts in chronological order."""
+  out = []
+  for e in events:
+    for fc in e.get_function_calls():
+      if fc.name == 'finish_task':
+        out.append(dict(fc.args or {}))
+  return out
+
+
+def _get_text_responses(events: list[Event]) -> list[str]:
+  """Concatenate text responses from all model events."""
+  texts = []
+  for e in events:
+    if not e.content or not e.content.parts:
+      continue
+    for p in e.content.parts:
+      if p.text and not p.thought:
+        texts.append(p.text)
+  return texts
+
+
+# ---------------------------------------------------------------------------
+# 1. LlmAgent root → single task sub-agent
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_chat_root_with_single_task_sub_agent(
+    request: pytest.FixtureRequest,
+):
+  """Chat coordinator delegates to one task sub-agent and reports its output."""
+  child = _make_task_agent(
+      name='child',
+      responses=[_finish_part({'result': 'child output'})],
+  )
+
+  root = LlmAgent(
+      name='root',
+      model=testing_utils.MockModel.create(
+          responses=[
+              _delegate_part('child', 'do the thing'),
+              'All done: child output.',
+          ]
+      ),
+      sub_agents=[child],
+  )
+
+  app = App(name=request.function.__name__, root_agent=root)
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  events = await runner.run_async(testing_utils.get_user_content('hi'))
+
+  finish_args = _collect_finish_outputs(events)
+  assert finish_args == [{'result': 'child output'}]
+  assert any(
+      'All done: child output.' in t for t in _get_text_responses(events)
+  )
+
+
+# ---------------------------------------------------------------------------
+# 2. LlmAgent root → multiple task sub-agents (sequential)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_chat_root_with_two_task_sub_agents_sequential(
+    request: pytest.FixtureRequest,
+):
+  """Chat coordinator delegates to two task sub-agents in one turn."""
+  collector = _make_task_agent(
+      name='collector',
+      responses=[_finish_part({'result': 'collected'})],
+  )
+  payer = _make_task_agent(
+      name='payer',
+      responses=[_finish_part({'result': 'paid'})],
+  )
+
+  root = LlmAgent(
+      name='root',
+      model=testing_utils.MockModel.create(
+          responses=[
+              _delegate_part('collector', 'collect'),
+              _delegate_part('payer', 'pay'),
+              'Order placed.',
+          ]
+      ),
+      sub_agents=[collector, payer],
+  )
+
+  app = App(name=request.function.__name__, root_agent=root)
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  events = await runner.run_async(testing_utils.get_user_content('place order'))
+
+  finish_args = _collect_finish_outputs(events)
+  assert finish_args == [{'result': 'collected'}, {'result': 'paid'}]
+  assert any('Order placed.' in t for t in _get_text_responses(events))
+
+
+# ---------------------------------------------------------------------------
+# 3. LlmAgent root → task sub-agent → nested task sub-agent
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.xfail(
+    reason=(
+        'Task-mode wrapper does not dispatch task-delegation FCs (only the '
+        'chat-mode wrapper does), so a task-mode middle agent cannot delegate '
+        'to its task sub-agent.  Documented limitation.'
+    ),
+    strict=True,
+)
+@pytest.mark.asyncio
+async def test_chat_root_with_nested_task_delegation(
+    request: pytest.FixtureRequest,
+):
+  """Task agent itself has a task sub-agent and delegates further."""
+  grandchild = _make_task_agent(
+      name='grandchild',
+      responses=[_finish_part({'result': 'leaf'})],
+  )
+
+  child = LlmAgent(
+      name='child',
+      model=testing_utils.MockModel.create(
+          responses=[
+              _delegate_part('grandchild', 'leaf work'),
+              _finish_part({'result': 'middle wraps leaf'}),
+          ]
+      ),
+      mode='task',
+      sub_agents=[grandchild],
+  )
+
+  root = LlmAgent(
+      name='root',
+      model=testing_utils.MockModel.create(
+          responses=[
+              _delegate_part('child', 'do the thing'),
+              'Top-level done.',
+          ]
+      ),
+      sub_agents=[child],
+  )
+
+  app = App(name=request.function.__name__, root_agent=root)
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  events = await runner.run_async(testing_utils.get_user_content('hi'))
+
+  finish_args = _collect_finish_outputs(events)
+  # grandchild fires first (deepest), then child.
+  assert finish_args == [
+      {'result': 'leaf'},
+      {'result': 'middle wraps leaf'},
+  ]
+  assert any('Top-level done.' in t for t in _get_text_responses(events))
+
+
+# ---------------------------------------------------------------------------
+# 4. Workflow with a single task-mode node (no FC delegation)
+# ---------------------------------------------------------------------------
+
+
+class _CaptureNode(BaseNode):
+  """Records its node_input for assertion."""
+
+  received: list[Any] = []
+
+  async def _run_impl(self, *, ctx, node_input):
+    type(self).received.append(node_input)
+    yield Event(output=node_input)
+
+
+@pytest.mark.asyncio
+async def test_workflow_accepts_task_mode_graph_node():
+  """A mode='task' LlmAgent can be used as a static workflow graph node."""
+  intake = _make_task_agent(name='intake', responses=[])
+  capture = _CaptureNode(name='capture')
+
+  wf = Workflow(name='wf', edges=[(START, intake), (intake, capture)])
+  assert wf is not None
+
+
+# ---------------------------------------------------------------------------
+# 6. Dynamic node: function node that dispatches a task agent via ctx.run_node
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_dynamic_dispatch_of_task_agent(
+    request: pytest.FixtureRequest,
+):
+  """A custom function node can dispatch a task agent and consume its output."""
+  task_agent = _make_task_agent(
+      name='task_agent',
+      responses=[_finish_part({'result': 'dynamic output'})],
+  )
+
+  @node(rerun_on_resume=True)
+  async def driver(*, ctx: Context, node_input: Any):
+    output = await ctx.run_node(task_agent, node_input='go')
+    yield Event(output=f'wrapped: {output}')
+
+  wf = Workflow(name='wf', edges=[(START, driver)])
+
+  app = App(name=request.function.__name__, root_agent=wf)
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  outputs = [e.output for e in events if e.output]
+  assert any(
+      isinstance(o, str) and 'dynamic output' in o for o in outputs
+  ), f'expected wrapped dynamic output, got: {outputs}'
+
+
+# ---------------------------------------------------------------------------
+# 7. Validation error -> retry: wrapper yields the error FR and lets the LLM
+#    emit a corrected finish_task on the next round.
+# ---------------------------------------------------------------------------
+
+
+class _StrictOutput(BaseModel):
+  name: str
+  age: int
+
+
+@pytest.mark.asyncio
+async def test_task_validation_error_drives_retry(
+    request: pytest.FixtureRequest,
+):
+  """Bad finish_task args produce an error FR; the LLM gets a retry."""
+  # First finish_task call has wrong types (age as string), second is correct.
+  child_model = testing_utils.MockModel.create(
+      responses=[
+          _finish_part({'name': 'Jane', 'age': 'thirty'}),
+          _finish_part({'name': 'Jane', 'age': 30}),
+      ]
+  )
+  child = LlmAgent(
+      name='child',
+      model=child_model,
+      mode='task',
+      output_schema=_StrictOutput,
+  )
+
+  root = LlmAgent(
+      name='root',
+      model=testing_utils.MockModel.create(
+          responses=[
+              _delegate_part('child', 'gather identity'),
+              'All set.',
+          ]
+      ),
+      sub_agents=[child],
+  )
+
+  app = App(name=request.function.__name__, root_agent=root)
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  events = await runner.run_async(testing_utils.get_user_content('hi'))
+
+  # The mock LLM was called twice for the child (the bad attempt + the
+  # corrected one), proving the wrapper looped instead of terminating
+  # on the first finish_task.
+  assert child_model.response_index == 1
+  finish_args = _collect_finish_outputs(events)
+  assert finish_args == [
+      {'name': 'Jane', 'age': 'thirty'},
+      {'name': 'Jane', 'age': 30},
+  ]
+  # The validation-error FR should be present in session for the LLM
+  # to see on its retry round.
+  error_frs = [
+      fr.response
+      for e in events
+      for fr in e.get_function_responses()
+      if fr.name == 'finish_task'
+      and isinstance(fr.response, dict)
+      and 'error' in fr.response
+  ]
+  assert len(error_frs) == 1, f'expected one error FR, got {error_frs}'
+
+
+# ---------------------------------------------------------------------------
+# 8. Cross-turn resumption: an unresolved task FC from a prior turn is
+#    re-dispatched by the chat coordinator on the next user turn, before
+#    the LLM is called.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_chat_coordinator_resumes_unresolved_task_fc(
+    request: pytest.FixtureRequest,
+):
+  """Pending task FC from a prior turn is dispatched before the new LLM call."""
+  child_model = testing_utils.MockModel.create(
+      responses=[_finish_part({'result': 'finished after resume'})]
+  )
+  child = LlmAgent(name='child', model=child_model, mode='task')
+
+  root_model = testing_utils.MockModel.create(
+      responses=[
+          # Only response needed: post-resume continuation after the
+          # pre-LLM scan dispatches the pending task and synthesizes its FR.
+          'Resumed and done.',
+      ]
+  )
+  root = LlmAgent(
+      name='root',
+      model=root_model,
+      sub_agents=[child],
+  )
+
+  # Seed the session with an unresolved task delegation FC authored by
+  # root from a "prior turn".  No matching FR exists.
+  from google.adk.sessions.in_memory_session_service import InMemorySessionService
+
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name=request.function.__name__,
+      user_id='u',
+  )
+  await session_service.append_event(
+      session=session,
+      event=Event(
+          invocation_id='prior-inv',
+          author='root',
+          content=types.Content(
+              role='model',
+              parts=[
+                  types.Part(
+                      function_call=types.FunctionCall(
+                          id='fc-pending',
+                          name='child',
+                          args={'request': 'leftover work'},
+                      )
+                  )
+              ],
+          ),
+      ),
+  )
+
+  from google.adk.runners import Runner
+
+  app = App(name=request.function.__name__, root_agent=root)
+  runner = Runner(app=app, session_service=session_service)
+
+  events = []
+  async for ev in runner.run_async(
+      user_id='u',
+      session_id=session.id,
+      new_message=testing_utils.get_user_content('continue'),
+  ):
+    events.append(ev)
+
+  # The child must have been dispatched once (resuming the pending FC).
+  assert (
+      child_model.response_index == 0
+  ), 'child LLM should have been called exactly once for the resumed task'
+  finish_args = _collect_finish_outputs(events)
+  assert {
+      'result': 'finished after resume'
+  } in finish_args, f'expected resumed task to finish; got {finish_args}'
+
+
+# ---------------------------------------------------------------------------
+# 9. Strict isolation filtering: a stranger event with a foreign
+#    isolation_scope must NOT appear in the task agent's LLM context.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_strict_isolation_filter_excludes_foreign_scope(
+    request: pytest.FixtureRequest,
+):
+  """Garbage-scoped events are excluded from the task agent's view."""
+  child_model = testing_utils.MockModel.create(
+      responses=[_finish_part({'result': 'ok'})]
+  )
+  child = LlmAgent(name='child', model=child_model, mode='task')
+
+  root = LlmAgent(
+      name='root',
+      model=testing_utils.MockModel.create(
+          responses=[
+              _delegate_part('child', 'do the thing'),
+              'Done.',
+          ]
+      ),
+      sub_agents=[child],
+  )
+
+  from google.adk.sessions.in_memory_session_service import InMemorySessionService
+
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name=request.function.__name__,
+      user_id='u',
+  )
+  # Seed a stranger event with a different scope.
+  stranger = Event(
+      invocation_id='stranger-inv',
+      author='someone_else',
+      content=types.Content(
+          role='user',
+          parts=[types.Part(text='SECRET-SHOULD-NOT-LEAK')],
+      ),
+  )
+  stranger.isolation_scope = 'garbage-scope'
+  session.events.append(stranger)
+
+  from google.adk.runners import Runner
+
+  app = App(name=request.function.__name__, root_agent=root)
+  runner = Runner(app=app, session_service=session_service)
+
+  async for _ in runner.run_async(
+      user_id='u',
+      session_id=session.id,
+      new_message=testing_utils.get_user_content('go'),
+  ):
+    pass
+
+  # Inspect the child's LLM request: SECRET text must not appear.
+  child_request = child_model.requests[0]
+  rendered = '\n'.join(
+      p.text or '' for c in child_request.contents or [] for p in c.parts or []
+  )
+  assert (
+      'SECRET-SHOULD-NOT-LEAK' not in rendered
+  ), 'stranger event leaked across isolation_scope filter'

tests/unittests/workflow/test_tool_node.py

@@ -0,0 +1,141 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for ToolNode input parsing and execution."""
+
+from typing import Any
+
+from google.adk.events.event import Event
+from google.adk.tools.base_tool import BaseTool
+from google.adk.workflow import START
+from google.adk.workflow._tool_node import _ToolNode as ToolNode
+from google.adk.workflow._workflow import Workflow
+from google.genai import types
+import pytest
+
+from . import workflow_testing_utils
+from .. import testing_utils
+
+
+class MockTool(BaseTool):
+  """A mock tool that returns the args it was called with."""
+
+  def __init__(self, name="mock_tool", description="Mock tool"):
+    super().__init__(name=name, description=description)
+
+  async def run_async(self, *, args: dict[str, Any], tool_context) -> Any:
+    return args
+
+
+async def _run_tool_node_wf(node_input: Any) -> list[Any]:
+  """Runs a workflow with a ToolNode that receives node_input."""
+  tool_node = ToolNode(tool=MockTool())
+
+  def start_node():
+    return Event(output=node_input)
+
+  wf = Workflow(
+      name="tool_node_test_wf",
+      edges=[
+          (START, start_node),
+          (start_node, tool_node),
+      ],
+  )
+  app_instance = testing_utils.App(name="test_app", root_agent=wf)
+  runner = testing_utils.InMemoryRunner(app=app_instance)
+  events = await runner.run_async("start")
+  return workflow_testing_utils.simplify_events_with_node(events)
+
+
+@pytest.mark.asyncio
+async def test_tool_node_accepts_dict():
+  """Tests that ToolNode accepts a dict as input and passes it to the tool."""
+  input_dict = {"param_a": 1, "param_b": "value"}
+  simplified = await _run_tool_node_wf(input_dict)
+  assert (
+      "tool_node_test_wf@1/mock_tool@1",
+      {"output": input_dict},
+  ) in simplified
+
+
+@pytest.mark.asyncio
+async def test_tool_node_accepts_none():
+  """Tests that ToolNode accepts None, converting it to an empty dict."""
+  simplified = await _run_tool_node_wf(None)
+  assert ("tool_node_test_wf@1/mock_tool@1", {"output": {}}) in simplified
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("empty_input", ["", "   ", "\n\t"])
+async def test_tool_node_accepts_empty_string(empty_input):
+  """Tests that ToolNode treats an empty/whitespace string as no arguments."""
+  simplified = await _run_tool_node_wf(empty_input)
+  assert ("tool_node_test_wf@1/mock_tool@1", {"output": {}}) in simplified
+
+
+@pytest.mark.asyncio
+async def test_tool_node_accepts_json_string():
+  """Tests that ToolNode accepts a valid JSON string representing a dict."""
+  json_str = '{"param_a": 1, "param_b": "value"}'
+  simplified = await _run_tool_node_wf(json_str)
+  assert (
+      "tool_node_test_wf@1/mock_tool@1",
+      {"output": {"param_a": 1, "param_b": "value"}},
+  ) in simplified
+
+
+@pytest.mark.asyncio
+async def test_tool_node_accepts_content_with_json_string():
+  """Tests that ToolNode accepts a types.Content containing a JSON string."""
+  json_str = '{"param_a": 1, "param_b": "value"}'
+  content = types.Content(
+      parts=[types.Part.from_text(text=json_str)], role="user"
+  )
+  simplified = await _run_tool_node_wf(content)
+  assert (
+      "tool_node_test_wf@1/mock_tool@1",
+      {"output": {"param_a": 1, "param_b": "value"}},
+  ) in simplified
+
+
+@pytest.mark.asyncio
+async def test_tool_node_rejects_non_dict_json_string():
+  """Tests that ToolNode raises TypeError if JSON string represents a non-dict (e.g. list)."""
+  json_str = "[1, 2, 3]"
+  with pytest.raises(
+      TypeError, match="The input to ToolNode must be a dictionary"
+  ):
+    await _run_tool_node_wf(json_str)
+
+
+@pytest.mark.asyncio
+async def test_tool_node_rejects_invalid_json_string():
+  """Tests that ToolNode raises TypeError if string input is not valid JSON."""
+  invalid_str = "not a json"
+  with pytest.raises(
+      TypeError, match="The input to ToolNode must be a dictionary"
+  ):
+    await _run_tool_node_wf(invalid_str)
+
+
+@pytest.mark.asyncio
+async def test_tool_node_rejects_non_dict_content():
+  """Tests that ToolNode raises TypeError if Content contains non-dict text."""
+  content = types.Content(
+      parts=[types.Part.from_text(text="not a json")], role="user"
+  )
+  with pytest.raises(
+      TypeError, match="The input to ToolNode must be a dictionary"
+  ):
+    await _run_tool_node_wf(content)

tests/unittests/workflow/test_workflow_agent_as_node.py

@@ -0,0 +1,104 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for BaseAgent instances used as nodes in a Workflow."""
+
+from typing import AsyncGenerator
+
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.invocation_context import InvocationContext as BaseInvocationContext
+from google.adk.events.event import Event
+from google.adk.runners import Runner
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.workflow import START
+from google.adk.workflow._workflow import Workflow
+from google.genai import types
+import pytest
+
+from .workflow_testing_utils import InputCapturingNode
+from .workflow_testing_utils import simplify_events_with_node
+
+
+class SimpleAgent(BaseAgent):
+  """A simple agent for testing."""
+
+  message: str = ''
+
+  async def _run_async_impl(
+      self, ctx: BaseInvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    """Yields a single event with a message."""
+    yield Event(
+        author=self.name,
+        invocation_id=ctx.invocation_id,
+        content=types.Content(parts=[types.Part(text=self.message)]),
+    )
+
+
+@pytest.mark.asyncio
+async def test_run_async_with_agent_nodes(request: pytest.FixtureRequest):
+  """BaseAgent nodes emit content events through the workflow."""
+  agent_a = SimpleAgent(name='AgentA', message='Hello')
+  agent_b = SimpleAgent(name='AgentB', message='World')
+  wf = Workflow(
+      name='wf_with_agents',
+      edges=[
+          (START, agent_a),
+          (agent_a, agent_b),
+      ],
+  )
+  ss = InMemorySessionService()
+  runner = Runner(app_name='test', node=wf, session_service=ss)
+  session = await ss.create_session(app_name='test', user_id='u')
+
+  msg = types.Content(parts=[types.Part(text='start')], role='user')
+  events: list[Event] = []
+  async for event in runner.run_async(
+      user_id='u', session_id=session.id, new_message=msg
+  ):
+    events.append(event)
+
+  assert simplify_events_with_node(events) == [
+      ('wf_with_agents@1/AgentA@1', 'Hello'),
+      ('wf_with_agents@1/AgentB@1', 'World'),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_run_async_with_agent_node_piping_data(
+    request: pytest.FixtureRequest,
+):
+  """AgentNode content is not piped as output to the next node."""
+  agent_a = SimpleAgent(name='AgentA', message='Hello')
+  node_b = InputCapturingNode(name='NodeB')
+  wf = Workflow(
+      name='wf_with_agent_piping',
+      edges=[
+          (START, agent_a),
+          (agent_a, node_b),
+      ],
+  )
+  ss = InMemorySessionService()
+  runner = Runner(app_name='test', node=wf, session_service=ss)
+  session = await ss.create_session(app_name='test', user_id='u')
+
+  msg = types.Content(parts=[types.Part(text='start')], role='user')
+  async for _ in runner.run_async(
+      user_id='u', session_id=session.id, new_message=msg
+  ):
+    pass
+
+  # AgentNode does not record content as node output, so the next node
+  # receives None as input.
+  assert node_b.received_inputs == [None]

tests/unittests/workflow/test_workflow_bytes.py

@@ -0,0 +1,137 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for Workflow handling of bytes and serialization."""
+
+from __future__ import annotations
+
+from typing import Any
+from typing import AsyncGenerator
+
+from google.adk.agents.context import Context
+from google.adk.events.event import Event
+from google.adk.runners import Runner
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.workflow import BaseNode
+from google.adk.workflow import START
+from google.adk.workflow._workflow import Workflow
+from google.genai import types
+from pydantic import ConfigDict
+from pydantic import Field
+import pytest
+
+# --- Helpers ---
+
+
+class _BytesOutputNode(BaseNode):
+  """Yields bytes content or raw bytes."""
+
+  raw_bytes: bool = False
+
+  async def _run_impl(
+      self, *, ctx: Context, node_input: Any
+  ) -> AsyncGenerator[Any, None]:
+    data = b"\x89PNG\r\n\x1a\n"
+    if self.raw_bytes:
+      yield data
+    else:
+      yield Event(
+          content=types.Content(
+              parts=[types.Part.from_bytes(data=data, mime_type="image/png")]
+          ),
+          output="bytes_sent",
+      )
+
+
+class _InputCapturingNode(BaseNode):
+  """Captures node_input for later assertion."""
+
+  model_config = ConfigDict(arbitrary_types_allowed=True)
+  received_inputs: list[Any] = Field(default_factory=list)
+
+  async def _run_impl(
+      self, *, ctx: Context, node_input: Any
+  ) -> AsyncGenerator[Any, None]:
+    self.received_inputs.append(node_input)
+    yield {"received": node_input}
+
+
+async def _run_workflow(wf, message="start"):
+  """Run a Workflow through Runner, return collected events."""
+  ss = InMemorySessionService()
+  runner = Runner(app_name="test", node=wf, session_service=ss)
+  session = await ss.create_session(app_name="test", user_id="u")
+  msg = types.Content(parts=[types.Part(text=message)], role="user")
+  events = []
+  async for event in runner.run_async(
+      user_id="u", session_id=session.id, new_message=msg
+  ):
+    events.append(event)
+  return events, ss, session
+
+
+# --- Tests ---
+
+
+@pytest.mark.asyncio
+async def test_bytes_in_content_output():
+  """Content with bytes propagates to downstream node."""
+  a = _BytesOutputNode(name="a", raw_bytes=False)
+  b = _InputCapturingNode(name="b")
+  wf = Workflow(name="wf", edges=[(START, a, b)])
+
+  events, _, _ = await _run_workflow(wf)
+
+  assert b.received_inputs == ["bytes_sent"]
+
+
+@pytest.mark.asyncio
+async def test_raw_bytes_output():
+  """Raw bytes output propagates to downstream node."""
+  a = _BytesOutputNode(name="a", raw_bytes=True)
+  b = _InputCapturingNode(name="b")
+  wf = Workflow(name="wf", edges=[(START, a, b)])
+
+  events, _, _ = await _run_workflow(wf)
+
+  assert len(b.received_inputs) == 1
+  assert isinstance(b.received_inputs[0], bytes)
+
+
+@pytest.mark.xfail(reason="Checkpoint/resume not yet in new Workflow.")
+@pytest.mark.asyncio
+async def test_bytes_in_node_input_serialization():
+  """Bytes in node input survive checkpoint/resume."""
+  assert False, "TODO"
+
+
+@pytest.mark.xfail(reason="Checkpoint/resume not yet in new Workflow.")
+@pytest.mark.asyncio
+async def test_bytes_in_typed_model_input():
+  """Bytes in Pydantic model input survive round-trip."""
+  assert False, "TODO"
+
+
+@pytest.mark.xfail(reason="Checkpoint/resume not yet in new Workflow.")
+@pytest.mark.asyncio
+async def test_bytes_in_trigger_buffer():
+  """Bytes in trigger buffer survive serialization."""
+  assert False, "TODO"
+
+
+@pytest.mark.xfail(reason="Checkpoint/resume not yet in new Workflow.")
+@pytest.mark.asyncio
+async def test_bytes_full_workflow_resume():
+  """Full resume with bytes data end-to-end."""
+  assert False, "TODO"

tests/unittests/workflow/test_workflow_concurrency.py

@@ -0,0 +1,132 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+from typing import Any
+from typing import AsyncGenerator
+
+from google.adk.agents.context import Context
+from google.adk.apps.app import App
+from google.adk.events.event import Event
+from google.adk.workflow import BaseNode
+from google.adk.workflow import START
+from google.adk.workflow._parallel_worker import _ParallelWorker as ParallelWorker
+from google.adk.workflow._workflow import Workflow
+from pydantic import Field
+import pytest
+from typing_extensions import override
+
+from .. import testing_utils
+from .workflow_testing_utils import create_parent_invocation_context
+
+
+@pytest.mark.asyncio
+async def test_max_concurrency_limits_running_nodes(
+    request: pytest.FixtureRequest,
+):
+  """Max concurrency limits the number of parallel graph-scheduled nodes.
+
+  Setup:
+    Workflow with 4 parallel nodes and max_concurrency=2.
+  Act:
+    - Start workflow in background.
+    - Release nodes one by one.
+  Assert:
+    - Initially only 2 nodes start.
+    - Releasing one node allows another to start.
+    - All nodes eventually complete.
+  """
+
+  class ConcurrencyWorkerNode(BaseNode):
+    """A node that signals when it starts and waits for a signal to finish."""
+
+    started_event: asyncio.Event
+    finish_event: asyncio.Event
+
+    @override
+    async def _run_impl(
+        self,
+        *,
+        ctx: Context,
+        node_input: Any,
+    ) -> AsyncGenerator[Any, None]:
+      self.started_event.set()
+      await self.finish_event.wait()
+      yield f'{self.name}_done'
+
+  class TerminalNode(BaseNode):
+
+    @override
+    async def _run_impl(
+        self, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Any, None]:
+      yield 'workflow_done'
+
+  # Given a workflow with 4 parallel nodes and max_concurrency=2
+  num_nodes = 4
+  max_concurrency = 2
+  started_events = [asyncio.Event() for _ in range(num_nodes)]
+  finish_events = [asyncio.Event() for _ in range(num_nodes)]
+
+  nodes = [
+      ConcurrencyWorkerNode(
+          name=f'Node{i}',
+          started_event=started_events[i],
+          finish_event=finish_events[i],
+      )
+      for i in range(num_nodes)
+  ]
+
+  terminal_node = TerminalNode(name='Terminal')
+  edges = [(START, tuple(nodes), terminal_node)]
+
+  agent = Workflow(
+      name='concurrency_agent',
+      max_concurrency=max_concurrency,
+      edges=edges,
+  )
+
+  app = App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  # When the workflow is run in background
+  async def run_agent():
+    return await runner.run_async(testing_utils.get_user_content('start'))
+
+  run_task = asyncio.create_task(run_agent())
+
+  # Then initially only max_concurrency nodes should start
+  await asyncio.sleep(0.1)
+  started_count = sum(1 for e in started_events if e.is_set())
+  assert started_count == max_concurrency
+
+  # When one node is released
+  for i in range(num_nodes):
+    if started_events[i].is_set():
+      finish_events[i].set()
+      break
+
+  # Then another node should start, bringing total to max_concurrency + 1
+  await asyncio.sleep(0.1)
+  started_count = sum(1 for e in started_events if e.is_set())
+  assert started_count == max_concurrency + 1
+
+  # When all remaining nodes are released
+  for e in finish_events:
+    e.set()
+
+  # Then all nodes should eventually complete
+  await run_task
+  started_count = sum(1 for e in started_events if e.is_set())
+  assert started_count == num_nodes

tests/unittests/workflow/test_workflow_function_tool_as_node.py

@@ -0,0 +1,75 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for FunctionTool nodes in a Workflow."""
+
+from google.adk.events.event import Event
+from google.adk.runners import Runner
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.tools.function_tool import FunctionTool
+from google.adk.workflow import START
+from google.adk.workflow._workflow import Workflow
+from google.genai import types
+import pytest
+
+from .workflow_testing_utils import simplify_events_with_node
+
+
+def _produce_input() -> None:
+  """Absorbs user content so downstream ToolNodes receive None."""
+  return None
+
+
+def _func_a() -> dict[str, str]:
+  """Returns a value from function A."""
+  return {'val': 'Hello'}
+
+
+def _func_b(val: str) -> str:
+  """Returns a value incorporating input from A."""
+  return f'{val}_world'
+
+
+@pytest.mark.asyncio
+async def test_run_async_with_function_tools():
+  """FunctionTool output is piped as input to the next FunctionTool."""
+  tool_a = FunctionTool(_func_a)
+  tool_b = FunctionTool(_func_b)
+  wf = Workflow(
+      name='wf_with_tools',
+      edges=[
+          (START, _produce_input, tool_a, tool_b),
+      ],
+  )
+  ss = InMemorySessionService()
+  runner = Runner(app_name='test', node=wf, session_service=ss)
+  session = await ss.create_session(app_name='test', user_id='u')
+
+  msg = types.Content(parts=[types.Part(text='start')], role='user')
+  events: list[Event] = []
+  async for event in runner.run_async(
+      user_id='u', session_id=session.id, new_message=msg
+  ):
+    events.append(event)
+
+  assert simplify_events_with_node(events) == [
+      (
+          'wf_with_tools@1/_func_a@1',
+          {'output': {'val': 'Hello'}},
+      ),
+      (
+          'wf_with_tools@1/_func_b@1',
+          {'output': 'Hello_world'},
+      ),
+  ]

tests/unittests/workflow/test_workflow_live.py

@@ -0,0 +1,653 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+from typing import Any
+from typing import AsyncGenerator
+
+from google.adk.agents.context import Context
+from google.adk.agents.live_request_queue import LiveRequestQueue
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.events.event import Event
+from google.adk.models.llm_response import LlmResponse
+from google.adk.runners import Runner
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.workflow._base_node import BaseNode
+from google.adk.workflow._base_node import START
+from google.adk.workflow._workflow import Workflow
+from google.genai import types
+from pydantic import Field
+import pytest
+
+from . import testing_utils
+
+# --- Mock Nodes and Agents for Testing Live Mode Design ---
+
+
+class _MockNonLiveNode(BaseNode):
+  """A standard non-live node whose signature does NOT accept live_request_queue."""
+
+  called: bool = False
+  actual_input: Any = None
+  shared_state: dict[str, Any] = Field(default_factory=dict)
+
+  def __init__(self, *, name: str):
+    super().__init__(name=name)
+
+  async def _run_impl(
+      self,
+      *,
+      ctx: Context,
+      node_input: Any,
+  ) -> AsyncGenerator[Any, None]:
+    self.called = True
+    self.actual_input = node_input
+    self.shared_state["called"] = True
+    self.shared_state["actual_input"] = node_input
+    yield Event(output=f"{self.name}_output")
+
+
+class _ConstantNode(BaseNode):
+  """A node that outputs a constant value."""
+
+  output_value: Any = None
+
+  def __init__(self, *, name: str, output_value: Any):
+    super().__init__(name=name)
+    self.output_value = output_value
+
+  async def _run_impl(
+      self,
+      *,
+      ctx: Context,
+      node_input: Any,
+  ) -> AsyncGenerator[Any, None]:
+    yield Event(output=self.output_value)
+
+
+# --- Live Workflow Unit Tests (TDD) ---
+
+
+@pytest.mark.asyncio
+async def test_hybrid_live_non_live_nodes():
+  """CUJ 1: A workflow has hybrid live & non-live nodes."""
+  mock_model1 = testing_utils.MockModel.create(
+      responses=[
+          LlmResponse(
+              content=testing_utils.ModelContent(
+                  [types.Part.from_text(text="Acknowledged: node1_start")]
+              )
+          ),
+          LlmResponse(
+              content=testing_utils.ModelContent(
+                  [types.Part.from_text(text="Acknowledged: node1_end")]
+              )
+          ),
+          LlmResponse(
+              content=testing_utils.ModelContent([
+                  types.Part.from_function_call(
+                      name="finish_task",
+                      args={"result": "LiveNode1_output"},
+                  )
+              ])
+          ),
+      ]
+  )
+  live_node1 = LlmAgent(
+      name="LiveNode1",
+      model=mock_model1,
+      mode="task",
+      instruction="Handle live interaction 1.",
+  )
+  non_live_node = _MockNonLiveNode(name="NonLiveNode")
+  mock_model2 = testing_utils.MockModel.create(
+      responses=[
+          LlmResponse(
+              content=testing_utils.ModelContent(
+                  [types.Part.from_text(text="Acknowledged: node2_start")]
+              )
+          ),
+          LlmResponse(
+              content=testing_utils.ModelContent(
+                  [types.Part.from_text(text="Acknowledged: node2_end")]
+              )
+          ),
+          LlmResponse(
+              content=testing_utils.ModelContent([
+                  types.Part.from_function_call(
+                      name="finish_task",
+                      args={"result": "LiveNode2_output"},
+                  )
+              ])
+          ),
+      ]
+  )
+  live_node2 = LlmAgent(
+      name="LiveNode2",
+      model=mock_model2,
+      mode="task",
+      instruction="Handle live interaction 2.",
+  )
+
+  wf = Workflow(
+      name="hybrid_workflow",
+      edges=[
+          (START, live_node1),
+          (live_node1, non_live_node),
+          (non_live_node, live_node2),
+      ],
+  )
+
+  live_queue = LiveRequestQueue()
+
+  # Pre-seed first live node's requests
+  live_queue.send_realtime(
+      types.Blob(data=b"node1_start", mime_type="audio/pcm")
+  )
+  live_queue.send_realtime(types.Blob(data=b"node1_end", mime_type="audio/pcm"))
+
+  ss = InMemorySessionService()
+  runner = Runner(app_name=wf.name, node=wf, session_service=ss)
+  session = await ss.create_session(app_name=wf.name, user_id="u")
+
+  events = []
+  async for event in runner.run_live(
+      user_id="u",
+      session_id=session.id,
+      live_request_queue=live_queue,
+      run_config=testing_utils.RunConfig(
+          get_session_config={"state_delta": {"__START__": "start_input"}}
+      ),
+  ):
+    events.append(event)
+    if event.output == "NonLiveNode_output":
+      # First live node and non-live node completed! Now feed the second live node's requests:
+      live_queue.send_realtime(
+          types.Blob(data=b"node2_start", mime_type="audio/pcm")
+      )
+      live_queue.send_realtime(
+          types.Blob(data=b"node2_end", mime_type="audio/pcm")
+      )
+
+  # 1. Assert exact outputs sequence
+  outputs = [e.output for e in events if e.output is not None]
+  assert outputs == [
+      {"result": "LiveNode1_output"},
+      "NonLiveNode_output",
+      {"result": "LiveNode2_output"},
+  ]
+  assert non_live_node.shared_state.get("actual_input") == {
+      "result": "LiveNode1_output"
+  }
+
+  # 2. Assert intermediate content events (conversational turns)
+  content_texts = [
+      p.text
+      for e in events
+      if e.content and e.content.parts and e.output is None
+      for p in e.content.parts
+      if p.text
+  ]
+  assert content_texts == [
+      "Acknowledged: node1_start",
+      "Acknowledged: node1_end",
+      "Acknowledged: node2_start",
+      "Acknowledged: node2_end",
+  ]
+
+  # 3. Assert live requests fed to the models
+  assert [b.data for b in mock_model1.live_blobs] == [
+      b"node1_start",
+      b"node1_end",
+  ]
+  assert [b.data for b in mock_model2.live_blobs] == [
+      b"node2_start",
+      b"node2_end",
+  ]
+
+
+@pytest.mark.asyncio
+async def test_nested_workflow_has_live_node():
+  """CUJ 2: A nested workflow has a live node."""
+  mock_model = testing_utils.MockModel.create(
+      responses=[
+          LlmResponse(
+              content=testing_utils.ModelContent(
+                  [types.Part.from_text(text="Acknowledged: inner_start")]
+              )
+          ),
+          LlmResponse(
+              content=testing_utils.ModelContent(
+                  [types.Part.from_text(text="Acknowledged: inner_end")]
+              )
+          ),
+          LlmResponse(
+              content=testing_utils.ModelContent([
+                  types.Part.from_function_call(
+                      name="finish_task",
+                      args={"result": "InnerLiveNode_output"},
+                  )
+              ])
+          ),
+      ]
+  )
+  live_node = LlmAgent(
+      name="InnerLiveNode",
+      model=mock_model,
+      mode="task",
+      instruction="Handle inner live interaction.",
+  )
+  inner_wf = Workflow(name="inner_wf", edges=[(START, live_node)])
+  outer_wf = Workflow(name="outer_wf", edges=[(START, inner_wf)])
+
+  live_queue = LiveRequestQueue()
+  live_queue.send_realtime(
+      types.Blob(data=b"inner_start", mime_type="audio/pcm")
+  )
+  live_queue.send_realtime(types.Blob(data=b"inner_end", mime_type="audio/pcm"))
+
+  ss = InMemorySessionService()
+  runner = Runner(app_name=outer_wf.name, node=outer_wf, session_service=ss)
+  session = await ss.create_session(app_name=outer_wf.name, user_id="u")
+
+  events = []
+  async for event in runner.run_live(
+      user_id="u",
+      session_id=session.id,
+      live_request_queue=live_queue,
+      run_config=testing_utils.RunConfig(
+          get_session_config={"state_delta": {"__START__": "start_input"}}
+      ),
+  ):
+    events.append(event)
+
+  # Assert exact outputs sequence
+  outputs = [e.output for e in events if e.output is not None]
+  assert outputs == [{"result": "InnerLiveNode_output"}]
+
+  # Assert content events
+  content_texts = [
+      p.text
+      for e in events
+      if e.content and e.content.parts and e.output is None
+      for p in e.content.parts
+      if p.text
+  ]
+  assert content_texts == [
+      "Acknowledged: inner_start",
+      "Acknowledged: inner_end",
+  ]
+  assert [b.data for b in mock_model.live_blobs] == [
+      b"inner_start",
+      b"inner_end",
+  ]
+
+
+@pytest.mark.asyncio
+async def test_nested_live_node_and_outer_live_node():
+  """CUJ 3: A nested workflow has live node & outer workflow then has a live node."""
+  mock_model_inner = testing_utils.MockModel.create(
+      responses=[
+          LlmResponse(
+              content=testing_utils.ModelContent(
+                  [types.Part.from_text(text="Acknowledged: inner_start")]
+              )
+          ),
+          LlmResponse(
+              content=testing_utils.ModelContent(
+                  [types.Part.from_text(text="Acknowledged: inner_end")]
+              )
+          ),
+          LlmResponse(
+              content=testing_utils.ModelContent([
+                  types.Part.from_function_call(
+                      name="finish_task",
+                      args={"result": "InnerLiveNode_output"},
+                  )
+              ])
+          ),
+      ]
+  )
+  inner_live = LlmAgent(
+      name="InnerLiveNode",
+      model=mock_model_inner,
+      mode="task",
+      instruction="Handle inner live interaction.",
+  )
+  inner_wf = Workflow(name="inner_wf", edges=[(START, inner_live)])
+
+  mock_model_outer = testing_utils.MockModel.create(
+      responses=[
+          LlmResponse(
+              content=testing_utils.ModelContent(
+                  [types.Part.from_text(text="Acknowledged: outer_start")]
+              )
+          ),
+          LlmResponse(
+              content=testing_utils.ModelContent(
+                  [types.Part.from_text(text="Acknowledged: outer_end")]
+              )
+          ),
+          LlmResponse(
+              content=testing_utils.ModelContent([
+                  types.Part.from_function_call(
+                      name="finish_task",
+                      args={"result": "OuterLiveNode_output"},
+                  )
+              ])
+          ),
+      ]
+  )
+  outer_live = LlmAgent(
+      name="OuterLiveNode",
+      model=mock_model_outer,
+      mode="task",
+      instruction="Handle outer live interaction.",
+  )
+  prep_node = _ConstantNode(name="PrepNode", output_value="prep_output")
+
+  wf = Workflow(
+      name="nested_sequential_live",
+      edges=[
+          (START, inner_wf),
+          (inner_wf, prep_node),
+          (prep_node, outer_live),
+      ],
+  )
+
+  live_queue = LiveRequestQueue()
+  live_queue.send_realtime(
+      types.Blob(data=b"inner_start", mime_type="audio/pcm")
+  )
+  live_queue.send_realtime(types.Blob(data=b"inner_end", mime_type="audio/pcm"))
+
+  ss = InMemorySessionService()
+  runner = Runner(app_name=wf.name, node=wf, session_service=ss)
+  session = await ss.create_session(app_name=wf.name, user_id="u")
+
+  events = []
+  async for event in runner.run_live(
+      user_id="u",
+      session_id=session.id,
+      live_request_queue=live_queue,
+      run_config=testing_utils.RunConfig(
+          get_session_config={"state_delta": {"__START__": "start_input"}}
+      ),
+  ):
+    events.append(event)
+    if event.output == "prep_output":
+      # Inner live node and prep node completed! Feed outer node's requests:
+      live_queue.send_realtime(
+          types.Blob(data=b"outer_start", mime_type="audio/pcm")
+      )
+      live_queue.send_realtime(
+          types.Blob(data=b"outer_end", mime_type="audio/pcm")
+      )
+
+  # Assert exact outputs sequence
+  outputs = [e.output for e in events if e.output is not None]
+  assert outputs == [
+      {"result": "InnerLiveNode_output"},
+      "prep_output",
+      {"result": "OuterLiveNode_output"},
+  ]
+
+  # Assert content events
+  content_texts = [
+      p.text
+      for e in events
+      if e.content and e.content.parts and e.output is None
+      for p in e.content.parts
+      if p.text
+  ]
+  assert content_texts == [
+      "Acknowledged: inner_start",
+      "Acknowledged: inner_end",
+      "Acknowledged: outer_start",
+      "Acknowledged: outer_end",
+  ]
+  assert [b.data for b in mock_model_inner.live_blobs] == [
+      b"inner_start",
+      b"inner_end",
+  ]
+  assert [b.data for b in mock_model_outer.live_blobs] == [
+      b"outer_start",
+      b"outer_end",
+  ]
+
+
+@pytest.mark.asyncio
+async def test_dynamic_node_scheduling_of_live_node():
+  """CUJ 4: A node in workflow dynamically schedules a live node using ctx.run_node()."""
+
+  class _DynamicLiveSchedulerNode(BaseNode):
+    """A node that dynamically schedules a child live node using ctx.run_node()."""
+
+    child_node: BaseNode | None = None
+    child_output: Any = None
+    shared_state: dict[str, Any] = Field(default_factory=dict)
+
+    def __init__(self, *, name: str, child_node: BaseNode):
+      super().__init__(name=name, rerun_on_resume=True)
+      self.child_node = child_node
+
+    async def _run_impl(
+        self,
+        *,
+        ctx: Context,
+        node_input: Any,
+    ) -> AsyncGenerator[Any, None]:
+      if self.child_node:
+        output = await ctx.run_node(self.child_node, node_input=node_input)
+        self.child_output = output
+        self.shared_state["child_output"] = output
+      yield Event(output=f"{self.name}_output")
+
+  mock_model = testing_utils.MockModel.create(
+      responses=[
+          LlmResponse(
+              content=testing_utils.ModelContent(
+                  [types.Part.from_text(text="Acknowledged: dynamic_start")]
+              )
+          ),
+          LlmResponse(
+              content=testing_utils.ModelContent(
+                  [types.Part.from_text(text="Acknowledged: dynamic_end")]
+              )
+          ),
+          LlmResponse(
+              content=testing_utils.ModelContent([
+                  types.Part.from_function_call(
+                      name="finish_task",
+                      args={"result": "DynamicLiveNode_output"},
+                  )
+              ])
+          ),
+      ]
+  )
+  live_node = LlmAgent(
+      name="DynamicLiveNode",
+      model=mock_model,
+      mode="task",
+      instruction="Handle dynamic live interaction.",
+  )
+  scheduler_node = _DynamicLiveSchedulerNode(
+      name="SchedulerNode", child_node=live_node
+  )
+
+  wf = Workflow(name="dynamic_wf", edges=[(START, scheduler_node)])
+
+  live_queue = LiveRequestQueue()
+  live_queue.send_realtime(
+      types.Blob(data=b"dynamic_start", mime_type="audio/pcm")
+  )
+  live_queue.send_realtime(
+      types.Blob(data=b"dynamic_end", mime_type="audio/pcm")
+  )
+
+  ss = InMemorySessionService()
+  runner = Runner(app_name=wf.name, node=wf, session_service=ss)
+  session = await ss.create_session(app_name=wf.name, user_id="u")
+
+  events = []
+  async for event in runner.run_live(
+      user_id="u",
+      session_id=session.id,
+      live_request_queue=live_queue,
+      run_config=testing_utils.RunConfig(
+          get_session_config={"state_delta": {"__START__": "start_input"}}
+      ),
+  ):
+    events.append(event)
+
+  # Assert exact outputs sequence
+  outputs = [e.output for e in events if e.output is not None]
+  assert outputs == [
+      {"result": "DynamicLiveNode_output"},
+      "SchedulerNode_output",
+  ]
+  assert scheduler_node.shared_state.get("child_output") == {
+      "result": "DynamicLiveNode_output"
+  }
+
+  # Assert content events
+  content_texts = [
+      p.text
+      for e in events
+      if e.content and e.content.parts and e.output is None
+      for p in e.content.parts
+      if p.text
+  ]
+  assert content_texts == [
+      "Acknowledged: dynamic_start",
+      "Acknowledged: dynamic_end",
+  ]
+  assert [b.data for b in mock_model.live_blobs] == [
+      b"dynamic_start",
+      b"dynamic_end",
+  ]
+
+
+@pytest.mark.asyncio
+async def test_live_node_output_passed_to_downstream():
+  """CUJ 5: Dedicated test verifying output of a live node is passed to the next node."""
+  mock_model = testing_utils.MockModel.create(
+      responses=[
+          LlmResponse(
+              content=testing_utils.ModelContent([
+                  types.Part.from_function_call(
+                      name="finish_task",
+                      args={"result": "LiveNode_output"},
+                  )
+              ])
+          ),
+      ]
+  )
+  live_node = LlmAgent(
+      name="LiveNode",
+      model=mock_model,
+      mode="task",
+      instruction="Handle live interaction.",
+  )
+  non_live_node = _MockNonLiveNode(name="NonLiveNode")
+
+  wf = Workflow(
+      name="dataflow_workflow",
+      edges=[
+          (START, live_node),
+          (live_node, non_live_node),
+      ],
+  )
+
+  live_queue = LiveRequestQueue()
+  live_queue.send_realtime(types.Blob(data=b"start_msg", mime_type="audio/pcm"))
+  live_queue.send_realtime(types.Blob(data=b"end_msg", mime_type="audio/pcm"))
+
+  ss = InMemorySessionService()
+  runner = Runner(app_name=wf.name, node=wf, session_service=ss)
+  session = await ss.create_session(app_name=wf.name, user_id="u")
+
+  events = []
+  async for event in runner.run_live(
+      user_id="u",
+      session_id=session.id,
+      live_request_queue=live_queue,
+      run_config=testing_utils.RunConfig(
+          get_session_config={"state_delta": {"__START__": "start_input"}}
+      ),
+  ):
+    events.append(event)
+
+  outputs = [e.output for e in events if e.output is not None]
+  assert outputs == [{"result": "LiveNode_output"}, "NonLiveNode_output"]
+  assert non_live_node.shared_state.get("actual_input") == {
+      "result": "LiveNode_output"
+  }, "The downstream node must receive the live node's exact output"
+  assert [b.data for b in mock_model.live_blobs] == [b"start_msg", b"end_msg"]
+
+
+@pytest.mark.asyncio
+async def test_single_turn_agent_runs_as_non_live_in_live_session():
+  """CUJ 6: A single_turn LlmAgent in a live session runs as non-live and consumes node_input."""
+  mock_model = testing_utils.MockModel.create(
+      responses=[
+          "SingleTurn_output",
+      ]
+  )
+  prep_node = _ConstantNode(
+      name="ConstantNode", output_value="initial_text_input"
+  )
+  single_turn_node = LlmAgent(
+      name="SingleTurnNode",
+      model=mock_model,
+      mode="single_turn",
+      instruction="Summarize the input.",
+  )
+  capture = _MockNonLiveNode(name="capture")
+
+  wf = Workflow(
+      name="single_turn_live_wf",
+      edges=[
+          (START, prep_node),
+          (prep_node, single_turn_node),
+          (single_turn_node, capture),
+      ],
+  )
+
+  live_queue = LiveRequestQueue()
+  live_queue.send_realtime(
+      types.Blob(data=b"ignored_audio", mime_type="audio/pcm")
+  )
+
+  ss = InMemorySessionService()
+  runner = Runner(app_name=wf.name, node=wf, session_service=ss)
+  session = await ss.create_session(app_name=wf.name, user_id="u")
+
+  events = []
+  async for event in runner.run_live(
+      user_id="u",
+      session_id=session.id,
+      live_request_queue=live_queue,
+  ):
+    events.append(event)
+
+  outputs = [e.output for e in events if e.output is not None]
+  assert outputs == ["initial_text_input", "capture_output"]
+  assert capture.shared_state.get("actual_input") == "SingleTurn_output"
+  # Verify that the model received the initial_text_input (node_input) and NOT the live queue audio
+  assert len(mock_model.requests) == 1
+  assert (
+      mock_model.requests[0].contents[0].parts[0].text == "initial_text_input"
+  )
+  assert mock_model.live_blobs == []

tests/unittests/workflow/test_workflow_llm_agent_interruptions.py

@@ -0,0 +1,933 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import asyncio
+import copy
+from typing import Any
+from typing import AsyncGenerator
+
+from google.adk.agents import LlmAgent
+from google.adk.agents.context import Context
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.run_config import RunConfig
+from google.adk.apps.app import App
+from google.adk.events.event import Event
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.sessions.session import Session
+from google.adk.tools.long_running_tool import LongRunningFunctionTool
+from google.adk.tools.tool_context import ToolContext
+from google.adk.workflow import Edge
+from google.adk.workflow import START
+from google.adk.workflow._workflow import Workflow
+from google.genai import types
+import pytest
+
+from tests.unittests import testing_utils
+from tests.unittests.workflow import workflow_testing_utils
+
+
+def long_running_tool_func():
+  """A test tool that simulates a long-running operation."""
+  return None
+
+
+@pytest.mark.asyncio
+async def test_workflow_pause_and_resume_simple(
+    request: pytest.FixtureRequest,
+):
+  """Tests that a workflow can pause and resume with a single LlmAgent node."""
+
+  mock_model = testing_utils.MockModel.create(
+      responses=[
+          types.Part.from_function_call(
+              name='long_running_tool_func',
+              args={},
+          ),
+          types.Part.from_text(text='LLM response after tool'),
+      ]
+  )
+
+  # 1. Create agent with LRO tool
+  node_a = LlmAgent(
+      name='my_agent',
+      model=mock_model,
+      tools=[LongRunningFunctionTool(func=long_running_tool_func)],
+  )
+
+  # 2. Create workflow with single node
+  wf = Workflow(
+      name='test_workflow_hitl',
+      edges=[
+          (START, node_a),
+      ],
+  )
+
+  app = App(
+      name=request.function.__name__,
+      root_agent=wf,
+  )
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  # 3. First run: should pause on the long-running function call.
+  user_event = testing_utils.get_user_content('start workflow')
+  events1 = await runner.run_async(user_event)
+
+  # Verify it paused on LRO
+  assert any(e.long_running_tool_ids for e in events1)
+
+  invocation_id = events1[0].invocation_id
+  fc_event = workflow_testing_utils.find_function_call_event(
+      events1, 'long_running_tool_func'
+  )
+  assert fc_event is not None
+  function_call_id = fc_event.content.parts[0].function_call.id
+
+  # 4. Prepare resume message
+  tool_response = testing_utils.UserContent(
+      types.Part(
+          function_response=types.FunctionResponse(
+              id=function_call_id,
+              name='long_running_tool_func',
+              response={'result': 'Final tool output'},
+          )
+      )
+  )
+
+  # 5. Resume with tool output.
+  events2 = await runner.run_async(
+      new_message=tool_response,
+      invocation_id=invocation_id,
+  )
+
+  # 6. Verify completion
+  content_texts = [
+      p.text
+      for e in events2
+      if e.content and e.content.parts
+      for p in e.content.parts
+      if p.text
+  ]
+
+  assert any('LLM response after tool' in t for t in content_texts)
+
+
+@pytest.mark.asyncio
+async def test_workflow_pause_and_resume_task_mode(
+    request: pytest.FixtureRequest,
+):
+  """Tests that a workflow can pause and resume with a single LlmAgent node in task mode."""
+
+  mock_model = testing_utils.MockModel.create(
+      responses=[
+          types.Part.from_function_call(
+              name='long_running_tool_func',
+              args={},
+          ),
+          types.Part.from_text(text='LLM response after tool in task mode'),
+      ]
+  )
+
+  # 1. Create agent with LRO tool and mode='task'
+  node_a = LlmAgent(
+      name='my_task_agent',
+      model=mock_model,
+      tools=[LongRunningFunctionTool(func=long_running_tool_func)],
+      mode='task',
+  )
+
+  # 2. Create workflow with single node
+  wf = Workflow(
+      name='test_workflow_task_hitl',
+      edges=[
+          (START, node_a),
+      ],
+  )
+
+  app = App(
+      name=request.function.__name__,
+      root_agent=wf,
+  )
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  # 3. First run: should pause on the long-running function call.
+  user_event = testing_utils.get_user_content('start workflow')
+  events1 = await runner.run_async(user_event)
+
+  # Verify it paused on LRO
+  assert any(e.long_running_tool_ids for e in events1)
+
+  invocation_id = events1[0].invocation_id
+  fc_event = workflow_testing_utils.find_function_call_event(
+      events1, 'long_running_tool_func'
+  )
+  assert fc_event is not None
+  function_call_id = fc_event.content.parts[0].function_call.id
+
+  # 4. Prepare resume message
+  tool_response = testing_utils.UserContent(
+      types.Part(
+          function_response=types.FunctionResponse(
+              id=function_call_id,
+              name='long_running_tool_func',
+              response={'result': 'Final tool output'},
+          )
+      )
+  )
+
+  # 5. Resume with tool output.
+  events2 = await runner.run_async(
+      new_message=tool_response,
+      invocation_id=invocation_id,
+  )
+
+  # 6. Verify completion
+  content_texts = [
+      p.text
+      for e in events2
+      if e.content and e.content.parts
+      for p in e.content.parts
+      if p.text
+  ]
+
+  assert any('LLM response after tool in task mode' in t for t in content_texts)
+
+
+@pytest.mark.asyncio
+async def test_workflow_pause_and_resume_tool_confirmation(
+    request: pytest.FixtureRequest,
+):
+  """Tests that a workflow can pause and resume with a tool requiring confirmation.
+
+  Setup: Workflow with a single LlmAgent node having a tool requiring confirmation.
+  Act:
+    - Run 1: Start workflow, tool requests confirmation.
+    - Run 2: Send confirmation response.
+  Assert:
+    - Run 1: Workflow pauses and yields confirmation request.
+    - Run 2: Workflow resumes and completes with LLM response.
+  """
+  from google.adk.flows.llm_flows.functions import REQUEST_CONFIRMATION_FUNCTION_CALL_NAME
+  from google.adk.tools.function_tool import FunctionTool
+
+  # Given a tool that requires confirmation and a mock model
+  def _simple_tool_func():
+    return {'result': 'tool executed'}
+
+  mock_model = testing_utils.MockModel.create(
+      responses=[
+          types.Part.from_function_call(
+              name='_simple_tool_func',
+              args={},
+          ),
+          types.Part.from_text(text='LLM response after confirmation'),
+      ]
+  )
+
+  node_a = LlmAgent(
+      name='my_agent',
+      model=mock_model,
+      tools=[FunctionTool(func=_simple_tool_func, require_confirmation=True)],
+  )
+
+  wf = Workflow(
+      name='test_workflow_confirmation',
+      edges=[(START, node_a)],
+  )
+
+  app = App(
+      name=request.function.__name__,
+      root_agent=wf,
+  )
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  # When the workflow is started
+  user_event = testing_utils.get_user_content('start workflow')
+  events1 = await runner.run_async(user_event)
+
+  # Then it should request confirmation
+  fc_event = None
+  for e in events1:
+    if e.content and e.content.parts:
+      for p in e.content.parts:
+        if (
+            p.function_call
+            and p.function_call.name == REQUEST_CONFIRMATION_FUNCTION_CALL_NAME
+        ):
+          fc_event = e
+          break
+
+  assert fc_event is not None, 'Did not find confirmation request event'
+
+  ask_for_confirmation_function_call_id = fc_event.content.parts[
+      0
+  ].function_call.id
+  invocation_id = events1[0].invocation_id
+
+  # When the user confirms the tool call
+  user_confirmation = testing_utils.UserContent(
+      types.Part(
+          function_response=types.FunctionResponse(
+              id=ask_for_confirmation_function_call_id,
+              name=REQUEST_CONFIRMATION_FUNCTION_CALL_NAME,
+              response={'confirmed': True},
+          )
+      )
+  )
+
+  events2 = await runner.run_async(
+      new_message=user_confirmation,
+      invocation_id=invocation_id,
+  )
+
+  # Then the workflow completes with the LLM response
+  content_texts = [
+      p.text
+      for e in events2
+      if e.content and e.content.parts
+      for p in e.content.parts
+      if p.text
+  ]
+
+  assert any('LLM response after confirmation' in t for t in content_texts)
+
+
+@pytest.mark.asyncio
+async def test_workflow_pause_and_resume_auth_node(
+    request: pytest.FixtureRequest,
+):
+  """Workflow pauses on missing credentials and resumes when provided.
+
+  Setup: Workflow with a single node requiring auth.
+  Act:
+    - Run 1: Start workflow without credentials.
+    - Run 2: Provide credentials via FunctionResponse.
+  Assert:
+    - Run 1: Workflow returns adk_request_credential request.
+    - Run 2: Workflow completes and yields event with credential.
+  """
+  from fastapi.openapi.models import APIKey
+  from fastapi.openapi.models import APIKeyIn
+  from google.adk.auth.auth_credential import AuthCredential
+  from google.adk.auth.auth_credential import AuthCredentialTypes
+  from google.adk.auth.auth_tool import AuthConfig
+  from google.adk.workflow import FunctionNode
+
+  # Given a workflow with a node requiring auth
+  auth_config = AuthConfig(
+      auth_scheme=APIKey(**{'in': APIKeyIn.header, 'name': 'X-Api-Key'}),
+      credential_key='my_key',
+  )
+
+  def fetch_weather(ctx):
+    cred = ctx.get_auth_response(auth_config)
+    api_key = cred.api_key if cred else 'unknown'
+    from google.adk import Event
+
+    yield Event(message=f'authed with {api_key}')
+
+  node_a = FunctionNode(
+      func=fetch_weather, auth_config=auth_config, rerun_on_resume=True
+  )
+
+  wf = Workflow(
+      name='test_workflow_auth_node',
+      edges=[(START, node_a)],
+  )
+
+  app = App(
+      name=request.function.__name__,
+      root_agent=wf,
+  )
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  # When the workflow is started without credentials
+  events1 = await runner.run_async(testing_utils.get_user_content('start'))
+
+  # Then it should pause and request credentials
+  auth_fc_events = [
+      e
+      for e in events1
+      if e.content
+      and e.content.parts
+      and e.content.parts[0].function_call
+      and e.content.parts[0].function_call.name == 'adk_request_credential'
+  ]
+  assert len(auth_fc_events) == 1
+  auth_fc_id = auth_fc_events[0].content.parts[0].function_call.id
+  invocation_id = events1[0].invocation_id
+
+  # When the user provides the credentials
+  auth_response = AuthConfig(
+      auth_scheme=auth_config.auth_scheme,
+      credential_key=auth_config.credential_key,
+      exchanged_auth_credential=AuthCredential(
+          auth_type=AuthCredentialTypes.API_KEY,
+          api_key='secret_key',
+      ),
+  )
+
+  user_credential_response = testing_utils.UserContent(
+      types.Part(
+          function_response=types.FunctionResponse(
+              id=auth_fc_id,
+              name='adk_request_credential',
+              response=auth_response.model_dump(
+                  exclude_none=True, by_alias=True
+              ),
+          ),
+      ),
+  )
+
+  # When the workflow is resumed
+  events2 = await runner.run_async(
+      new_message=user_credential_response,
+      invocation_id=invocation_id,
+  )
+
+  # Then the workflow should resume and complete
+  content_texts = [
+      p.text
+      for e in events2
+      if e.content and e.content.parts
+      for p in e.content.parts
+      if p.text
+  ]
+  assert any('authed with secret_key' in t for t in content_texts)
+
+
+@pytest.mark.asyncio
+async def test_workflow_pause_and_resume_parent_interruption(
+    request: pytest.FixtureRequest,
+):
+  """Tests multi-agent workflow where parent produces an interruption."""
+
+  # Child agent (does nothing special)
+  child_agent = LlmAgent(
+      name='child_agent',
+      model=testing_utils.MockModel.create(
+          responses=[
+              types.Part.from_function_call(
+                  name='finish_task',
+                  args={'result': 'Child done'},
+              )
+          ]
+      ),
+      mode='task',
+  )
+
+  # Parent agent calls LRO tool first, then delegates to child
+  fc = types.Part.from_function_call(name='long_running_tool_func', args={})
+  call_child = types.Part.from_function_call(
+      name='child_agent',
+      args={'request': 'Start child task'},
+  )
+
+  parent_model = testing_utils.MockModel.create(
+      responses=[
+          fc,  # First call LRO
+          call_child,  # Then call child
+          'Parent all done',  # Finally finish
+      ]
+  )
+
+  parent_agent = LlmAgent(
+      name='parent_agent',
+      model=parent_model,
+      tools=[
+          LongRunningFunctionTool(func=long_running_tool_func),
+      ],
+      sub_agents=[child_agent],
+      mode='task',
+  )
+
+  wf = Workflow(
+      name='test_workflow_parent_hitl',
+      edges=[
+          (START, parent_agent),
+      ],
+  )
+
+  app = App(
+      name=request.function.__name__,
+      root_agent=wf,
+  )
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  # Run 1: Should pause on LRO
+  events1 = await runner.run_async(testing_utils.get_user_content('start'))
+  assert any(e.long_running_tool_ids for e in events1)
+
+  invocation_id = events1[0].invocation_id
+  fc_event = workflow_testing_utils.find_function_call_event(
+      events1, 'long_running_tool_func'
+  )
+  assert fc_event is not None
+  function_call_id = fc_event.content.parts[0].function_call.id
+
+  # Resume with tool output
+  tool_response = testing_utils.UserContent(
+      types.Part(
+          function_response=types.FunctionResponse(
+              id=function_call_id,
+              name='long_running_tool_func',
+              response={'result': 'LRO done'},
+          )
+      )
+  )
+
+  events2 = await runner.run_async(
+      new_message=tool_response,
+      invocation_id=invocation_id,
+  )
+
+  # Verify completion
+  content_texts = [
+      p.text
+      for e in events2
+      if e.content and e.content.parts
+      for p in e.content.parts
+      if p.text
+  ]
+  assert any('Parent all done' in t for t in content_texts)
+
+
+@pytest.mark.xfail(reason='Task agents cannot have sub-agents in workflow')
+@pytest.mark.asyncio
+async def test_workflow_pause_and_resume_child_interruption(
+    request: pytest.FixtureRequest,
+):
+  """Tests multi-agent workflow where child produces an interruption."""
+
+  # Child agent calls LRO tool
+  fc = types.Part.from_function_call(name='long_running_tool_func', args={})
+  child_model = testing_utils.MockModel.create(
+      responses=[
+          fc,
+          types.Part.from_function_call(
+              name='finish_task',
+              args={'result': 'Child done after tool'},
+          ),
+      ]
+  )
+
+  child_agent = LlmAgent(
+      name='child_agent',
+      model=child_model,
+      tools=[LongRunningFunctionTool(func=long_running_tool_func)],
+      mode='task',
+  )
+
+  # Parent agent delegates to child first
+  call_child = types.Part.from_function_call(
+      name='child_agent',
+      args={'request': 'Start child task'},
+  )
+  parent_model = testing_utils.MockModel.create(
+      responses=[
+          call_child,
+          call_child,
+          'Parent all done',
+      ]
+  )
+
+  parent_agent = LlmAgent(
+      name='parent_agent',
+      model=parent_model,
+      sub_agents=[child_agent],
+      mode='task',
+  )
+
+  wf = Workflow(
+      name='test_workflow_child_hitl',
+      edges=[
+          (START, parent_agent),
+      ],
+  )
+
+  app = App(
+      name=request.function.__name__,
+      root_agent=wf,
+  )
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  # Run 1: Should enter parent, then child, then child pauses on LRO!
+  events1 = await runner.run_async(testing_utils.get_user_content('start'))
+  assert any(
+      p.function_call and p.function_call.name == 'child_agent'
+      for e in events1
+      if e.content
+      for p in e.content.parts
+  )
+
+  invocation_id = events1[0].invocation_id
+  fc_event = workflow_testing_utils.find_function_call_event(
+      events1, 'long_running_tool_func'
+  )
+  assert fc_event is not None
+  function_call_id = fc_event.content.parts[0].function_call.id
+
+  # Resume with tool output
+  tool_response = testing_utils.UserContent(
+      types.Part(
+          function_response=types.FunctionResponse(
+              id=function_call_id,
+              name='long_running_tool_func',
+              response={'result': 'LRO done'},
+          )
+      )
+  )
+
+  events2 = await runner.run_async(
+      new_message=tool_response,
+      invocation_id=invocation_id,
+  )
+
+  # Verify completion
+  content_texts = [
+      p.text
+      for e in events2
+      if e.content and e.content.parts
+      for p in e.content.parts
+      if p.text
+  ]
+  assert any('Parent all done' in t for t in content_texts)
+
+
+def _append_function_response(
+    session, invocation_id, branch, fc_id, func_name, response
+):
+  """Helper to append a FunctionResponse event to a session."""
+  session.events.append(
+      Event(
+          invocation_id=invocation_id,
+          author='user',
+          branch=branch,
+          content=types.Content(
+              role='user',
+              parts=[
+                  types.Part(
+                      function_response=types.FunctionResponse(
+                          id=fc_id,
+                          name=func_name,
+                          response=response,
+                      )
+                  )
+              ],
+          ),
+      )
+  )
+
+
+@pytest.mark.asyncio
+async def test_workflow_resume_inputs_fallback_branch(monkeypatch):
+  """Resume inputs find function name in a different branch.
+
+  Setup: Session contains a FunctionCall in branch_A.
+  Act: Run the wrapper with resume_inputs for that FunctionCall in branch_B.
+  Assert: The wrapper successfully finds the function name and completes.
+  """
+
+  # Arrange
+  mock_model = testing_utils.MockModel.create(
+      responses=[types.Part.from_text(text='I am done')]
+  )
+  agent = LlmAgent(name='test_agent', model=mock_model)
+
+  # Create a dummy context and session
+
+  session = Session(id='test_session', appName='test_app', userId='test_user')
+  # Add an event with function call in branch 'branch_A'
+  session.events.append(
+      Event(
+          invocation_id='test_inv',
+          branch='branch_A',
+          content=types.Content(
+              parts=[
+                  types.Part(
+                      function_call=types.FunctionCall(
+                          id='fc_123',
+                          name='my_target_func',
+                          args={},
+                      )
+                  )
+              ]
+          ),
+      )
+  )
+
+  # Create invocation context with branch 'branch_B'
+
+  session_service = InMemorySessionService()
+
+  ic = InvocationContext(
+      session=session,
+      branch='branch_B',
+      session_service=session_service,
+      invocation_id='test_inv',
+      run_config=RunConfig(),
+      agent=agent,
+  )
+
+  # Create context with resume_inputs
+  ctx = Context(
+      ic,
+      node_path='test_agent',
+      run_id='1',
+      resume_inputs={'fc_123': {'result': 'ok'}},
+  )
+
+  # Mock prepare functions
+  class DummyAgentCtx:
+
+    def __init__(self, ic):
+      self._ic = ic
+
+    def get_invocation_context(self):
+      return self._ic
+
+  from google.adk.workflow import _llm_agent_wrapper
+
+  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
+  )
+
+  # Simulate Runner adding the event to the correct branch!
+  _append_function_response(
+      session,
+      invocation_id='test_inv',
+      branch='branch_A',
+      fc_id='fc_123',
+      func_name='my_target_func',
+      response={'result': 'ok'},
+  )
+
+  # Act - Run the function directly
+  from google.adk.workflow._llm_agent_wrapper import run_llm_agent_as_node
+
+  gen = run_llm_agent_as_node(agent, ctx=ctx, node_input='start')
+  try:
+    await gen.__anext__()
+  except StopAsyncIteration:
+    pass
+
+  # Assert - Verify that the event is there with correct branch
+  # Initial events had 1 item + 1 simulated by Runner = 2!
+  assert len(session.events) == 2
+  event = session.events[-1]  # The newly injected event!
+  assert (
+      event.branch == 'branch_A'
+  )  # Injected in the branch where call was made!
+  assert event.content.parts[0].function_response.name == 'my_target_func'
+
+
+@pytest.mark.asyncio
+async def test_workflow_resume_inputs_multiple_branches(monkeypatch):
+  """Resume inputs handle multiple items targeting different branches.
+
+  Setup: Session contains FunctionCalls in branch_A and branch_B.
+  Act: Run the wrapper with resume_inputs for both in branch_C.
+  Assert: The wrapper successfully finds function names for both.
+  """
+
+  # Arrange
+  mock_model = testing_utils.MockModel.create(
+      responses=[types.Part.from_text(text='I am done')]
+  )
+  agent = LlmAgent(name='test_agent', model=mock_model)
+
+  session = Session(id='test_session', appName='test_app', userId='test_user')
+
+  # Add event 1 in branch_A
+  session.events.append(
+      Event(
+          invocation_id='test_inv',
+          branch='branch_A',
+          content=types.Content(
+              parts=[
+                  types.Part(
+                      function_call=types.FunctionCall(
+                          id='fc_A',
+                          name='func_A',
+                          args={},
+                      )
+                  )
+              ]
+          ),
+      )
+  )
+
+  # Add event 2 in branch_B
+  session.events.append(
+      Event(
+          invocation_id='test_inv',
+          branch='branch_B',
+          content=types.Content(
+              parts=[
+                  types.Part(
+                      function_call=types.FunctionCall(
+                          id='fc_B',
+                          name='func_B',
+                          args={},
+                      )
+                  )
+              ]
+          ),
+      )
+  )
+
+  session_service = InMemorySessionService()
+
+  # Current branch is branch_C
+  ic = InvocationContext(
+      session=session,
+      branch='branch_C',
+      session_service=session_service,
+      invocation_id='test_inv',
+      run_config=RunConfig(),
+      agent=agent,
+  )
+
+  ctx = Context(
+      ic,
+      node_path='test_agent',
+      run_id='1',
+      resume_inputs={'fc_A': {'result': 'ok_A'}, 'fc_B': {'result': 'ok_B'}},
+  )
+
+  class DummyAgentCtx:
+
+    def __init__(self, ic):
+      self._ic = ic
+
+    def get_invocation_context(self):
+      return self._ic
+
+  from google.adk.workflow import _llm_agent_wrapper
+
+  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
+  )
+
+  # Simulate Runner adding the events to the correct branches!
+  _append_function_response(
+      session,
+      invocation_id='test_inv',
+      branch='branch_A',
+      fc_id='fc_A',
+      func_name='func_A',
+      response={'result': 'ok_A'},
+  )
+  _append_function_response(
+      session,
+      invocation_id='test_inv',
+      branch='branch_B',
+      fc_id='fc_B',
+      func_name='func_B',
+      response={'result': 'ok_B'},
+  )
+
+  # Act - Run the function directly
+  from google.adk.workflow._llm_agent_wrapper import run_llm_agent_as_node
+
+  gen = run_llm_agent_as_node(agent, ctx=ctx, node_input='start')
+  try:
+    await gen.__anext__()
+  except StopAsyncIteration:
+    pass
+
+  # Assert - Verify that the events are there with correct branches
+  # Initial events had 2 items + 2 simulated by Runner = 4!
+  assert len(session.events) == 4
+
+  # Check both exist in the injected events
+  injected_events = session.events[2:]
+  branches = [e.branch for e in injected_events]
+  names = [e.content.parts[0].function_response.name for e in injected_events]
+
+  assert 'branch_A' in branches
+  assert 'branch_B' in branches
+  assert 'func_A' in names
+  assert 'func_B' in names
+
+
+@pytest.mark.asyncio
+async def test_workflow_task_mode_plain_text_resume_auto_routing(
+    request: pytest.FixtureRequest,
+):
+  """Tests that task mode agent can pause for text and resume with text without explicit invocation_id."""
+
+  # LLM behavior:
+  # Round 1: Outputs text asking for info.
+  # Round 2: After user replies with text, calls finish_task with result.
+  mock_model = testing_utils.MockModel.create(
+      responses=[
+          types.Part.from_text(text='Please provide the secret code:'),
+          types.Part.from_function_call(
+              name='finish_task',
+              args={'result': 'Success with code'},
+          ),
+      ]
+  )
+
+  node_a = LlmAgent(
+      name='my_task_agent',
+      model=mock_model,
+      mode='task',
+  )
+
+  wf = Workflow(
+      name='test_workflow_plain_text_resume',
+      edges=[
+          (START, node_a),
+      ],
+  )
+
+  app = App(
+      name=request.function.__name__,
+      root_agent=wf,
+  )
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  # 1. First run: should yield the prompt and pause.
+  events1 = await runner.run_async('start workflow')
+
+  # Verify it yielded the prompt
+  texts1 = [
+      p.text
+      for e in events1
+      if e.content and e.content.parts
+      for p in e.content.parts
+      if p.text
+  ]
+  assert 'Please provide the secret code:' in texts1
+
+  # 2. Second run: resume with plain text, NOT passing invocation_id!
+  # It should automatically resolve invocation_id and isolation_scope.
+  events2 = await runner.run_async('my_secret_code_123')
+
+  # Verify completion
+  # The last event should have output set from finish_task args
+  assert any(e.output == {'result': 'Success with code'} for e in events2)

tests/unittests/workflow/test_workflow_node.py

@@ -0,0 +1,450 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for @node decorator and behavior."""
+
+from __future__ import annotations
+
+from unittest import mock
+
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.apps import App
+from google.adk.runners import Runner
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.tools.base_tool import BaseTool
+from google.adk.workflow import FunctionNode
+from google.adk.workflow import START
+from google.adk.workflow._base_node import BaseNode
+from google.adk.workflow._node import node
+from google.adk.workflow._node import Node
+from google.adk.workflow._parallel_worker import _ParallelWorker as ParallelWorker
+from google.adk.workflow._retry_config import RetryConfig
+from google.adk.workflow._tool_node import _ToolNode as ToolNode
+from google.adk.workflow._workflow import Workflow
+from google.genai import types
+import pytest
+
+from .. import testing_utils
+
+ANY = mock.ANY
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+async def _run_workflow(wf, message="start"):
+  """Run a Workflow through Runner, return collected events."""
+  ss = InMemorySessionService()
+  runner = Runner(app_name="test", node=wf, session_service=ss)
+  session = await ss.create_session(app_name="test", user_id="u")
+  msg = types.Content(parts=[types.Part(text=message)], role="user")
+  events = []
+  async for event in runner.run_async(
+      user_id="u", session_id=session.id, new_message=msg
+  ):
+    events.append(event)
+  return events, ss, session
+
+
+def _output_by_node(events):
+  """Extract (node_name_from_path, output) for child node events."""
+  results = []
+  for e in events:
+    if e.output is not None and e.node_info.path and "/" in e.node_info.path:
+      node_name = e.node_info.path.rsplit("/", 1)[-1]
+      if "@" in node_name:
+        node_name = node_name.rsplit("@", 1)[0]
+      results.append((node_name, e.output))
+  return results
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_node_decorator():
+  """Tests that @node decorator can wrap a function and override its name."""
+
+  @node(name="decorated_node")
+  def my_func():
+    return "Hello from decorated_func"
+
+  assert my_func.name == "decorated_node"
+
+  wf = Workflow(
+      name="test_agent",
+      edges=[
+          (START, my_func),
+      ],
+  )
+  events, _, _ = await _run_workflow(wf)
+
+  by_node = _output_by_node(events)
+  assert ("decorated_node", "Hello from decorated_func") in by_node
+
+
+def test_node_parallel_worker_instance():
+  """Tests that node() can wrap a node in ParallelWorker."""
+
+  @node(parallel_worker=True)
+  def my_func(node_input):
+    return node_input
+
+  assert isinstance(my_func, ParallelWorker)
+  assert my_func.name == "my_func"
+
+  def other_func(x):
+    return x
+
+  parallel_node = node(other_func, parallel_worker=True)
+  assert isinstance(parallel_node, ParallelWorker)
+  assert parallel_node.name == "other_func"
+
+
+@pytest.mark.asyncio
+async def test_node_parallel_worker_execution():
+  """Tests that a node with parallel_worker=True correctly processes inputs."""
+
+  @node(parallel_worker=True)
+  async def my_func(node_input):
+    return node_input * 2
+
+  async def producer_func() -> list[int]:
+    return [1, 2, 3]
+
+  wf = Workflow(
+      name="test_agent",
+      edges=[
+          (START, producer_func),
+          (producer_func, my_func),
+      ],
+  )
+  events, _, _ = await _run_workflow(wf)
+
+  by_node = _output_by_node(events)
+  assert ("producer_func", [1, 2, 3]) in by_node
+  assert ("my_func", [2, 4, 6]) in by_node
+
+
+def test_node_decorator_rerun_on_resume():
+  """Tests that @node decorator can override rerun_on_resume."""
+
+  @node(name="decorated_node", rerun_on_resume=True)
+  def my_func():
+    return "Hello from decorated_func"
+
+  assert isinstance(my_func, FunctionNode)
+  assert my_func.rerun_on_resume
+
+  @node()
+  def my_func2():
+    return "Hello from decorated_func2"
+
+  assert isinstance(my_func2, FunctionNode)
+  assert not my_func2.rerun_on_resume
+
+
+def test_node_decorator_parameter_binding():
+  """Tests that @node decorator can configure parameter_binding."""
+
+  @node(parameter_binding="node_input")
+  def my_func(foo: str):
+    return f"Hello {foo}"
+
+  assert isinstance(my_func, FunctionNode)
+  assert my_func.parameter_binding == "node_input"
+
+
+@pytest.mark.asyncio
+async def test_node_decorator_parameter_binding_execution():
+  """Tests execution of a node with parameter_binding='node_input'."""
+
+  async def producer_func() -> dict[str, Any]:
+    return {"foo": "hello", "bar": 42}
+
+  @node(parameter_binding="node_input")
+  def my_func(foo: str, bar: int):
+    return f"{foo}:{bar}"
+
+  wf = Workflow(
+      name="test_agent",
+      edges=[
+          (START, producer_func),
+          (producer_func, my_func),
+      ],
+  )
+  events, _, _ = await _run_workflow(wf)
+
+  by_node = _output_by_node(events)
+  assert ("my_func", "hello:42") in by_node
+
+
+def test_node_function_with_base_node():
+  """Tests that node() function returns a copied node when given a BaseNode."""
+
+  @node(name="original")
+  def original():
+    pass
+
+  wrapped = node(original, name="overridden", rerun_on_resume=True)
+
+  assert isinstance(wrapped, FunctionNode)
+  assert wrapped is not original
+  assert wrapped.name == "overridden"
+  assert wrapped.rerun_on_resume
+
+
+class MyTool(BaseTool):
+  name = "tool"
+  description = "desc"
+
+  async def _run_async_impl(self):
+    return "done"
+
+
+def test_node_no_unnecessary_wrap():
+  """Tests that node() does not wrap LlmAgent, Agent, Tool, or func in OverridingNode."""
+
+  llm_agent = LlmAgent(name="llm")
+  llm_node = node(llm_agent, name="overridden_llm")
+
+  assert isinstance(llm_node, LlmAgent)
+  assert llm_node.name == "overridden_llm"
+  assert llm_node.mode == "single_turn"
+
+  agent = BaseAgent(name="agent")
+  agent_node_inst = node(agent, name="overridden_agent", rerun_on_resume=True)
+  assert isinstance(agent_node_inst, BaseAgent)
+  assert agent_node_inst.name == "overridden_agent"
+  assert agent_node_inst.rerun_on_resume
+
+  tool_inst = MyTool(name="tool", description="desc")
+  t_node = node(tool_inst, name="overridden_tool")
+  assert isinstance(t_node, ToolNode)
+  assert t_node.name == "overridden_tool"
+
+  def my_func():
+    pass
+
+  f_node = node(my_func, name="overridden_func", rerun_on_resume=True)
+  assert isinstance(f_node, FunctionNode)
+  assert f_node.name == "overridden_func"
+  assert f_node.rerun_on_resume
+
+
+class StatefulTool(BaseTool):
+  """A tool that modifies state via tool_context."""
+
+  async def run_async(self, *, args, tool_context):
+    tool_context.state["tool_key"] = "tool_value"
+    tool_context.state["tool_count"] = 10
+    return {"status": "ok"}
+
+
+from .workflow_testing_utils import simplify_events_with_node
+
+
+@pytest.mark.asyncio
+async def test_tool_node_state_delta():
+  """Tests that state set via tool_context.state in ToolNode is persisted."""
+
+  tool_node = ToolNode(
+      tool=StatefulTool(name="stateful_tool", description="Sets state values"),
+  )
+
+  def read_state(tool_key: str, tool_count: int) -> str:
+    return f"tool_key={tool_key}, tool_count={tool_count}"
+
+  def start_node():
+    return {}
+
+  wf = Workflow(
+      name="test_tool_node_state_delta",
+      edges=[
+          (START, start_node),
+          (start_node, tool_node),
+          (tool_node, read_state),
+      ],
+  )
+
+  events, _, _ = await _run_workflow(wf)
+
+  simplified = simplify_events_with_node(
+      events, include_workflow_output=True, include_state_delta=True
+  )
+
+  assert (
+      "test_tool_node_state_delta@1/stateful_tool@1",
+      {"output": {"status": "ok"}},
+  ) in [(e[0], {"output": e[1].get("output")}) for e in simplified]
+
+  assert (
+      "test_tool_node_state_delta@1/read_state@1",
+      {
+          "output": "tool_key=tool_value, tool_count=10",
+      },
+  ) in [(e[0], {"output": e[1].get("output")}) for e in simplified]
+
+
+class _CustomNode(Node):
+  custom_val: str = "hello"
+  rerun_on_resume: bool = True
+
+  async def run_node_impl(self, *, ctx, node_input):
+    yield f"subclass: {self.custom_val} -> {node_input}"
+
+
+def test_node_subclassing_model_copy_preserves_identity():
+  """Tests that Node.model_copy preserves the subclass class identity."""
+  node_inst = _CustomNode(
+      name="subclass", parallel_worker=True, custom_val="barrier"
+  )
+  assert node_inst.parallel_worker is True
+
+  cloned = node_inst.model_copy()
+  assert isinstance(cloned, _CustomNode)
+  assert cloned.custom_val == "barrier"
+  assert cloned.parallel_worker is True
+  assert isinstance(cloned._inner_node, ParallelWorker)
+  # Confirm inner node wraps a clone of _CustomNode preserving identity!
+  assert isinstance(cloned._inner_node._node, _CustomNode)
+  assert cloned._inner_node._node.custom_val == "barrier"
+  assert cloned._inner_node._node.parallel_worker is False
+
+
+@pytest.mark.asyncio
+async def test_node_subclassing_execution_with_parallel_worker():
+  """Tests that a subclassed Node with parallel_worker=True executes successfully."""
+  subclass_node = _CustomNode(
+      name="subclass", parallel_worker=True, custom_val="workflow"
+  )
+
+  async def producer():
+    return ["input1", "input2"]
+
+  wf = Workflow(
+      name="test_agent",
+      edges=[
+          (START, producer),
+          (producer, subclass_node),
+      ],
+  )
+
+  events, _, _ = await _run_workflow(wf)
+  by_node = _output_by_node(events)
+
+  assert ("producer", ["input1", "input2"]) in by_node
+  assert (
+      "subclass",
+      ["subclass: workflow -> input1", "subclass: workflow -> input2"],
+  ) in by_node
+
+
+def test_node_decorator_parallel_worker_max_parallel_workers():
+  """Tests that node() correctly sets max_parallel_workers on ParallelWorker."""
+
+  @node(parallel_worker=True, max_parallel_workers=3)
+  def my_func(node_input):
+    return node_input
+
+  assert isinstance(my_func, ParallelWorker)
+  assert my_func.max_parallel_workers == 3
+
+
+def test_node_decorator_invalid_max_parallel_workers():
+  """Tests that node() raises ValueError if max_parallel_workers is set without parallel_worker."""
+  with pytest.raises(
+      ValueError,
+      match="max_parallel_workers can only be set when parallel_worker is True",
+  ):
+
+    @node(parallel_worker=False, max_parallel_workers=3)
+    def my_func(node_input):
+      return node_input
+
+
+def test_node_subclass_invalid_max_parallel_workers():
+  """Tests that Node subclass raises ValidationError if max_parallel_workers is set without parallel_worker."""
+  from pydantic import ValidationError
+
+  with pytest.raises(ValidationError) as exc_info:
+    _CustomNode(name="subclass", parallel_worker=False, max_parallel_workers=3)
+
+  assert (
+      "max_parallel_workers can only be set when parallel_worker is True"
+      in str(exc_info.value)
+  )
+
+
+def test_node_subclassing_model_copy_preserves_max_parallel_workers():
+  """Tests that Node.model_copy preserves max_parallel_workers."""
+  node_inst = _CustomNode(
+      name="subclass",
+      parallel_worker=True,
+      max_parallel_workers=5,
+      custom_val="barrier",
+  )
+  assert node_inst.parallel_worker is True
+  assert node_inst.max_parallel_workers == 5
+  assert node_inst._inner_node.max_parallel_workers == 5
+
+  cloned = node_inst.model_copy()
+  assert isinstance(cloned, _CustomNode)
+  assert cloned.parallel_worker is True
+  assert cloned.max_parallel_workers == 5
+  assert isinstance(cloned._inner_node, ParallelWorker)
+  assert cloned._inner_node.max_parallel_workers == 5
+  assert cloned._inner_node._node.parallel_worker is False
+
+
+def test_node_decorator_invalid_max_parallel_workers_less_than_one():
+  """Tests that node() raises ValueError if max_parallel_workers is less than 1."""
+  with pytest.raises(
+      ValueError,
+      match="max_parallel_workers must be greater than or equal to 1",
+  ):
+
+    @node(parallel_worker=True, max_parallel_workers=0)
+    def my_func(node_input):
+      return node_input
+
+
+def test_node_subclass_invalid_max_parallel_workers_less_than_one():
+  """Tests that Node subclass raises ValidationError if max_parallel_workers is less than 1."""
+  from pydantic import ValidationError
+
+  with pytest.raises(ValidationError) as exc_info:
+    _CustomNode(name="subclass", parallel_worker=True, max_parallel_workers=0)
+
+  assert "max_parallel_workers must be greater than or equal to 1" in str(
+      exc_info.value
+  )
+
+
+def test_parallel_worker_invalid_max_parallel_workers_less_than_one():
+  """Tests that ParallelWorker constructor raises ValueError if max_parallel_workers is less than 1."""
+
+  def dummy_func(x):
+    return x
+
+  with pytest.raises(
+      ValueError,
+      match="max_parallel_workers must be greater than or equal to 1",
+  ):
+    ParallelWorker(node=dummy_func, max_parallel_workers=0)

tests/unittests/workflow/test_workflow_node_timeout.py

@@ -0,0 +1,204 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for node timeout behavior."""
+
+from __future__ import annotations
+
+import asyncio
+
+from google.adk.apps import App
+from google.adk.runners import Runner
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.workflow import START
+from google.adk.workflow._errors import NodeTimeoutError
+from google.adk.workflow._node import node
+from google.adk.workflow._retry_config import RetryConfig
+from google.adk.workflow._workflow import Workflow
+from google.genai import types
+import pytest
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+async def _run_workflow(wf, message='start'):
+  """Run a Workflow through Runner, return collected events."""
+  ss = InMemorySessionService()
+  runner = Runner(app_name='test', node=wf, session_service=ss)
+
+  session = await ss.create_session(app_name='test', user_id='u')
+  msg = types.Content(parts=[types.Part(text=message)], role='user')
+  events = []
+  async for event in runner.run_async(
+      user_id='u', session_id=session.id, new_message=msg
+  ):
+    events.append(event)
+  return events, ss, session
+
+
+def _output_by_node(events):
+  """Extract (node_name_from_path, output) for child node events."""
+  results = []
+  for e in events:
+    if e.output is not None and e.node_info.path and '/' in e.node_info.path:
+      node_name = e.node_info.path.rsplit('/', 1)[-1]
+      if '@' in node_name:
+        node_name = node_name.rsplit('@', 1)[0]
+      results.append((node_name, e.output))
+  return results
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_node_completes_within_timeout():
+  """A node that finishes before the timeout should succeed normally."""
+
+  @node(timeout=5.0)
+  async def my_slow_node():
+    await asyncio.sleep(0.01)
+    return 'done'
+
+  wf = Workflow(
+      name='test_workflow',
+      edges=[
+          (START, my_slow_node),
+      ],
+  )
+  events, _, _ = await _run_workflow(wf)
+  by_node = _output_by_node(events)
+
+  assert ('my_slow_node', 'done') in by_node
+
+
+@pytest.mark.asyncio
+async def test_node_exceeds_timeout():
+  """A node that exceeds its timeout should fail."""
+
+  from google.adk.workflow import FunctionNode
+
+  async def raw_slow_func():
+    await asyncio.sleep(1.0)
+    return 'done'
+
+  my_too_slow_node = FunctionNode(
+      name='my_too_slow_node', func=raw_slow_func, timeout=0.05
+  )
+
+  wf = Workflow(
+      name='test_workflow',
+      edges=[
+          (START, my_too_slow_node),
+      ],
+  )
+  with pytest.raises(NodeTimeoutError) as exc_info:
+    await _run_workflow(wf)
+
+  assert 'my_too_slow_node' in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+async def test_node_no_timeout():
+  """A node with timeout=None should run without any time limit."""
+
+  @node(timeout=None)
+  async def my_no_timeout_node():
+    await asyncio.sleep(0.01)
+    return 'done'
+
+  wf = Workflow(
+      name='test_workflow',
+      edges=[
+          (START, my_no_timeout_node),
+      ],
+  )
+  events, _, _ = await _run_workflow(wf)
+  by_node = _output_by_node(events)
+
+  assert ('my_no_timeout_node', 'done') in by_node
+
+
+@pytest.mark.asyncio
+async def test_node_timeout_with_retry():
+  """A timed-out node should be retried if retry_config is set."""
+  run_count = 0
+
+  @node(
+      timeout=0.05,
+      retry_config=RetryConfig(max_attempts=3, initial_delay=0.0, jitter=0.0),
+  )
+  async def node_a():
+    nonlocal run_count
+    run_count += 1
+    if run_count == 1:
+      await asyncio.sleep(1.0)
+    return 'success'
+
+  wf = Workflow(
+      name='test_workflow',
+      edges=[
+          (START, node_a),
+      ],
+  )
+  events, _, _ = await _run_workflow(wf)
+
+  by_node = _output_by_node(events)
+  # Verify that the final result was successfully obtained.
+  assert ('node_a', 'success') in by_node
+
+  # Verify that the node was actually executed more than once (i.e., retried).
+  assert run_count == 2, f'Expected run_count == 2, got {run_count}'
+
+
+@pytest.mark.asyncio
+async def test_nested_workflow_timeout():
+  """A nested workflow that exceeds its timeout in the outer workflow should fail.
+
+  Setup: outer_wf -> inner_wf -> slow_node. inner_wf has timeout=0.05.
+  Act: Run the outer workflow.
+  Assert: Execution raises NodeTimeoutError referencing inner_wf.
+  """
+  import sys
+
+  if sys.version_info < (3, 11):
+    pytest.skip('asyncio.timeout requires Python 3.11+')
+
+  # Given an outer workflow containing a slow inner workflow with a timeout
+  @node()
+  async def slow_node():
+    await asyncio.sleep(1.0)
+    return 'done'
+
+  inner_wf = Workflow(
+      name='inner_wf',
+      edges=[(START, slow_node)],
+      timeout=0.05,
+  )
+
+  outer_wf = Workflow(
+      name='outer_wf',
+      edges=[(START, inner_wf)],
+  )
+
+  # When the outer workflow is executed
+  # Then it should raise NodeTimeoutError referencing the inner workflow
+  with pytest.raises(NodeTimeoutError) as exc_info:
+    await _run_workflow(outer_wf)
+
+  assert 'inner_wf' in str(exc_info.value)

tests/unittests/workflow/test_workflow_output_deduplication.py

@@ -0,0 +1,161 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for nested workflow output event deduplication.
+
+Verifies that nested workflows produce only leaf terminal events and
+resolve output via terminal path resolution from the graph structure.
+"""
+
+from typing import Any
+
+from google.adk.apps.app import App
+from google.adk.events.event import Event
+from google.adk.workflow._workflow import Workflow
+import pytest
+
+from . import testing_utils
+
+
+def _is_checkpoint(event: Event) -> bool:
+  """Returns True if event is an agent state checkpoint or end_of_agent."""
+  if event.actions.agent_state is not None:
+    return True
+  if event.actions.end_of_agent:
+    return True
+  return False
+
+
+def _output_events(events: list[Event]) -> list[Event]:
+  """Returns only events with output data (not checkpoint/state)."""
+  return [e for e in events if not _is_checkpoint(e) and e.output is not None]
+
+
+async def test_two_level_nesting_deduplicates(
+    request,
+):
+  """Two-level nesting emits only the leaf's output event, not finalize events.
+
+  Setup: outer → inner → leaf.
+  Assert: exactly 1 output event from 'outer/inner/leaf', no
+    duplicate finalize events from inner or outer.
+  """
+
+  async def leaf(node_input: Any):
+    return 'leaf_data'
+
+  inner = Workflow(name='inner', edges=[('START', leaf)])
+  outer = Workflow(name='outer', edges=[('START', inner)])
+
+  app = App(name=request.function.__name__, root_agent=outer)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('hi'))
+  out_events = _output_events(events)
+
+  assert len(out_events) == 1
+  assert out_events[0].output == 'leaf_data'
+  assert out_events[0].node_info.path == 'outer@1/inner@1/leaf@1'
+
+
+async def test_nested_with_output_schema_validates_at_read_time(
+    request,
+):
+  """Nested workflow with output_schema validates without emitting extra events.
+
+  Setup: outer → inner(output_schema=str) → leaf.
+  Assert: 1 output event with validated data, no extra finalize event.
+  """
+
+  async def leaf(node_input: Any):
+    return 'raw_data'
+
+  inner = Workflow(
+      name='inner',
+      edges=[('START', leaf)],
+      output_schema=str,
+  )
+  outer = Workflow(name='outer', edges=[('START', inner)])
+
+  app = App(name=request.function.__name__, root_agent=outer)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('hi'))
+  out_events = _output_events(events)
+
+  assert len(out_events) == 1
+  assert out_events[0].output == 'raw_data'
+
+
+async def test_multiple_terminals_in_nested_workflow_raises(
+    request,
+):
+  """Fan-out with no join raises ValueError for multiple terminal outputs.
+
+  Setup: outer → inner → (branch_a, branch_b).
+  Assert: ValueError because inner has two terminal nodes producing output.
+  """
+
+  async def branch_a(node_input: Any):
+    return 'a_out'
+
+  async def branch_b(node_input: Any):
+    return 'b_out'
+
+  inner = Workflow(
+      name='inner',
+      edges=[('START', (branch_a, branch_b))],
+  )
+  outer = Workflow(name='outer', edges=[('START', inner)])
+
+  app = App(name=request.function.__name__, root_agent=outer)
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  with pytest.raises(ValueError, match='multiple terminal nodes'):
+    await runner.run_async(testing_utils.get_user_content('hi'))
+
+
+async def test_non_terminal_output_not_exposed_as_workflow_output(
+    request,
+):
+  """Downstream node receives terminal output, not intermediate node output.
+
+  Setup: outer → inner(step_a → step_b) → consume.
+  Assert: consume receives step_b's 'final' (terminal), not step_a's
+    'intermediate'.
+  """
+
+  async def step_a(node_input: Any):
+    return 'intermediate'
+
+  async def step_b(node_input: str):
+    return 'final'
+
+  inner = Workflow(name='inner', edges=[('START', step_a, step_b)])
+
+  async def consume(node_input: str):
+    return f'got: {node_input}'
+
+  outer = Workflow(name='outer', edges=[('START', inner, consume)])
+
+  app = App(name=request.function.__name__, root_agent=outer)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('hi'))
+  out_events = _output_events(events)
+
+  consume_events = [
+      e
+      for e in out_events
+      if e.node_info.name and e.node_info.name.split('@')[0] == 'consume'
+  ]
+  assert len(consume_events) == 1
+  assert consume_events[0].output == 'got: final'

tests/unittests/workflow/test_workflow_routes.py

@@ -0,0 +1,685 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Testings for the Workflow routes."""
+
+from typing import Any
+
+from google.adk.agents.context import Context
+from google.adk.apps.app import App
+from google.adk.workflow import Edge
+from google.adk.workflow import START
+from google.adk.workflow._graph import DEFAULT_ROUTE
+from google.adk.workflow._graph import Graph
+from google.adk.workflow._join_node import JoinNode
+from google.adk.workflow._workflow import Workflow
+import pytest
+
+from .. import testing_utils
+from .workflow_testing_utils import simplify_events_with_node
+from .workflow_testing_utils import TestingNode
+
+
+@pytest.mark.asyncio
+async def test_run_async_with_edge_routes(request: pytest.FixtureRequest):
+  route_holder = {'route': 'route_b'}
+
+  def dynamic_router(_ctx: Context, _node_input: Any):
+    return route_holder['route']
+
+  node_a = TestingNode(name='NodeA', output='A', route=dynamic_router)
+  node_b = TestingNode(name='NodeB', output='B')
+  node_c = TestingNode(name='NodeC', output='C')
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(
+              from_node=node_a,
+              to_node=node_b,
+              route='route_b',
+          ),
+          Edge(
+              from_node=node_a,
+              to_node=node_c,
+              route='route_c',
+          ),
+      ],
+  )
+  agent = Workflow(
+      name='test_workflow_agent',
+      graph=graph,
+  )
+
+  # Test case for route_b
+  route_holder['route'] = 'route_b'
+  app = App(name=request.function.__name__ + '_b', root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events_b = await runner.run_async(testing_utils.get_user_content('start'))
+  assert simplify_events_with_node(events_b) == [
+      ('test_workflow_agent@1/NodeA@1', {'output': 'A'}),
+      ('test_workflow_agent@1/NodeB@1', {'output': 'B'}),
+  ]
+
+  # Test case for route_c
+  route_holder['route'] = 'route_c'
+  app_c = App(name=request.function.__name__ + '_c', root_agent=agent)
+  runner_c = testing_utils.InMemoryRunner(app=app_c)
+  events_c = await runner_c.run_async(testing_utils.get_user_content('start'))
+  assert simplify_events_with_node(events_c) == [
+      ('test_workflow_agent@1/NodeA@1', {'output': 'A'}),
+      ('test_workflow_agent@1/NodeC@1', {'output': 'C'}),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_output_route_int(request: pytest.FixtureRequest):
+  node_a = TestingNode(name='NodeA', route=1)
+  node_b = TestingNode(name='NodeB', output='B')
+  node_c = TestingNode(name='NodeC', output='C')
+  agent = Workflow(
+      name='test_workflow_agent_route_int',
+      edges=[
+          (START, node_a),
+          (node_a, {1: node_b, 2: node_c}),
+      ],
+  )
+  app = App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  assert simplify_events_with_node(events) == [
+      (
+          'test_workflow_agent_route_int@1/NodeA@1',
+          {'output': None},
+      ),
+      (
+          'test_workflow_agent_route_int@1/NodeB@1',
+          {'output': 'B'},
+      ),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_output_route_bool(request: pytest.FixtureRequest):
+  node_a = TestingNode(name='NodeA', route=True)
+  node_b = TestingNode(name='NodeB', output='B')
+  node_c = TestingNode(name='NodeC', output='C')
+  agent = Workflow(
+      name='test_workflow_agent_route_bool',
+      edges=[
+          (START, node_a),
+          (node_a, {True: node_b, False: node_c}),
+      ],
+  )
+  app = App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  assert simplify_events_with_node(events) == [
+      (
+          'test_workflow_agent_route_bool@1/NodeA@1',
+          {'output': None},
+      ),
+      (
+          'test_workflow_agent_route_bool@1/NodeB@1',
+          {'output': 'B'},
+      ),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_wait_for_output_with_route_only_completes_successfully(
+    request: pytest.FixtureRequest,
+):
+  """A node with wait_for_output=True that yields only a route should complete and continue the workflow."""
+  node_a = TestingNode(name='NodeA', route='go_next', wait_for_output=True)
+  node_b = TestingNode(name='NodeB', output='B_done')
+
+  agent = Workflow(
+      name='test_wait_for_output_route',
+      edges=[
+          (START, node_a),
+          (node_a, {'go_next': node_b}),
+      ],
+  )
+  app = App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  # NodeA should yield no output, and NodeB should yield its output.
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  assert simplify_events_with_node(events) == [
+      (
+          'test_wait_for_output_route@1/NodeA@1',
+          {'output': None},
+      ),
+      (
+          'test_wait_for_output_route@1/NodeB@1',
+          {'output': 'B_done'},
+      ),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_output_route_no_data(request: pytest.FixtureRequest):
+  node_a = TestingNode(name='NodeA', route='route_b')
+  node_b = TestingNode(name='NodeB', output='B')
+  node_c = TestingNode(name='NodeC', output='C')
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(
+              from_node=node_a,
+              to_node=node_b,
+              route='route_b',
+          ),
+          Edge(
+              from_node=node_a,
+              to_node=node_c,
+              route='route_c',
+          ),
+      ],
+  )
+  agent = Workflow(
+      name='test_workflow_agent_route_no_data',
+      graph=graph,
+  )
+  app = App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  assert simplify_events_with_node(events) == [
+      (
+          'test_workflow_agent_route_no_data@1/NodeA@1',
+          {'output': None},
+      ),
+      (
+          'test_workflow_agent_route_no_data@1/NodeB@1',
+          {'output': 'B'},
+      ),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_run_async_with_list_of_routes(request: pytest.FixtureRequest):
+  node_a = TestingNode(name='NodeA', output='A', route=['route_b', 'route_c'])
+  node_b = TestingNode(name='NodeB', output='B')
+  node_c = TestingNode(name='NodeC', output='C')
+  node_d = TestingNode(name='NodeD', output='D')
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(
+              from_node=node_a,
+              to_node=node_b,
+              route='route_b',
+          ),
+          Edge(
+              from_node=node_a,
+              to_node=node_c,
+              route='route_c',
+          ),
+          Edge(
+              from_node=node_a,
+              to_node=node_d,
+              route='route_d',
+          ),
+      ],
+  )
+  agent = Workflow(
+      name='test_workflow_agent_list_routes',
+      graph=graph,
+  )
+
+  app = App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  simplified_events = simplify_events_with_node(events)
+
+  assert len(simplified_events) == 3
+  assert simplified_events[0] == (
+      'test_workflow_agent_list_routes@1/NodeA@1',
+      {'output': 'A'},
+  )
+
+  # Check that the other two events are from NodeB and NodeC, in any order.
+  other_events = simplified_events[1:]
+  expected_other_events = [
+      (
+          'test_workflow_agent_list_routes@1/NodeB@1',
+          {'output': 'B'},
+      ),
+      (
+          'test_workflow_agent_list_routes@1/NodeC@1',
+          {'output': 'C'},
+      ),
+  ]
+  assert len(other_events) == len(expected_other_events)
+  assert all(item in other_events for item in expected_other_events)
+
+
+@pytest.mark.asyncio
+async def test_run_async_with_default_route(request: pytest.FixtureRequest):
+  node_a = TestingNode(name='NodeA', output='A', route='unmatched_route')
+  node_b = TestingNode(name='NodeB', output='B')
+  node_c = TestingNode(name='NodeC', output='C')
+  node_d = TestingNode(name='NodeD', output='D')
+
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(
+              from_node=node_a,
+              to_node=node_b,
+              route='route_b',
+          ),
+          # This edge has the DEFAULT_ROUTE tag.
+          Edge(
+              from_node=node_a,
+              to_node=node_c,
+              route=DEFAULT_ROUTE,
+          ),
+          # This edge has no route tag, so it should always be triggered.
+          Edge(
+              from_node=node_a,
+              to_node=node_d,
+          ),
+      ],
+  )
+  agent = Workflow(
+      name='test_workflow_agent_default_route',
+      graph=graph,
+  )
+
+  app = App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  simplified_events = simplify_events_with_node(events)
+
+  assert len(simplified_events) == 3
+  assert simplified_events[0] == (
+      'test_workflow_agent_default_route@1/NodeA@1',
+      {'output': 'A'},
+  )
+
+  # Check that NodeC (default route) and NodeD (untagged) are triggered.
+  other_events = simplified_events[1:]
+  expected_other_events = [
+      (
+          'test_workflow_agent_default_route@1/NodeC@1',
+          {'output': 'C'},
+      ),
+      (
+          'test_workflow_agent_default_route@1/NodeD@1',
+          {'output': 'D'},
+      ),
+  ]
+  assert len(other_events) == len(expected_other_events)
+  assert all(item in other_events for item in expected_other_events)
+
+
+@pytest.mark.asyncio
+async def test_run_async_default_route_not_triggered_if_match(
+    request: pytest.FixtureRequest,
+):
+  node_a = TestingNode(name='NodeA', output='A', route='route_b')
+  node_b = TestingNode(name='NodeB', output='B')
+  node_c = TestingNode(name='NodeC', output='C')
+
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(
+              from_node=node_a,
+              to_node=node_b,
+              route='route_b',
+          ),
+          # This edge has the DEFAULT_ROUTE tag.
+          Edge(
+              from_node=node_a,
+              to_node=node_c,
+              route=DEFAULT_ROUTE,
+          ),
+      ],
+  )
+  agent = Workflow(
+      name='test_workflow_agent_default_route_not_triggered',
+      graph=graph,
+  )
+
+  app = App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  simplified_events = simplify_events_with_node(events)
+
+  assert simplified_events == [
+      (
+          'test_workflow_agent_default_route_not_triggered@1/NodeA@1',
+          {'output': 'A'},
+      ),
+      (
+          'test_workflow_agent_default_route_not_triggered@1/NodeB@1',
+          {'output': 'B'},
+      ),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_run_async_with_untagged_edges(request: pytest.FixtureRequest):
+  node_a = TestingNode(name='NodeA', output='A', route='route_b')
+  node_b = TestingNode(name='NodeB', output='B')
+  node_c = TestingNode(name='NodeC', output='C')
+  node_d = TestingNode(name='NodeD', output='D')
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(
+              from_node=node_a,
+              to_node=node_b,
+              route='route_b',
+          ),
+          Edge(
+              from_node=node_a,
+              to_node=node_c,
+              route='route_c',
+          ),
+          # This edge has no route tag, so it should always be triggered.
+          Edge(
+              from_node=node_a,
+              to_node=node_d,
+          ),
+      ],
+  )
+  agent = Workflow(
+      name='test_workflow_agent_untagged_edges',
+      graph=graph,
+  )
+
+  app = App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  simplified_events = simplify_events_with_node(events)
+
+  assert len(simplified_events) == 3
+  assert simplified_events[0] == (
+      'test_workflow_agent_untagged_edges@1/NodeA@1',
+      {'output': 'A'},
+  )
+
+  # Check that NodeB and NodeD are triggered.
+  other_events = simplified_events[1:]
+  expected_other_events = [
+      (
+          'test_workflow_agent_untagged_edges@1/NodeB@1',
+          {'output': 'B'},
+      ),
+      (
+          'test_workflow_agent_untagged_edges@1/NodeD@1',
+          {'output': 'D'},
+      ),
+  ]
+  assert len(other_events) == len(expected_other_events)
+  assert all(item in other_events for item in expected_other_events)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'emitted_route, expected_target',
+    [
+        ('route_a', 'Target'),
+        ('route_b', 'Target'),
+        ('route_c', 'Other'),
+    ],
+    ids=['route_a_matches_list', 'route_b_matches_list', 'route_c_single'],
+)
+async def test_edge_with_multiple_routes(
+    request: pytest.FixtureRequest, emitted_route, expected_target
+):
+  """Tests that an edge with a list of routes matches any of them."""
+  node_router = TestingNode(name='Router', output='R', route=emitted_route)
+  node_target = TestingNode(name='Target', output='T')
+  node_other = TestingNode(name='Other', output='O')
+
+  agent = Workflow(
+      name='test_multi_route',
+      edges=[
+          (START, node_router),
+          Edge(
+              from_node=node_router,
+              to_node=node_target,
+              route=['route_a', 'route_b'],
+          ),
+          (node_router, {'route_c': node_other}),
+      ],
+  )
+
+  app = App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+  assert simplify_events_with_node(events) == [
+      ('test_multi_route@1/Router@1', {'output': 'R'}),
+      (
+          f'test_multi_route@1/{expected_target}@1',
+          {
+              'output': 'T' if expected_target == 'Target' else 'O',
+          },
+      ),
+  ]
+
+
+# --- Routing map integration tests ---
+
+
+@pytest.mark.asyncio
+async def test_routing_map_with_default_route(
+    request: pytest.FixtureRequest,
+):
+  """Tests that DEFAULT_ROUTE works as a fallback in routing maps."""
+  node_a = TestingNode(name='NodeA', output='A', route='unmatched_route')
+  node_b = TestingNode(name='NodeB', output='B')
+  node_c = TestingNode(name='NodeC', output='C')
+
+  agent = Workflow(
+      name='test_routing_map_default',
+      edges=[
+          (START, node_a),
+          (node_a, {'route_b': node_b, DEFAULT_ROUTE: node_c}),
+      ],
+  )
+
+  app = App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+  assert simplify_events_with_node(events) == [
+      (
+          'test_routing_map_default@1/NodeA@1',
+          {'output': 'A'},
+      ),
+      (
+          'test_routing_map_default@1/NodeC@1',
+          {'output': 'C'},
+      ),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_routing_map_mixed_with_other_formats(
+    request: pytest.FixtureRequest,
+):
+  """Tests that routing maps coexist with tuples and Edge objects."""
+  node_a = TestingNode(name='NodeA', output='A', route='route_b')
+  node_b = TestingNode(name='NodeB', output='B')
+  node_c = TestingNode(name='NodeC', output='C')
+  node_d = TestingNode(name='NodeD', output='D')
+
+  agent = Workflow(
+      name='test_routing_map_mixed',
+      edges=[
+          (START, node_a),
+          (node_a, {'route_b': node_b, 'route_c': node_c}),
+          (node_b, node_d),  # unconditional 2-tuple
+          Edge(from_node=node_c, to_node=node_d),  # Edge object
+      ],
+  )
+
+  app = App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+  assert simplify_events_with_node(events) == [
+      (
+          'test_routing_map_mixed@1/NodeA@1',
+          {'output': 'A'},
+      ),
+      (
+          'test_routing_map_mixed@1/NodeB@1',
+          {'output': 'B'},
+      ),
+      (
+          'test_routing_map_mixed@1/NodeD@1',
+          {'output': 'D'},
+      ),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_routing_map_fan_out_runs_both_targets(
+    request: pytest.FixtureRequest,
+):
+  """Tests that fan-out in routing maps triggers both targets at runtime."""
+  node_a = TestingNode(name='NodeA', output='A', route='route_x')
+  node_b = TestingNode(name='NodeB', output='B')
+  node_c = TestingNode(name='NodeC', output='C')
+  gate = JoinNode(name='Gate')
+
+  agent = Workflow(
+      name='test_routing_map_fan_out',
+      edges=[
+          (START, node_a),
+          (node_a, {'route_x': (node_b, node_c)}),
+          (node_b, gate),
+          (node_c, gate),
+      ],
+  )
+
+  app = App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+  simplified = simplify_events_with_node(events)
+
+  # NodeB and NodeC should both be triggered, in any order.
+  # Gate node also produces output combining the two.
+  outputs = [
+      e
+      for e in simplified
+      if isinstance(e[1], dict) and e[1].get('output') is not None
+  ]
+
+  assert len(outputs) == 4
+  assert outputs[0] == (
+      'test_routing_map_fan_out@1/NodeA@1',
+      {'output': 'A'},
+  )
+
+  # Gate should be last
+  assert outputs[-1] == (
+      'test_routing_map_fan_out@1/Gate@1',
+      {'output': {'NodeB': 'B', 'NodeC': 'C'}},
+  )
+
+  other = outputs[1:3]
+  expected = [
+      (
+          'test_routing_map_fan_out@1/NodeB@1',
+          {'output': 'B'},
+      ),
+      (
+          'test_routing_map_fan_out@1/NodeC@1',
+          {'output': 'C'},
+      ),
+  ]
+  assert len(other) == len(expected)
+  assert all(item in other for item in expected)
+
+
+@pytest.mark.asyncio
+async def test_fan_in_with_route(request: pytest.FixtureRequest):
+  """Fan-in with conditional routes — both route to same target."""
+  a = TestingNode(name='a', output='A', route='route1')
+  b = TestingNode(name='b', output='B', route='route1')
+  c = TestingNode(name='c', output='C')
+  wf = Workflow(
+      name='wf',
+      edges=[
+          (START, a),
+          (START, b),
+          ((a, b), {'route1': c}),
+      ],
+  )
+
+  app = App(name=request.function.__name__, root_agent=wf)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+  simplified = simplify_events_with_node(events)
+
+  # 'c' should receive from both 'a' and 'b' and run twice.
+  c_events = [e for e in simplified if e[0].split('/')[-1].split('@')[0] == 'c']
+  assert len(c_events) == 2
+
+
+@pytest.mark.asyncio
+async def test_fan_out_with_route(request: pytest.FixtureRequest):
+  """Fan-out via route to multiple terminals raises ValueError."""
+  router = TestingNode(name='r', output='R', route='route1')
+  b = TestingNode(name='b', output='B')
+  c = TestingNode(name='c', output='C')
+  wf = Workflow(
+      name='wf',
+      edges=[
+          (START, router),
+          (router, {'route1': (b, c)}),
+      ],
+  )
+
+  app = App(name=request.function.__name__, root_agent=wf)
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  with pytest.raises(ValueError, match='multiple terminal nodes'):
+    await runner.run_async(testing_utils.get_user_content('start'))
+
+
+@pytest.mark.asyncio
+async def test_fan_in_out_with_route(request: pytest.FixtureRequest):
+  """Fan-in/out with routes to multiple terminals raises ValueError."""
+  a = TestingNode(name='a', output='A', route='route1')
+  b = TestingNode(name='b', output='B', route='route1')
+  c = TestingNode(name='c', output='C')
+  d = TestingNode(name='d', output='D')
+  wf = Workflow(
+      name='wf',
+      edges=[
+          (START, a),
+          (START, b),
+          ((a, b), {'route1': (c, d)}),
+      ],
+  )
+
+  app = App(name=request.function.__name__, root_agent=wf)
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  with pytest.raises(ValueError, match='multiple terminal nodes'):
+    await runner.run_async(testing_utils.get_user_content('start'))

tests/unittests/workflow/test_workflow_schema.py

@@ -0,0 +1,841 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for Workflow schema validation (input_schema, output_schema)."""
+
+from __future__ import annotations
+
+from google.adk.apps.app import App
+from google.adk.events.event import Event
+from google.adk.runners import Runner
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.workflow import BaseNode
+from google.adk.workflow import JoinNode
+from google.adk.workflow import START
+from google.adk.workflow._workflow import Workflow
+from google.genai import types
+from pydantic import BaseModel
+from pydantic import ValidationError
+import pytest
+
+from .. import testing_utils
+from .workflow_testing_utils import create_parent_invocation_context
+
+
+class _OutputModel(BaseModel):
+  name: str
+  value: int
+
+
+class _OtherModel(BaseModel):
+  name: str
+  value: int
+  extra: str = 'default'
+
+
+@pytest.mark.asyncio
+async def test_workflow_output_schema_validates_terminal(
+    request: pytest.FixtureRequest,
+):
+  """Workflow.output_schema validates when downstream reads the output."""
+
+  def produce() -> dict:
+    return {'name': 'result', 'value': 10}
+
+  def consume(node_input: dict) -> str:
+    return f"got {node_input['name']}"
+
+  inner = Workflow(
+      name='wf',
+      edges=[(START, produce)],
+      output_schema=_OutputModel,
+  )
+  outer = Workflow(
+      name='outer',
+      edges=[(START, inner, consume)],
+  )
+  app = App(name=request.function.__name__, root_agent=outer)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+  data_events = [e for e in events if isinstance(e, Event) and e.output]
+  assert any(e.output == 'got result' for e in data_events)
+
+
+@pytest.mark.asyncio
+async def test_workflow_output_schema_rejects_invalid(
+    request: pytest.FixtureRequest,
+):
+  """Workflow.output_schema rejects invalid output when downstream reads."""
+
+  def produce_bad() -> dict:
+    return {'wrong_field': 'oops'}
+
+  def consume(node_input: dict) -> str:
+    return 'should not reach'
+
+  inner = Workflow(
+      name='wf',
+      edges=[(START, produce_bad)],
+      output_schema=_OutputModel,
+  )
+  outer = Workflow(
+      name='outer',
+      edges=[(START, inner, consume)],
+  )
+  app = App(name=request.function.__name__, root_agent=outer)
+  runner = testing_utils.InMemoryRunner(app=app)
+  with pytest.raises(ValueError):
+    await runner.run_async(testing_utils.get_user_content('start'))
+
+
+@pytest.mark.asyncio
+async def test_workflow_output_schema_coerces_defaults(
+    request: pytest.FixtureRequest,
+):
+  """Workflow.output_schema coerces terminal output (fills defaults)."""
+
+  def produce() -> dict:
+    return {'name': 'x', 'value': 1}
+
+  def consume(node_input: dict) -> dict:
+    return node_input
+
+  inner = Workflow(
+      name='wf',
+      edges=[(START, produce)],
+      output_schema=_OtherModel,
+  )
+  outer = Workflow(
+      name='outer',
+      edges=[(START, inner, consume)],
+  )
+  app = App(name=request.function.__name__, root_agent=outer)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+  data_events = [e for e in events if isinstance(e, Event) and e.output]
+  consume_events = [e for e in data_events if e.node_info.name == 'consume']
+  assert len(consume_events) == 1
+  assert consume_events[0].output == {
+      'name': 'x',
+      'value': 1,
+      'extra': 'default',
+  }
+
+
+@pytest.mark.asyncio
+async def test_nested_workflow_output_schema(
+    request: pytest.FixtureRequest,
+):
+  """Nested workflow's output_schema validates before passing to parent."""
+
+  def inner_produce() -> dict:
+    return {'name': 'inner', 'value': 3}
+
+  inner_workflow = Workflow(
+      name='inner_wf',
+      edges=[(START, inner_produce)],
+      output_schema=_OutputModel,
+  )
+
+  def outer_consume(node_input: dict) -> str:
+    return f"got {node_input['name']}"
+
+  outer_workflow = Workflow(
+      name='outer_wf',
+      edges=[
+          (START, inner_workflow),
+          (inner_workflow, outer_consume),
+      ],
+  )
+  app = App(name=request.function.__name__, root_agent=outer_workflow)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  data_events = [e for e in events if isinstance(e, Event) and e.output]
+  assert any(e.output == 'got inner' for e in data_events)
+
+
+@pytest.mark.asyncio
+async def test_workflow_output_schema_validates_multiple_terminals(
+    request: pytest.FixtureRequest,
+):
+  """Each terminal output is validated when downstream reads."""
+
+  class _JoinedModel(BaseModel):
+    branch_a: _OtherModel
+    branch_b: _OtherModel
+
+  def branch_a() -> dict:
+    return {'name': 'from_a', 'value': 1}
+
+  def branch_b() -> dict:
+    return {'name': 'from_b', 'value': 2}
+
+  join = JoinNode(name='join')
+
+  inner = Workflow(
+      name='wf',
+      edges=[
+          (START, branch_a),
+          (START, branch_b),
+          (branch_a, join),
+          (branch_b, join),
+      ],
+      output_schema=_JoinedModel,
+  )
+
+  def consume(node_input: dict) -> dict:
+    return node_input
+
+  outer = Workflow(
+      name='outer',
+      edges=[(START, inner, consume)],
+  )
+  app = App(name=request.function.__name__, root_agent=outer)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+  consume_events = [
+      e
+      for e in events
+      if isinstance(e, Event) and e.output and e.node_info.name == 'consume'
+  ]
+  assert len(consume_events) == 1
+
+  # Both terminal outputs should have 'extra' filled by _OtherModel default.
+  output = consume_events[0].output
+  assert output['branch_a'] == {
+      'name': 'from_a',
+      'value': 1,
+      'extra': 'default',
+  }
+  assert output['branch_b'] == {
+      'name': 'from_b',
+      'value': 2,
+      'extra': 'default',
+  }
+
+
+@pytest.mark.asyncio
+async def test_workflow_output_schema_rejects_invalid_among_multiple_terminals(
+    request: pytest.FixtureRequest,
+):
+  """One invalid terminal among multiple raises validation error."""
+
+  class _JoinedModel(BaseModel):
+    branch_good: _OutputModel
+    branch_bad: _OutputModel
+
+  def branch_good() -> dict:
+    return {'name': 'ok', 'value': 1}
+
+  def branch_bad() -> dict:
+    return {'wrong_field': 'oops'}
+
+  join = JoinNode(name='join')
+
+  inner = Workflow(
+      name='wf',
+      edges=[
+          (START, branch_good),
+          (START, branch_bad),
+          (branch_good, join),
+          (branch_bad, join),
+      ],
+      output_schema=_JoinedModel,
+  )
+
+  def consume(node_input: dict) -> str:
+    return 'should not reach'
+
+  outer = Workflow(
+      name='outer',
+      edges=[(START, inner, consume)],
+  )
+  app = App(name=request.function.__name__, root_agent=outer)
+  runner = testing_utils.InMemoryRunner(app=app)
+  with pytest.raises(ValueError):
+    await runner.run_async(testing_utils.get_user_content('start'))
+
+
+# ── Primitive and generic type output_schema ─────────────────────────
+
+
+@pytest.mark.asyncio
+async def test_workflow_output_schema_int_coerces(
+    request: pytest.FixtureRequest,
+):
+  """Workflow output_schema=int coerces string to int at read time."""
+
+  def produce() -> str:
+    return '42'
+
+  def consume(node_input: int) -> int:
+    return node_input
+
+  inner = Workflow(
+      name='wf',
+      edges=[(START, produce)],
+      output_schema=int,
+  )
+  outer = Workflow(name='outer', edges=[(START, inner, consume)])
+  app = App(name=request.function.__name__, root_agent=outer)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+  consume_events = [
+      e
+      for e in events
+      if isinstance(e, Event)
+      and e.output is not None
+      and e.node_info.name == 'consume'
+  ]
+  assert len(consume_events) == 1
+  assert consume_events[0].output == 42
+
+
+@pytest.mark.asyncio
+async def test_workflow_output_schema_int_rejects_invalid(
+    request: pytest.FixtureRequest,
+):
+  """Workflow output_schema=int rejects non-coercible value at read time."""
+
+  def produce() -> dict:
+    return {'key': 'value'}
+
+  def consume(node_input: int) -> int:
+    return node_input
+
+  inner = Workflow(
+      name='wf',
+      edges=[(START, produce)],
+      output_schema=int,
+  )
+  outer = Workflow(name='outer', edges=[(START, inner, consume)])
+  app = App(name=request.function.__name__, root_agent=outer)
+  runner = testing_utils.InMemoryRunner(app=app)
+  with pytest.raises(ValueError):
+    await runner.run_async(testing_utils.get_user_content('start'))
+
+
+@pytest.mark.asyncio
+async def test_workflow_output_schema_list_of_str(
+    request: pytest.FixtureRequest,
+):
+  """Workflow output_schema=list[str] validates list output at read time."""
+
+  def produce() -> list:
+    return ['a', 'b', 'c']
+
+  def consume(node_input: list) -> list:
+    return node_input
+
+  inner = Workflow(
+      name='wf',
+      edges=[(START, produce)],
+      output_schema=list[str],
+  )
+  outer = Workflow(name='outer', edges=[(START, inner, consume)])
+  app = App(name=request.function.__name__, root_agent=outer)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+  consume_events = [
+      e
+      for e in events
+      if isinstance(e, Event)
+      and e.output is not None
+      and e.node_info.name == 'consume'
+  ]
+  assert len(consume_events) == 1
+  assert consume_events[0].output == ['a', 'b', 'c']
+
+
+@pytest.mark.asyncio
+async def test_workflow_output_schema_list_of_basemodel(
+    request: pytest.FixtureRequest,
+):
+  """Workflow output_schema=list[BaseModel] validates and serializes."""
+
+  def produce() -> list:
+    return [
+        {'name': 'x', 'value': 1},
+        {'name': 'y', 'value': 2},
+    ]
+
+  def consume(node_input: list) -> list:
+    return node_input
+
+  inner = Workflow(
+      name='wf',
+      edges=[(START, produce)],
+      output_schema=list[_OutputModel],
+  )
+  outer = Workflow(name='outer', edges=[(START, inner, consume)])
+  app = App(name=request.function.__name__, root_agent=outer)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+  consume_events = [
+      e
+      for e in events
+      if isinstance(e, Event)
+      and e.output is not None
+      and e.node_info.name == 'consume'
+  ]
+  assert len(consume_events) == 1
+  assert consume_events[0].output == [
+      {'name': 'x', 'value': 1},
+      {'name': 'y', 'value': 2},
+  ]
+
+
+# ── End-to-end: input_schema + output_schema combined ──────────────
+
+
+class _TaskOutput(BaseModel):
+  result: str
+  score: int
+
+
+class _ReviewResult(BaseModel):
+  result: str
+  score: int
+  reviewer: str = 'auto'
+
+
+@pytest.mark.asyncio
+async def test_e2e_input_and_output_schema_pipeline(
+    request: pytest.FixtureRequest,
+):
+  """output_schema on producer + input_schema on consumer validates both."""
+
+  def produce() -> _TaskOutput:
+    return {'result': 'done', 'score': 88}
+
+  def consume(node_input: _TaskOutput) -> str:
+    return f'result={node_input.result}, score={node_input.score}'
+
+  agent = Workflow(
+      name='wf',
+      edges=[(START, produce), (produce, consume)],
+  )
+  app = App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  data_events = [e for e in events if isinstance(e, Event) and e.output]
+  assert any(e.output == 'result=done, score=88' for e in data_events)
+
+
+@pytest.mark.asyncio
+async def test_e2e_output_schema_fails_before_input_schema(
+    request: pytest.FixtureRequest,
+):
+  """Producer output_schema failure prevents consumer from running."""
+
+  def produce_bad() -> _TaskOutput:
+    return {'wrong': 'shape'}
+
+  def consume(node_input: _TaskOutput) -> str:
+    return 'should not reach'
+
+  agent = Workflow(
+      name='wf',
+      edges=[(START, produce_bad), (produce_bad, consume)],
+  )
+  app = App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app)
+  with pytest.raises(ValueError):
+    await runner.run_async(testing_utils.get_user_content('start'))
+
+
+@pytest.mark.asyncio
+async def test_e2e_fan_out_join_with_schemas(
+    request: pytest.FixtureRequest,
+):
+  """Fan-out -> join -> consume with input/output schemas."""
+
+  class _JoinedResult(BaseModel):
+    analyzer: dict
+    summarizer: dict
+
+  def analyzer() -> _TaskOutput:
+    return {'result': 'analysis complete', 'score': 92}
+
+  def summarizer() -> _TaskOutput:
+    return {'result': 'summary complete', 'score': 78}
+
+  join = JoinNode(name='join')
+
+  def reviewer(node_input: dict) -> _ReviewResult:
+    a = node_input['analyzer']
+    s = node_input['summarizer']
+    return {
+        'result': f"{a['result']} + {s['result']}",
+        'score': (a['score'] + s['score']) // 2,
+    }
+
+  agent = Workflow(
+      name='wf',
+      edges=[
+          (START, analyzer),
+          (START, summarizer),
+          (analyzer, join),
+          (summarizer, join),
+          (join, reviewer),
+      ],
+  )
+  app = App(name=request.function.__name__, root_agent=agent)
+  runner = testing_utils.InMemoryRunner(app=app)
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  terminal = [
+      e
+      for e in events
+      if isinstance(e, Event)
+      and e.output is not None
+      and e.node_info.name == 'reviewer'
+  ]
+  assert len(terminal) == 1
+  assert terminal[0].output == {
+      'result': 'analysis complete + summary complete',
+      'score': 85,
+      'reviewer': 'auto',
+  }
+  assert 'wf' in terminal[0].node_info.path
+  assert 'reviewer' in terminal[0].node_info.path
+
+
+@pytest.mark.asyncio
+async def test_start_node_with_str_input_schema():
+  """input_schema=str parses user text."""
+
+  class _AssertingNode(BaseNode):
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Any, None]:
+      assert node_input == 'hello'
+      yield 'done'
+
+  node = _AssertingNode(name='node', input_schema=str)
+  wf = Workflow(name='wf', edges=[(START, node)])
+
+  ss = InMemorySessionService()
+  runner = Runner(app_name='test', node=wf, session_service=ss)
+  session = await ss.create_session(app_name='test', user_id='u')
+
+  msg = types.Content(parts=[types.Part(text='hello')], role='user')
+  events = []
+  async for event in runner.run_async(
+      user_id='u', session_id=session.id, new_message=msg
+  ):
+    events.append(event)
+
+  data_events = [e for e in events if isinstance(e, Event) and e.output]
+  assert any(e.output == 'done' for e in data_events)
+
+
+@pytest.mark.asyncio
+async def test_start_node_with_int_input_schema():
+  """input_schema=int parses user text to int."""
+
+  class _AssertingNode(BaseNode):
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Any, None]:
+      assert node_input == 42
+      yield 'done'
+
+  node = _AssertingNode(name='node', input_schema=int)
+  wf = Workflow(name='wf', edges=[(START, node)])
+
+  ss = InMemorySessionService()
+  runner = Runner(app_name='test', node=wf, session_service=ss)
+  session = await ss.create_session(app_name='test', user_id='u')
+
+  msg = types.Content(parts=[types.Part(text='42')], role='user')
+  events = []
+
+  async for event in runner.run_async(
+      user_id='u', session_id=session.id, new_message=msg
+  ):
+    events.append(event)
+
+  data_events = [e for e in events if isinstance(e, Event) and e.output]
+  assert any(e.output == 'done' for e in data_events)
+
+
+@pytest.mark.asyncio
+async def test_start_node_with_int_list_input_schema():
+  """input_schema=list[int] parses JSON list."""
+
+  class _AssertingNode(BaseNode):
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Any, None]:
+      assert node_input == [1, 2, 3]
+      yield 'done'
+
+  node = _AssertingNode(name='node', input_schema=list[int])
+  wf = Workflow(name='wf', edges=[(START, node)])
+
+  ss = InMemorySessionService()
+  runner = Runner(app_name='test', node=wf, session_service=ss)
+  session = await ss.create_session(app_name='test', user_id='u')
+
+  msg = types.Content(parts=[types.Part(text='[1, 2, 3]')], role='user')
+  events = []
+
+  async for event in runner.run_async(
+      user_id='u', session_id=session.id, new_message=msg
+  ):
+    events.append(event)
+
+  data_events = [e for e in events if isinstance(e, Event) and e.output]
+  assert any(e.output == 'done' for e in data_events)
+
+
+@pytest.mark.asyncio
+async def test_start_node_with_invalid_input_schema():
+  """Invalid input against schema raises error."""
+
+  class _MyModel(BaseModel):
+    age: int
+
+  class _AssertingNode(BaseNode):
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Any, None]:
+      yield 'done'
+
+  node = _AssertingNode(name='node', input_schema=_MyModel)
+  wf = Workflow(name='wf', edges=[(START, node)])
+
+  ss = InMemorySessionService()
+  runner = Runner(app_name='test', node=wf, session_service=ss)
+  session = await ss.create_session(app_name='test', user_id='u')
+
+  # Pass invalid input (Content instead of dict with age)
+  msg = types.Content(parts=[types.Part(text='hello')], role='user')
+
+  # We expect it to raise ValidationError
+  with pytest.raises(ValidationError):
+    async for event in runner.run_async(
+        user_id='u', session_id=session.id, new_message=msg
+    ):
+      pass
+
+
+@pytest.mark.asyncio
+async def test_start_node_receives_parsed_user_content_with_schema():
+  """Parsed input replaces raw Content for first node."""
+
+  class _AssertingNode(BaseNode):
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Any, None]:
+      assert isinstance(node_input, str)
+      assert node_input == 'hello'
+      yield 'done'
+
+  node = _AssertingNode(name='node', input_schema=str)
+  wf = Workflow(name='wf', edges=[(START, node)])
+
+  ss = InMemorySessionService()
+  runner = Runner(app_name='test', node=wf, session_service=ss)
+  session = await ss.create_session(app_name='test', user_id='u')
+
+  msg = types.Content(parts=[types.Part(text='hello')], role='user')
+  events = []
+  async for event in runner.run_async(
+      user_id='u', session_id=session.id, new_message=msg
+  ):
+    events.append(event)
+
+  data_events = [e for e in events if isinstance(e, Event) and e.output]
+  assert any(e.output == 'done' for e in data_events)
+
+
+@pytest.mark.asyncio
+async def test_workflow_with_invalid_output_schema():
+  """Workflow raises ValidationError if terminal output doesn't match output_schema."""
+
+  from pydantic import ValidationError
+
+  class _MyModel(BaseModel):
+    name: str
+
+  class _MyNode(BaseNode):
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Any, None]:
+      yield {'age': 10}
+
+  node = _MyNode(name='node')
+  wf = Workflow(name='wf', edges=[(START, node)], output_schema=_MyModel)
+
+  ss = InMemorySessionService()
+  runner = Runner(app_name='test', node=wf, session_service=ss)
+  session = await ss.create_session(app_name='test', user_id='u')
+
+  msg = types.Content(parts=[types.Part(text='hello')], role='user')
+
+  with pytest.raises(ValidationError):
+    async for event in runner.run_async(
+        user_id='u', session_id=session.id, new_message=msg
+    ):
+      pass
+
+
+@pytest.mark.asyncio
+async def test_node_returns_content_json_parsed():
+  """Node output as types.Content containing JSON is parsed if output_schema is defined."""
+
+  class _MyModel(BaseModel):
+    name: str
+    age: int
+
+  class _MyNode(BaseNode):
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Any, None]:
+      yield types.Content(
+          parts=[types.Part(text='{"name": "Alice", "age": 30}')]
+      )
+
+  node = _MyNode(name='node', output_schema=_MyModel)
+  wf = Workflow(name='wf', edges=[(START, node)])
+
+  ss = InMemorySessionService()
+  runner = Runner(app_name='test', node=wf, session_service=ss)
+  session = await ss.create_session(app_name='test', user_id='u')
+
+  msg = types.Content(parts=[types.Part(text='hello')], role='user')
+  events = []
+
+  async for event in runner.run_async(
+      user_id='u', session_id=session.id, new_message=msg
+  ):
+    events.append(event)
+
+  data_events = [e for e in events if isinstance(e, Event) and e.output]
+
+  assert len(data_events) == 1
+  assert data_events[0].output == {'name': 'Alice', 'age': 30}
+
+
+@pytest.mark.asyncio
+async def test_node_returns_raw_string_not_parsed():
+  """Node output as raw JSON string is NOT parsed if output_schema is defined."""
+  from pydantic import ValidationError
+
+  class _MyModel(BaseModel):
+    name: str
+    age: int
+
+  class _MyNode(BaseNode):
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Any, None]:
+      # This should fail validation because it's a string, not a dict/model
+      yield '{"name": "Alice", "age": 30}'
+
+  node = _MyNode(name='node', output_schema=_MyModel)
+  wf = Workflow(name='wf', edges=[(START, node)])
+
+  ss = InMemorySessionService()
+  runner = Runner(app_name='test', node=wf, session_service=ss)
+  session = await ss.create_session(app_name='test', user_id='u')
+
+  msg = types.Content(parts=[types.Part(text='hello')], role='user')
+
+  with pytest.raises(ValidationError):
+    async for _ in runner.run_async(
+        user_id='u', session_id=session.id, new_message=msg
+    ):
+      pass
+
+
+@pytest.mark.asyncio
+async def test_output_schema_enforced_by_runtime_without_manual_validation():
+  """Runtime enforces output_schema even when _run_impl doesn't call _validate_output_data."""
+  from pydantic import ValidationError
+
+  class _MyModel(BaseModel):
+    name: str
+    age: int
+
+  class _NaiveNode(BaseNode):
+    """Node that yields raw data without any manual validation."""
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Any, None]:
+      yield {'color': 'red'}  # Does NOT match _MyModel
+
+  node = _NaiveNode(name='node', output_schema=_MyModel)
+  wf = Workflow(name='wf', edges=[(START, node)])
+
+  ss = InMemorySessionService()
+  runner = Runner(app_name='test', node=wf, session_service=ss)
+  session = await ss.create_session(app_name='test', user_id='u')
+
+  msg = types.Content(parts=[types.Part(text='hello')], role='user')
+
+  with pytest.raises(ValidationError):
+    async for _ in runner.run_async(
+        user_id='u', session_id=session.id, new_message=msg
+    ):
+      pass
+
+
+@pytest.mark.asyncio
+async def test_output_schema_enforced_for_valid_raw_yield():
+  """Runtime validates and coerces valid raw yields against output_schema."""
+
+  class _MyModel(BaseModel):
+    name: str
+    age: int
+
+  class _NaiveNode(BaseNode):
+    """Node that yields valid raw data without manual validation."""
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Any, None]:
+      yield {'name': 'Alice', 'age': 30}
+
+  node = _NaiveNode(name='node', output_schema=_MyModel)
+  wf = Workflow(name='wf', edges=[(START, node)])
+
+  ss = InMemorySessionService()
+  runner = Runner(app_name='test', node=wf, session_service=ss)
+  session = await ss.create_session(app_name='test', user_id='u')
+
+  msg = types.Content(parts=[types.Part(text='hello')], role='user')
+  events = []
+
+  async for event in runner.run_async(
+      user_id='u', session_id=session.id, new_message=msg
+  ):
+    events.append(event)
+
+  data_events = [e for e in events if isinstance(e, Event) and e.output]
+  assert len(data_events) == 1
+  assert data_events[0].output == {'name': 'Alice', 'age': 30}

tests/unittests/workflow/testing_utils.py

@@ -0,0 +1,507 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import contextlib
+import copy
+from typing import Any
+from typing import AsyncGenerator
+from typing import Generator
+from typing import Optional
+
+from google.adk.agents.context import Context as WorkflowContext
+from google.adk.agents.invocation_context import InvocationContext as BaseInvocationContext
+from google.adk.agents.live_request_queue import LiveRequestQueue
+from google.adk.agents.llm_agent import Agent
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.agents.run_config import RunConfig
+from google.adk.apps.app import App
+from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
+from google.adk.events.event import Event
+from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
+from google.adk.models.base_llm import BaseLlm
+from google.adk.models.base_llm_connection import BaseLlmConnection
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from google.adk.plugins.base_plugin import BasePlugin
+from google.adk.plugins.plugin_manager import PluginManager
+from google.adk.runners import InMemoryRunner as AfInMemoryRunner
+from google.adk.runners import Runner
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.sessions.session import Session
+from google.adk.utils.context_utils import Aclosing
+from google.genai import types
+from google.genai.types import Part
+from typing_extensions import override
+
+
+def create_test_agent(name: str = 'test_agent') -> LlmAgent:
+  """Create a simple test agent for use in unit tests.
+
+  Args:
+    name: The name of the test agent.
+
+  Returns:
+    A configured LlmAgent instance suitable for testing.
+  """
+  return LlmAgent(name=name)
+
+
+class UserContent(types.Content):
+
+  def __init__(self, text_or_part: str):
+    parts = [
+        types.Part.from_text(text=text_or_part)
+        if isinstance(text_or_part, str)
+        else text_or_part
+    ]
+    super().__init__(role='user', parts=parts)
+
+
+class ModelContent(types.Content):
+
+  def __init__(self, parts: list[types.Part]):
+    super().__init__(role='model', parts=parts)
+
+
+async def create_invocation_context(
+    agent: Agent,
+    user_content: str = '',
+    run_config: RunConfig = None,
+    plugins: list[BasePlugin] = [],
+):
+  invocation_id = 'test_id'
+  artifact_service = InMemoryArtifactService()
+  session_service = InMemorySessionService()
+  memory_service = InMemoryMemoryService()
+  invocation_context = BaseInvocationContext(
+      artifact_service=artifact_service,
+      session_service=session_service,
+      memory_service=memory_service,
+      plugin_manager=PluginManager(plugins=plugins),
+      invocation_id=invocation_id,
+      agent=agent,
+      session=await session_service.create_session(
+          app_name='test_app', user_id='test_user'
+      ),
+      user_content=types.Content(
+          role='user', parts=[types.Part.from_text(text=user_content)]
+      ),
+      run_config=run_config or RunConfig(),
+  )
+  if user_content:
+    append_user_content(
+        invocation_context, [types.Part.from_text(text=user_content)]
+    )
+  return invocation_context
+
+
+async def create_workflow_context(
+    agent,
+    user_content='',
+) -> WorkflowContext:
+  """Create a WorkflowContext for isolated node testing.
+
+  Constructs the minimal InvocationContext and wraps it in a
+  WorkflowContext so that individual nodes can be tested in
+  isolation without running the full _SingleLlmAgent pipeline.
+  """
+  invocation_context = await create_invocation_context(agent, user_content)
+  return WorkflowContext(
+      invocation_context=invocation_context,
+      node_path='test',
+      run_id='test-execution',
+  )
+
+
+def append_user_content(
+    invocation_context: BaseInvocationContext, parts: list[types.Part]
+) -> Event:
+  session = invocation_context.session
+  event = Event(
+      invocation_id=invocation_context.invocation_id,
+      author='user',
+      content=types.Content(role='user', parts=parts),
+  )
+  session.events.append(event)
+  return event
+
+
+# Extracts the contents from the events and transform them into a list of
+# (author, simplified_content) tuples.
+def simplify_events(events: list[Event]) -> list[tuple[str, types.Part]]:
+  res = []
+  for event in events:
+    if event.content:
+      author = event.author
+      res.append((author, simplify_content(event.content)))
+  return res
+
+
+END_OF_AGENT = 'end_of_agent'
+
+
+# Extracts the contents from the events and transform them into a list of
+# (author, simplified_content OR AgentState OR "end_of_agent") tuples.
+#
+# Could be used to compare events for testing resumability.
+def simplify_resumable_app_events(
+    events: list[Event],
+) -> list[(str, types.Part | str)]:
+  results = []
+  for event in events:
+    if event.content:
+      results.append((event.author, simplify_content(event.content)))
+    elif event.actions.end_of_agent:
+      results.append((event.author, END_OF_AGENT))
+    elif event.actions.agent_state is not None:
+      agent_state = event.actions.agent_state
+      if isinstance(agent_state, dict):
+        nodes = agent_state.get('nodes', {})
+        agent_state = {
+            'node_states': {
+                node_name: node_state.get('status')
+                for node_name, node_state in nodes.items()
+            }
+        }
+      results.append((event.author, agent_state))
+  return results
+
+
+# Simplifies the contents into a list of (author, simplified_content) tuples.
+def simplify_contents(contents: list[types.Content]) -> list[(str, types.Part)]:
+  return [(content.role, simplify_content(content)) for content in contents]
+
+
+# Simplifies the content so it's easier to assert.
+# - If there is only one part, return part
+# - If the only part is pure text, return stripped_text
+# - If there are multiple parts, return parts
+# - remove function_call_id if it exists
+def simplify_content(
+    content: types.Content,
+) -> str | types.Part | list[types.Part]:
+  content = copy.deepcopy(content)
+  for part in content.parts:
+    if part.function_call and part.function_call.id:
+      part.function_call.id = None
+    if part.function_response and part.function_response.id:
+      part.function_response.id = None
+  if len(content.parts) == 1:
+    if content.parts[0].text:
+      return content.parts[0].text.strip()
+    else:
+      return content.parts[0]
+  return content.parts
+
+
+def get_user_content(message: types.ContentUnion) -> types.Content:
+  return message if isinstance(message, types.Content) else UserContent(message)
+
+
+class TestInMemoryRunner(AfInMemoryRunner):
+  """InMemoryRunner that is tailored for tests, features async run method.
+
+  app_name is hardcoded as InMemoryRunner in the parent class.
+  """
+
+  async def run_async_with_new_session(
+      self, new_message: types.ContentUnion
+  ) -> list[Event]:
+
+    collected_events: list[Event] = []
+    async for event in self.run_async_with_new_session_agen(new_message):
+      collected_events.append(event)
+
+    return collected_events
+
+  async def run_async_with_new_session_agen(
+      self, new_message: types.ContentUnion
+  ) -> AsyncGenerator[Event, None]:
+    session = await self.session_service.create_session(
+        app_name='InMemoryRunner', user_id='test_user'
+    )
+    agen = self.run_async(
+        user_id=session.user_id,
+        session_id=session.id,
+        new_message=get_user_content(new_message),
+    )
+    async with Aclosing(agen):
+      async for event in agen:
+        yield event
+
+
+class InMemoryRunner:
+  """InMemoryRunner that is tailored for tests."""
+
+  def __init__(
+      self,
+      root_agent: Optional[Agent | LlmAgent] = None,
+      response_modalities: list[str] = None,
+      plugins: list[BasePlugin] = [],
+      app: Optional[App] = None,
+      node: Any = None,
+  ):
+    """Initializes the InMemoryRunner.
+
+    Args:
+      root_agent: The root agent to run, won't be used if app is provided.
+      response_modalities: The response modalities of the runner.
+      plugins: The plugins to use in the runner, won't be used if app is
+        provided.
+      app: The app to use in the runner.
+      node: The root node to run.
+    """
+    self._app = app
+    if node:
+      self.app_name = node.name
+      self.root_agent = None
+      self.runner = Runner(
+          node=node,
+          artifact_service=InMemoryArtifactService(),
+          session_service=InMemorySessionService(),
+          memory_service=InMemoryMemoryService(),
+          plugins=plugins,
+      )
+    elif not app:
+      self.app_name = 'test_app'
+      self.root_agent = root_agent
+      self.runner = Runner(
+          app_name='test_app',
+          agent=root_agent,
+          artifact_service=InMemoryArtifactService(),
+          session_service=InMemorySessionService(),
+          memory_service=InMemoryMemoryService(),
+          plugins=plugins,
+      )
+    else:
+      self.app_name = app.name
+      self.root_agent = app.root_agent
+      self.runner = Runner(
+          app=app,
+          artifact_service=InMemoryArtifactService(),
+          session_service=InMemorySessionService(),
+          memory_service=InMemoryMemoryService(),
+      )
+    self.session_id = None
+
+  @property
+  def session(self) -> Session:
+    if not self.session_id:
+      session = self.runner.session_service.create_session_sync(
+          app_name=self.app_name, user_id='test_user'
+      )
+      self.session_id = session.id
+      return session
+    return self.runner.session_service.get_session_sync(
+        app_name=self.app_name, user_id='test_user', session_id=self.session_id
+    )
+
+  def run(self, new_message: types.ContentUnion) -> list[Event]:
+    return list(
+        self.runner.run(
+            user_id=self.session.user_id,
+            session_id=self.session.id,
+            new_message=get_user_content(new_message),
+        )
+    )
+
+  @property
+  def is_resumable(self) -> bool:
+    """Returns whether the app is configured for resumable HITL."""
+    if hasattr(self, '_app') and self._app:
+      cfg = getattr(self._app, 'resumability_config', None)
+      return cfg is not None and cfg.is_resumable
+    return False
+
+  async def run_async(
+      self,
+      new_message: Optional[types.ContentUnion] = None,
+      invocation_id: Optional[str] = None,
+  ) -> list[Event]:
+    # For non-resumable apps, don't reuse invocation_id on resume.
+    # State reconstruction relies on scanning events from *previous*
+    # invocations, so the resume call must get a fresh invocation_id.
+    if invocation_id and not self.is_resumable:
+      invocation_id = None
+    events = []
+    async for event in self.runner.run_async(
+        user_id=self.session.user_id,
+        session_id=self.session.id,
+        invocation_id=invocation_id,
+        new_message=get_user_content(new_message) if new_message else None,
+    ):
+      events.append(event)
+    return events
+
+  def run_live(
+      self, live_request_queue: LiveRequestQueue, run_config: RunConfig = None
+  ) -> list[Event]:
+    collected_responses = []
+
+    async def consume_responses(session: Session):
+      run_res = self.runner.run_live(
+          session=session,
+          live_request_queue=live_request_queue,
+          run_config=run_config or RunConfig(),
+      )
+
+      async for response in run_res:
+        collected_responses.append(response)
+        # When we have enough response, we should return
+        if len(collected_responses) >= 1:
+          return
+
+    try:
+      session = self.session
+      asyncio.run(consume_responses(session))
+    except asyncio.TimeoutError:
+      print('Returning any partial results collected so far.')
+
+    return collected_responses
+
+
+class MockModel(BaseLlm):
+  model: str = 'mock'
+
+  requests: list[LlmRequest] = []
+  live_blobs: list[types.Blob] = []
+  live_contents: list[types.Content] = []
+  responses: list[LlmResponse]
+  error: Exception | None = None
+  response_index: int = -1
+
+  # Whether the mock model should wait for realtime input (blobs or content)
+  # to be sent before yielding pre-defined responses in live mode.
+  wait_for_realtime_input: bool = False
+
+  @classmethod
+  def create(
+      cls,
+      responses: (
+          list[types.Part]
+          | list[LlmResponse]
+          | list[str]
+          | list[list[types.Part]]
+      ),
+      error: Exception | None = None,
+      wait_for_realtime_input: bool = False,
+  ):
+    if error and not responses:
+      return cls(
+          responses=[],
+          error=error,
+          wait_for_realtime_input=wait_for_realtime_input,
+      )
+    if not responses:
+      return cls(responses=[], wait_for_realtime_input=wait_for_realtime_input)
+    elif isinstance(responses[0], LlmResponse):
+      # responses is list[LlmResponse]
+      return cls(
+          responses=responses, wait_for_realtime_input=wait_for_realtime_input
+      )
+    else:
+      responses = [
+          LlmResponse(content=ModelContent(item))
+          if isinstance(item, list) and isinstance(item[0], types.Part)
+          # responses is list[list[Part]]
+          else LlmResponse(
+              content=ModelContent(
+                  # responses is list[str] or list[Part]
+                  [Part(text=item) if isinstance(item, str) else item]
+              )
+          )
+          for item in responses
+          if item
+      ]
+
+      return cls(
+          responses=responses, wait_for_realtime_input=wait_for_realtime_input
+      )
+
+  @classmethod
+  @override
+  def supported_models(cls) -> list[str]:
+    return ['mock']
+
+  def generate_content(
+      self, llm_request: LlmRequest, stream: bool = False
+  ) -> Generator[LlmResponse, None, None]:
+    if self.error is not None:
+      raise self.error
+    # Increasement of the index has to happen before the yield.
+    self.response_index += 1
+    self.requests.append(llm_request)
+    # yield LlmResponse(content=self.responses[self.response_index])
+    yield self.responses[self.response_index]
+
+  @override
+  async def generate_content_async(
+      self, llm_request: LlmRequest, stream: bool = False
+  ) -> AsyncGenerator[LlmResponse, None]:
+    if self.error is not None:
+      raise self.error
+    # Increasement of the index has to happen before the yield.
+    self.response_index += 1
+    self.requests.append(llm_request)
+    yield self.responses[self.response_index]
+
+  @contextlib.asynccontextmanager
+  async def connect(self, llm_request: LlmRequest) -> BaseLlmConnection:
+    """Creates a live connection to the LLM."""
+    self.requests.append(llm_request)
+    yield MockLlmConnection(
+        self.responses,
+        self,
+        wait_for_realtime_input=self.wait_for_realtime_input,
+    )
+
+
+class MockLlmConnection(BaseLlmConnection):
+
+  def __init__(
+      self,
+      llm_responses: list[LlmResponse],
+      mock_model: MockModel,
+      wait_for_realtime_input: bool = False,
+  ):
+    self.llm_responses = llm_responses
+    self.mock_model = mock_model
+    self.wait_for_realtime_input = wait_for_realtime_input
+    self._input_event = asyncio.Event()
+
+  async def send_history(self, history: list[types.Content]):
+    pass
+
+  async def send_content(self, content: types.Content):
+    self.mock_model.live_contents.append(content)
+    self._input_event.set()
+
+  async def send(self, data):
+    pass
+
+  async def send_realtime(self, blob: types.Blob):
+    self.mock_model.live_blobs.append(blob)
+    self._input_event.set()
+
+  async def receive(self) -> AsyncGenerator[LlmResponse, None]:
+    """Yield each of the pre-defined LlmResponses."""
+    if self.wait_for_realtime_input:
+      await self._input_event.wait()
+
+    for response in self.llm_responses:
+      yield response
+
+  async def close(self):
+    pass

tests/unittests/workflow/utils/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/unittests/workflow/utils/test_graph_parser.py

@@ -0,0 +1,381 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for Graph parser utility."""
+
+from google.adk.workflow import Edge
+from google.adk.workflow import FunctionNode
+from google.adk.workflow import START
+from google.adk.workflow._graph import DEFAULT_ROUTE
+from google.adk.workflow.utils._graph_parser import parse_edge_items
+import pytest
+
+from ..workflow_testing_utils import TestingNode
+
+
+def test_parse_edge_items_with_node_reuse() -> None:
+  """Tests that node reuse during parsing works and deduplicates wrapped nodes."""
+
+  def my_node_func() -> None:
+    pass
+
+  node_b = TestingNode(name='NodeB')
+  edges = parse_edge_items([
+      (START, my_node_func),
+      (my_node_func, node_b),
+  ])
+
+  assert len(edges) == 2
+  edge1, edge2 = edges
+  assert edge1.from_node == START
+  assert isinstance(edge1.to_node, FunctionNode)
+  assert edge1.to_node.name == 'my_node_func'
+
+  assert isinstance(edge2.from_node, FunctionNode)
+  assert edge2.from_node.name == 'my_node_func'
+  assert edge2.to_node == node_b
+
+  # Verify exact same object instance was reused
+  assert edge1.to_node is edge2.from_node
+
+
+def test_routing_map_basic() -> None:
+  """Tests that a string-keyed routing map expands to correct edges."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  node_c = TestingNode(name='NodeC')
+  edges = parse_edge_items([
+      (START, node_a),
+      (node_a, {'route_b': node_b, 'route_c': node_c}),
+  ])
+
+  assert len(edges) == 3  # START->A, A->B(route_b), A->C(route_c)
+
+  routed_edges = [e for e in edges if e.route is not None]
+  assert len(routed_edges) == 2
+
+  routes_and_targets = {(e.route, e.to_node.name) for e in routed_edges}
+  assert routes_and_targets == {('route_b', 'NodeB'), ('route_c', 'NodeC')}
+
+  for e in routed_edges:
+    assert e.from_node.name == 'NodeA'
+
+
+def test_routing_map_int_keys() -> None:
+  """Tests that integer route keys work in routing maps."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  node_c = TestingNode(name='NodeC')
+  edges = parse_edge_items([
+      (START, node_a),
+      (node_a, {1: node_b, 2: node_c}),
+  ])
+
+  routed_edges = [e for e in edges if e.route is not None]
+  assert len(routed_edges) == 2
+  routes = [e.route for e in routed_edges]
+  assert 1 in routes
+  assert 2 in routes
+
+
+def test_routing_map_bool_keys() -> None:
+  """Tests that boolean route keys work in routing maps."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  node_c = TestingNode(name='NodeC')
+  edges = parse_edge_items([
+      (START, node_a),
+      (node_a, {True: node_b, False: node_c}),
+  ])
+
+  routed_edges = [e for e in edges if e.route is not None]
+  assert len(routed_edges) == 2
+  routes = [e.route for e in routed_edges]
+  assert True in routes
+  assert False in routes
+
+
+def test_routing_map_with_fan_in_source() -> None:
+  """Tests that fan-in on the source side works with routing maps."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  node_c = TestingNode(name='NodeC')
+  node_d = TestingNode(name='NodeD')
+  edges = parse_edge_items([
+      (START, node_a),
+      (START, node_b),
+      ((node_a, node_b), {'route_x': node_c, 'route_y': node_d}),
+  ])
+
+  # 2 from START + 4 from fan-in (A->C, A->D, B->C, B->D)
+  assert len(edges) == 6
+
+  fan_in_edges = [e for e in edges if e.from_node.name in ('NodeA', 'NodeB')]
+  assert len(fan_in_edges) == 4
+
+  combos = {(e.from_node.name, e.to_node.name, e.route) for e in fan_in_edges}
+  assert combos == {
+      ('NodeA', 'NodeC', 'route_x'),
+      ('NodeA', 'NodeD', 'route_y'),
+      ('NodeB', 'NodeC', 'route_x'),
+      ('NodeB', 'NodeD', 'route_y'),
+  }
+
+
+def test_routing_map_with_callable_target() -> None:
+  """Tests that callable values in routing maps get wrapped via build_node."""
+  node_a = TestingNode(name='NodeA')
+
+  def my_target_func() -> None:
+    pass
+
+  edges = parse_edge_items([
+      (START, node_a),
+      (node_a, {'route_x': my_target_func}),
+  ])
+
+  target_edge = next(e for e in edges if e.route == 'route_x')
+  assert isinstance(target_edge.to_node, FunctionNode)
+  assert target_edge.to_node.name == 'my_target_func'
+
+
+def test_routing_map_node_reuse() -> None:
+  """Tests that the same callable used in a map and elsewhere is deduplicated."""
+
+  def my_func() -> None:
+    pass
+
+  node_b = TestingNode(name='NodeB')
+  edges = parse_edge_items([
+      (START, my_func),
+      (my_func, {'route_x': node_b}),
+  ])
+
+  # my_func should be wrapped once and reused.
+  assert len(edges) == 2
+  assert edges[0].to_node is edges[1].from_node
+  assert isinstance(edges[0].to_node, FunctionNode)
+
+
+def test_routing_map_empty_dict_raises() -> None:
+  """Tests that an empty routing map raises ValueError."""
+  node_a = TestingNode(name='NodeA')
+  with pytest.raises(
+      ValueError,
+      match=r'Routing map must not be empty',
+  ):
+    parse_edge_items([
+        (START, node_a),
+        (node_a, {}),
+    ])
+
+
+def test_routing_map_invalid_key_raises() -> None:
+  """Tests that a non-RouteValue key raises ValueError."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  with pytest.raises(
+      ValueError,
+      match=r'Invalid routing map key',
+  ):
+    parse_edge_items([
+        (START, node_a),
+        (node_a, {1.5: node_b}),
+    ])
+
+
+def test_routing_map_invalid_value_raises() -> None:
+  """Tests that a non-NodeLike value raises ValueError."""
+  node_a = TestingNode(name='NodeA')
+  with pytest.raises(
+      ValueError,
+      match=r'Invalid routing map value',
+  ):
+    parse_edge_items([
+        (START, node_a),
+        (node_a, {'route_x': 42}),
+    ])
+
+
+def test_routing_map_fan_out_target() -> None:
+  """Tests that a tuple value in a routing map creates fan-out edges."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  node_c = TestingNode(name='NodeC')
+  edges = parse_edge_items([
+      (START, node_a),
+      (node_a, {'route_x': (node_b, node_c)}),
+  ])
+
+  # START->A, A->B(route_x), A->C(route_x)
+  assert len(edges) == 3
+
+  routed_edges = [e for e in edges if e.route is not None]
+  assert len(routed_edges) == 2
+
+  # Both fan-out edges share the same route and source.
+  for e in routed_edges:
+    assert e.from_node.name == 'NodeA'
+    assert e.route == 'route_x'
+
+  targets = {e.to_node.name for e in routed_edges}
+  assert targets == {'NodeB', 'NodeC'}
+
+
+def test_routing_map_fan_out_invalid_element_raises() -> None:
+  """Tests that a non-NodeLike element inside a fan-out tuple raises."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  with pytest.raises(
+      ValueError,
+      match=r'Invalid node in fan-out tuple',
+  ):
+    parse_edge_items([
+        (START, node_a),
+        (node_a, {'route_x': (node_b, 42)}),
+    ])
+
+
+# --- Routing map as chain element tests ---
+
+
+def test_routing_map_chain_ending_with_dict() -> None:
+  """Tests a chain ending with a routing map creates correct edges."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  node_c = TestingNode(name='NodeC')
+  edges = parse_edge_items([
+      (START, node_a, {'r1': node_b, 'r2': node_c}),
+  ])
+
+  # START->A (None), A->B (r1), A->C (r2)
+  assert len(edges) == 3
+
+  start_edge = next(e for e in edges if e.from_node.name == '__START__')
+  assert start_edge.to_node.name == 'NodeA'
+  assert start_edge.route is None
+
+  routed_edges = [e for e in edges if e.route is not None]
+  assert len(routed_edges) == 2
+  routes_and_targets = {(e.route, e.to_node.name) for e in routed_edges}
+  assert routes_and_targets == {('r1', 'NodeB'), ('r2', 'NodeC')}
+  for e in routed_edges:
+    assert e.from_node.name == 'NodeA'
+
+
+def test_routing_map_mid_chain_with_fan_in() -> None:
+  """Tests routing map mid-chain with fan-in to the next element."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  node_c = TestingNode(name='NodeC')
+  node_d = TestingNode(name='NodeD')
+  edges = parse_edge_items([
+      (START, node_a, {'r1': node_b, 'r2': node_c}, node_d),
+  ])
+
+  # START->A (None), A->B (r1), A->C (r2), B->D (None), C->D (None)
+  assert len(edges) == 5
+
+  routed_edges = sorted(
+      [e for e in edges if e.route is not None],
+      key=lambda e: e.to_node.name,
+  )
+  assert len(routed_edges) == 2
+  assert routed_edges[0].from_node.name == 'NodeA'
+  assert routed_edges[0].to_node.name == 'NodeB'
+  assert routed_edges[0].route == 'r1'
+  assert routed_edges[1].from_node.name == 'NodeA'
+  assert routed_edges[1].to_node.name == 'NodeC'
+  assert routed_edges[1].route == 'r2'
+
+  fan_in_edges = sorted(
+      [e for e in edges if e.to_node.name == 'NodeD'],
+      key=lambda e: e.from_node.name,
+  )
+  assert len(fan_in_edges) == 2
+  assert fan_in_edges[0].from_node.name == 'NodeB'
+  assert fan_in_edges[0].route is None
+  assert fan_in_edges[1].from_node.name == 'NodeC'
+  assert fan_in_edges[1].route is None
+
+
+def test_routing_map_mid_chain_fan_out_values() -> None:
+  """Tests routing map with fan-out tuple values, followed by fan-in."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  node_c = TestingNode(name='NodeC')
+  node_d = TestingNode(name='NodeD')
+  edges = parse_edge_items([
+      (START, node_a, {'r1': (node_b, node_c)}, node_d),
+  ])
+
+  # START->A (None), A->B (r1), A->C (r1), B->D (None), C->D (None)
+  assert len(edges) == 5
+
+  routed_edges = [e for e in edges if e.route is not None]
+  assert len(routed_edges) == 2
+  for e in routed_edges:
+    assert e.from_node.name == 'NodeA'
+    assert e.route == 'r1'
+
+  fan_in_edges = [e for e in edges if e.to_node.name == 'NodeD']
+  assert len(fan_in_edges) == 2
+  fan_in_sources = {e.from_node.name for e in fan_in_edges}
+  assert fan_in_sources == {'NodeB', 'NodeC'}
+
+
+def test_routing_map_consecutive_dicts_raises() -> None:
+  """Tests that consecutive routing maps in a chain are rejected."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  node_c = TestingNode(name='NodeC')
+  node_d = TestingNode(name='NodeD')
+  with pytest.raises(
+      ValueError, match=r'Consecutive routing maps are not allowed'
+  ):
+    parse_edge_items([
+        (START, node_a, {'r1': node_b, 'r2': node_c}, {'r3': node_d}),
+    ])
+
+
+def test_routing_map_empty_dict_in_chain_raises() -> None:
+  """Tests that an empty routing map in a chain raises ValueError."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  with pytest.raises(ValueError, match=r'Routing map must not be empty'):
+    parse_edge_items([
+        (START, node_a, {}, node_b),
+    ])
+
+
+def test_routing_map_invalid_key_in_chain_raises() -> None:
+  """Tests that invalid routing map keys in a chain raise ValueError."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  with pytest.raises(ValueError, match=r'Invalid routing map key'):
+    parse_edge_items([
+        (START, node_a, {1.5: node_b}),
+    ])
+
+
+def test_routing_map_2_tuple_backward_compat() -> None:
+  """Ensures existing 2-tuple routing map syntax still works."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  node_c = TestingNode(name='NodeC')
+  edges = parse_edge_items([
+      (START, node_a),
+      (node_a, {'r1': node_b, 'r2': node_c}),
+  ])
+  assert len(edges) == 3

tests/unittests/workflow/utils/test_graph_validation.py

@@ -0,0 +1,390 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for Graph validation utility."""
+
+import logging
+
+from google.adk.workflow import Edge
+from google.adk.workflow import START
+from google.adk.workflow._graph import DEFAULT_ROUTE
+from google.adk.workflow._graph import Graph
+from google.adk.workflow.utils._graph_validation import validate_graph
+from pydantic import BaseModel
+import pytest
+
+from ..workflow_testing_utils import TestingNode
+
+
+def test_missing_start_node() -> None:
+  """Tests that a graph missing the START node fails validation."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  graph = Graph(
+      edges=[Edge(from_node=node_a, to_node=node_b)],
+  )
+  with pytest.raises(
+      ValueError,
+      match=(
+          r"Graph validation failed\. START node \(name: '__START__'\) not"
+          r' found in graph nodes\.'
+      ),
+  ):
+    validate_graph(graph.nodes, graph.edges)
+
+
+def test_unreachable_node() -> None:
+  """Tests that a graph with an unreachable node fails validation."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')  # Unreachable
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(from_node=node_b, to_node=node_a),
+      ],
+  )
+  with pytest.raises(
+      ValueError,
+      match=(
+          r'Graph validation failed\. The following nodes are unreachable'
+          r" from START: \['NodeB'\]"
+      ),
+  ):
+    validate_graph(graph.nodes, graph.edges)
+
+
+def test_disconnected_routed_subgraph_is_unreachable() -> None:
+  """Tests that a disconnected subgraph with routed edges fails validation."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  node_c = TestingNode(name='NodeC')
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(from_node=node_b, to_node=node_c, route='x'),
+          Edge(from_node=node_c, to_node=node_b, route='y'),
+      ],
+  )
+  with pytest.raises(
+      ValueError,
+      match=(
+          r'Graph validation failed\. The following nodes are unreachable'
+          r" from START: \['NodeB', 'NodeC'\]"
+      ),
+  ):
+    validate_graph(graph.nodes, graph.edges)
+
+
+@pytest.mark.parametrize(
+    'routes',
+    [
+        (None, None),
+        ('route1', 'route1'),
+        ('route1', 'route2'),
+        ('route1', None),
+    ],
+)
+def test_duplicate_edges_fail_validation(
+    routes: tuple[str | None, str | None],
+) -> None:
+  """Tests that duplicate edges fail validation, regardless of routes."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(
+              from_node=node_a,
+              to_node=node_b,
+              route=routes[0],
+          ),
+          Edge(
+              from_node=node_a,
+              to_node=node_b,
+              route=routes[1],
+          ),
+      ],
+  )
+  with pytest.raises(
+      ValueError,
+      match=(
+          r'Graph validation failed\. Duplicate edge found: from=NodeA,'
+          r' to=NodeB'
+      ),
+  ):
+    validate_graph(graph.nodes, graph.edges)
+
+
+def test_routed_start_edge_fails_validation() -> None:
+  """Tests that routed edges from START node fail validation."""
+  node_a = TestingNode(name='NodeA')
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a, route='route1'),
+      ],
+  )
+  with pytest.raises(
+      ValueError,
+      match=r'Graph validation failed\. Edges from START must not have routes',
+  ):
+    validate_graph(graph.nodes, graph.edges)
+
+
+def test_start_node_with_incoming_edge() -> None:
+  """Tests graph with incoming edge to START node fails validation."""
+  node_a = TestingNode(name='NodeA')
+  graph = Graph(
+      edges=[
+          Edge(from_node=node_a, to_node=START),
+          Edge(from_node=START, to_node=node_a),
+      ],
+  )
+  with pytest.raises(
+      ValueError,
+      match=(
+          r'Graph validation failed\. START node must not have incoming edges\.'
+      ),
+  ):
+    validate_graph(graph.nodes, graph.edges)
+
+
+def test_multiple_default_routes_fail_validation() -> None:
+  """Tests that multiple DEFAULT_ROUTE edges from a node fail validation."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  node_c = TestingNode(name='NodeC')
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(from_node=node_a, to_node=node_b, route=DEFAULT_ROUTE),
+          Edge(from_node=node_a, to_node=node_c, route=DEFAULT_ROUTE),
+      ],
+  )
+  with pytest.raises(
+      ValueError,
+      match=(
+          r'Graph validation failed\. Multiple DEFAULT_ROUTE edges found from'
+          r' node NodeA to NodeB and NodeC'
+      ),
+  ):
+    validate_graph(graph.nodes, graph.edges)
+
+
+def test_single_default_route_passes_validation() -> None:
+  """Tests that a single DEFAULT_ROUTE edge from a node passes validation."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  node_c = TestingNode(name='NodeC')
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(from_node=node_a, to_node=node_b, route=DEFAULT_ROUTE),
+          Edge(from_node=node_a, to_node=node_c, route='another_route'),
+      ],
+  )
+  validate_graph(graph.nodes, graph.edges)  # Should not raise
+
+
+def test_duplicate_node_names_fail_validation() -> None:
+  """Tests that duplicate nodes raise error."""
+  node_a1 = TestingNode(name='NodeA')
+  node_a2 = TestingNode(name='NodeA')
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a1),
+          Edge(from_node=node_a1, to_node=node_a2),
+      ],
+  )
+  with pytest.raises(
+      ValueError,
+      match=(
+          r"Graph validation failed\. Duplicate node names found: \['NodeA'\]\."
+          r' This means multiple distinct node objects have the same name\. If'
+          r' you intended to reuse the same node, ensure you pass the exact'
+          r' same object instance\. If you intended to have distinct nodes,'
+          r' ensure they have unique names\.'
+      ),
+  ):
+    validate_graph(graph.nodes, graph.edges)
+
+
+def test_unconditional_cycle_fails_validation() -> None:
+  """Tests that a cycle of unconditional edges (route=None) fails."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(from_node=node_a, to_node=node_b),
+          Edge(from_node=node_b, to_node=node_a),
+      ],
+  )
+  with pytest.raises(
+      ValueError,
+      match=r'Graph validation failed\. Unconditional cycle detected:',
+  ):
+    validate_graph(graph.nodes, graph.edges)
+
+
+def test_unconditional_self_loop_fails_validation() -> None:
+  """Tests that an unconditional self-loop (A -> A) fails."""
+  node_a = TestingNode(name='NodeA')
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(from_node=node_a, to_node=node_a),
+      ],
+  )
+  with pytest.raises(
+      ValueError,
+      match=r'Graph validation failed\. Unconditional cycle detected:',
+  ):
+    validate_graph(graph.nodes, graph.edges)
+
+
+def test_longer_unconditional_cycle_fails_validation() -> None:
+  """Tests that a longer unconditional cycle (A -> B -> C -> A) fails."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  node_c = TestingNode(name='NodeC')
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(from_node=node_a, to_node=node_b),
+          Edge(from_node=node_b, to_node=node_c),
+          Edge(from_node=node_c, to_node=node_a),
+      ],
+  )
+  with pytest.raises(
+      ValueError,
+      match=r'Graph validation failed\. Unconditional cycle detected:',
+  ):
+    validate_graph(graph.nodes, graph.edges)
+
+
+def test_conditional_cycle_passes_validation() -> None:
+  """Tests that a cycle with a routed edge (loop pattern) passes."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(from_node=node_a, to_node=node_b),
+          Edge(from_node=node_b, to_node=node_a, route='retry'),
+      ],
+  )
+  validate_graph(
+      graph.nodes, graph.edges
+  )  # Should not raise — routed back-edge
+
+
+def test_conditional_self_loop_passes_validation() -> None:
+  """Tests that a self-loop with a route passes validation."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(from_node=node_a, to_node=node_a, route='continue'),
+          Edge(from_node=node_a, to_node=node_b, route='done'),
+      ],
+  )
+  validate_graph(
+      graph.nodes, graph.edges
+  )  # Should not raise — routed self-loop
+
+
+def test_dag_with_diamond_passes_validation() -> None:
+  """Tests that a DAG with a diamond shape passes validation."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  node_c = TestingNode(name='NodeC')
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(from_node=START, to_node=node_b),
+          Edge(from_node=node_a, to_node=node_c),
+          Edge(from_node=node_b, to_node=node_c),
+      ],
+  )
+  validate_graph(graph.nodes, graph.edges)  # Should not raise
+
+
+class ModelA(BaseModel):
+  x: int
+
+
+class ModelB(BaseModel):
+  x: int
+
+
+def test_schema_match_passes() -> None:
+  """Tests that edges with matching schemas pass validation."""
+  node_a = TestingNode(name='NodeA', output_schema=ModelA)
+  node_b = TestingNode(name='NodeB', input_schema=ModelA)
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(from_node=node_a, to_node=node_b),
+      ],
+  )
+  validate_graph(graph.nodes, graph.edges)  # Should not raise
+
+
+def test_schema_mismatch_raises() -> None:
+  """Tests that edges with mismatching schemas fail validation."""
+  node_a = TestingNode(name='NodeA', output_schema=ModelA)
+  node_b = TestingNode(name='NodeB', input_schema=ModelB)
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(from_node=node_a, to_node=node_b),
+      ],
+  )
+  with pytest.raises(
+      ValueError,
+      match=r'Graph validation failed\. Schema mismatch on edge',
+  ):
+    validate_graph(graph.nodes, graph.edges)
+
+
+def test_schema_missing_passes() -> None:
+  """Tests that edges with missing schemas pass validation."""
+  node_a = TestingNode(name='NodeA', output_schema=ModelA)
+  node_b = TestingNode(name='NodeB')  # No input schema
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(from_node=node_a, to_node=node_b),
+      ],
+  )
+  validate_graph(graph.nodes, graph.edges)  # Should not raise
+
+
+def test_chat_agent_wiring_validation_only_runs_on_llm_agent() -> None:
+  """Tests that _validate_chat_agent_wiring checks non-LlmAgent nodes safely."""
+  node_a = TestingNode(name='NodeA')
+  node_b = TestingNode(name='NodeB')
+  # Set mode='chat' on a non-LlmAgent node
+  object.__setattr__(node_b, 'mode', 'chat')
+
+  graph = Graph(
+      edges=[
+          Edge(from_node=START, to_node=node_a),
+          Edge(from_node=node_a, to_node=node_b),
+      ],
+  )
+  validate_graph(
+      graph.nodes, graph.edges
+  )  # Should not raise because node_b is a TestingNode, not LlmAgent

tests/unittests/workflow/utils/test_rehydration_utils.py

@@ -0,0 +1,389 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from google.adk.events.event import Event
+from google.adk.events.event import NodeInfo
+from google.adk.events.request_input import RequestInput
+from google.adk.workflow._base_node import BaseNode
+from google.adk.workflow.utils._rehydration_utils import _ChildScanState
+from google.adk.workflow.utils._rehydration_utils import _process_rehydrated_output
+from google.adk.workflow.utils._rehydration_utils import _reconstruct_node_states
+from google.adk.workflow.utils._rehydration_utils import _unwrap_response
+from google.adk.workflow.utils._rehydration_utils import _validate_resume_response
+from google.adk.workflow.utils._rehydration_utils import _wrap_response
+from google.adk.workflow.utils._workflow_hitl_utils import create_request_input_event
+from google.genai import types
+from pydantic import BaseModel
+import pytest
+
+# --- _wrap_response ---
+
+
+class TestWrapResponse:
+
+  def test_dict_returned_as_is(self):
+    d = {"foo": "bar"}
+    assert _wrap_response(d) is d
+
+  def test_string_wrapped(self):
+    assert _wrap_response("hello") == {"result": "hello"}
+
+  def test_int_wrapped(self):
+    assert _wrap_response(42) == {"result": 42}
+
+  def test_none_wrapped(self):
+    assert _wrap_response(None) == {"result": None}
+
+  def test_list_wrapped(self):
+    assert _wrap_response([1, 2]) == {"result": [1, 2]}
+
+
+# --- _unwrap_response ---
+
+
+class TestUnwrapResponse:
+
+  def test_single_result_key_string(self):
+    assert _unwrap_response({"result": "hello"}) == "hello"
+
+  def test_single_result_key_int(self):
+    assert _unwrap_response({"result": 42}) == 42
+
+  def test_single_result_key_none(self):
+    assert _unwrap_response({"result": None}) is None
+
+  def test_dict_without_result_key_unchanged(self):
+    d = {"foo": "bar"}
+    assert _unwrap_response(d) == {"foo": "bar"}
+
+  def test_dict_with_multiple_keys_unchanged(self):
+    d = {"result": "x", "other": "y"}
+    assert _unwrap_response(d) == {"result": "x", "other": "y"}
+
+  def test_non_dict_unchanged(self):
+    assert _unwrap_response("hello") == "hello"
+    assert _unwrap_response(42) == 42
+    assert _unwrap_response(None) is None
+
+  def test_json_string_parsed_to_dict(self):
+    """Web frontend sends {"result": '{"approved": false}'}."""
+    assert _unwrap_response({"result": '{"approved": false}'}) == {
+        "approved": False
+    }
+
+  def test_json_string_parsed_to_list(self):
+    assert _unwrap_response({"result": "[1, 2, 3]"}) == [1, 2, 3]
+
+  def test_json_string_parsed_to_number(self):
+    assert _unwrap_response({"result": "42"}) == 42
+
+  def test_json_string_parsed_to_bool(self):
+    assert _unwrap_response({"result": "true"}) is True
+
+  def test_non_json_string_stays_string(self):
+    assert _unwrap_response({"result": "plain text"}) == "plain text"
+
+  def test_roundtrip_wrap_unwrap_string(self):
+    assert _unwrap_response(_wrap_response("hello")) == "hello"
+
+  def test_roundtrip_wrap_unwrap_dict(self):
+    """Dicts are not wrapped, so unwrap is a no-op."""
+    d = {"foo": "bar"}
+    assert _unwrap_response(_wrap_response(d)) == d
+
+
+# --- _process_rehydrated_output ---
+
+
+class TestProcessRehydratedOutput:
+
+  def test_extracts_plain_text_without_schema(self):
+    node = BaseNode(name="dummy")
+    content = types.Content(parts=[types.Part(text="hello world")])
+    assert _process_rehydrated_output(node, content) == "hello world"
+
+  def test_returns_plain_text_even_if_json_when_no_schema(self):
+    node = BaseNode(name="dummy")
+    content = types.Content(parts=[types.Part(text='{"foo": "bar"}')])
+    assert _process_rehydrated_output(node, content) == '{"foo": "bar"}'
+
+  def test_parses_json_text_with_output_schema(self):
+    class MySchema(BaseModel):
+      foo: str
+
+    node = BaseNode(name="dummy", output_schema=MySchema)
+    content = types.Content(parts=[types.Part(text='{"foo": "bar"}')])
+    assert _process_rehydrated_output(node, content) == {"foo": "bar"}
+
+  def test_joins_multiple_parts(self):
+    node = BaseNode(name="dummy")
+    content = types.Content(
+        parts=[types.Part(text="hello "), types.Part(text="world")]
+    )
+    assert _process_rehydrated_output(node, content) == "hello world"
+
+  def test_filters_thought_parts(self):
+    class MySchema(BaseModel):
+      answer: int
+
+    node = BaseNode(name="dummy", output_schema=MySchema)
+    content = types.Content(
+        parts=[
+            types.Part(text="thinking...", thought=True),
+            types.Part(text='{"answer": 42}'),
+        ]
+    )
+    assert _process_rehydrated_output(node, content) == {"answer": 42}
+
+  def test_returns_none_for_empty_text(self):
+    node = BaseNode(name="dummy")
+    content = types.Content(parts=[types.Part(text="  ")])
+    assert _process_rehydrated_output(node, content) is None
+
+  def test_gracefully_falls_back_on_schema_mismatch(self, caplog):
+    class MySchema(BaseModel):
+      foo: str
+      bar: int  # Required field that is missing in the stored output
+
+    node = BaseNode(name="dummy", output_schema=MySchema)
+    content = types.Content(parts=[types.Part(text='{"foo": "only"}')])
+
+    # Should NOT raise ValueError, but fallback to unvalidated parsed dict
+    res = _process_rehydrated_output(node, content)
+    assert res == {"foo": "only"}
+    assert (
+        "Validation failed for rehydrated output against schema" in caplog.text
+    )
+
+  def test_raises_value_error_if_not_valid_json_on_schema_mismatch(self):
+    class MySchema(BaseModel):
+      foo: str
+
+    node = BaseNode(name="dummy", output_schema=MySchema)
+    content = types.Content(parts=[types.Part(text="invalid json")])
+
+    # Should raise ValueError because it's not valid JSON
+    with pytest.raises(
+        ValueError,
+        match="Validation failed for rehydrated output against schema",
+    ):
+      _process_rehydrated_output(node, content)
+
+
+# --- _validate_resume_response ---
+
+
+class TestValidateResumeResponse:
+
+  def test_none_schema_returns_data(self):
+    assert _validate_resume_response("hello", None) == "hello"
+
+  def test_str_to_int_coercion(self):
+    assert _validate_resume_response("42", {"type": "integer"}) == 42
+
+  def test_str_to_float_coercion(self):
+    assert _validate_resume_response("42.5", {"type": "number"}) == 42.5
+
+  def test_str_to_bool_true(self):
+    assert _validate_resume_response("true", {"type": "boolean"}) is True
+    assert _validate_resume_response("1", {"type": "boolean"}) is True
+
+  def test_str_to_bool_false(self):
+    assert _validate_resume_response("false", {"type": "boolean"}) is False
+    assert _validate_resume_response("0", {"type": "boolean"}) is False
+
+  def test_invalid_coercion_raises_value_error(self):
+    with pytest.raises(ValueError):
+      _validate_resume_response("abc", {"type": "integer"})
+
+  def test_object_schema_validates_dict_type(self):
+    schema = {"type": "object"}
+    assert _validate_resume_response({"name": "Alice"}, schema) == {
+        "name": "Alice"
+    }
+
+    with pytest.raises(ValueError, match="Failed to coerce data to object"):
+      _validate_resume_response("not a dict", schema)
+
+  def test_array_schema_validates_list_type(self):
+    schema = {"type": "array"}
+    assert _validate_resume_response([1, 2], schema) == [1, 2]
+
+    with pytest.raises(ValueError, match="Failed to coerce data to array"):
+      _validate_resume_response("not a list", schema)
+
+  def test_pydantic_type_validation(self):
+    class User(BaseModel):
+      name: str
+      age: int
+
+    assert _validate_resume_response(
+        {"name": "Alice", "age": 30}, User
+    ) == User(name="Alice", age=30)
+
+
+# --- _reconstruct_node_states ---
+
+
+class TestScanNodeEvents:
+
+  def test_scan_empty_events(self):
+    results = _reconstruct_node_states([], "/wf@1", invocation_id="test_id")
+    assert results == {}
+
+  def test_scan_direct_child_output(self):
+    event = Event(
+        node_info=NodeInfo(path="/wf@1/node_a@1"),
+        output="node_a output",
+        invocation_id="test_id",
+    )
+    results = _reconstruct_node_states(
+        [event], "/wf@1", invocation_id="test_id", group_by_direct_child=True
+    )
+
+    assert "node_a@1" in results
+    assert results["node_a@1"].output == "node_a output"
+    assert results["node_a@1"].run_id == "1"
+
+  def test_scan_message_as_output(self):
+    content = types.Content(parts=[types.Part(text="hello")])
+    event = Event(
+        node_info=NodeInfo(path="/wf@1/node_a@1"),
+        content=content,
+        invocation_id="test_id",
+    )
+    event.node_info.message_as_output = True
+
+    results = _reconstruct_node_states(
+        [event], "/wf@1", invocation_id="test_id", group_by_direct_child=True
+    )
+
+    assert "node_a@1" in results
+    assert results["node_a@1"].output == content
+
+  def test_scan_descendant_interrupts(self):
+    event = Event(
+        node_info=NodeInfo(path="/wf@1/node_a@1/sub_node@1"),
+        long_running_tool_ids={"interrupt-1"},
+        invocation_id="test_id",
+    )
+    results = _reconstruct_node_states(
+        [event], "/wf@1", invocation_id="test_id", group_by_direct_child=True
+    )
+
+    assert "node_a@1" in results
+    assert "interrupt-1" in results["node_a@1"].interrupt_ids
+
+  def test_scan_resolve_interrupts(self):
+    event_int = Event(
+        node_info=NodeInfo(path="/wf@1/node_a@1"),
+        long_running_tool_ids={"interrupt-1"},
+        invocation_id="test_id",
+    )
+    event_fr = Event(
+        author="user",
+        content=types.Content(
+            parts=[
+                types.Part(
+                    function_response=types.FunctionResponse(
+                        id="interrupt-1",
+                        name="adk_request_input",
+                        response={"result": "user answer"},
+                    )
+                )
+            ]
+        ),
+        invocation_id="test_id",
+    )
+
+    # Act
+    results = _reconstruct_node_states(
+        [event_int, event_fr],
+        "/wf@1",
+        invocation_id="test_id",
+        group_by_direct_child=True,
+    )
+
+    # Assert
+    assert "node_a@1" in results
+    assert "interrupt-1" in results["node_a@1"].resolved_ids
+    assert (
+        results["node_a@1"].resolved_responses["interrupt-1"] == "user answer"
+    )
+
+  def test_scan_matches_specific_node_path_without_child_grouping(self):
+    """Scanning matches events for a specific node path when not grouping by direct child."""
+    event = Event(
+        node_info=NodeInfo(path="/wf@1/node_a@1"),
+        output="node_a output",
+        invocation_id="test_id",
+    )
+
+    # Act
+    results = _reconstruct_node_states(
+        [event],
+        "/wf@1/node_a@1",
+        invocation_id="test_id",
+        group_by_direct_child=False,
+    )
+
+    # Assert
+    assert "/wf@1/node_a@1" in results
+    assert results["/wf@1/node_a@1"].output == "node_a output"
+
+  def test_scan_validates_and_coerces_response_against_schema(self):
+    """Scanning validates and coerces user response data against the provided schema."""
+
+    class MySchema(BaseModel):
+      count: int
+
+    ri = RequestInput(
+        interrupt_id="interrupt-1",
+        response_schema=MySchema,
+    )
+    event_int = create_request_input_event(ri)
+    event_int.node_info = NodeInfo(path="/wf@1/node_a@1")
+    event_int.invocation_id = "test_id"
+
+    event_fr = Event(
+        author="user",
+        content=types.Content(
+            parts=[
+                types.Part(
+                    function_response=types.FunctionResponse(
+                        id="interrupt-1",
+                        name="adk_request_input",
+                        response={"result": '{"count": "42"}'},
+                    )
+                )
+            ]
+        ),
+        invocation_id="test_id",
+    )
+
+    # Act
+    results = _reconstruct_node_states(
+        [event_int, event_fr],
+        "/wf@1",
+        invocation_id="test_id",
+        group_by_direct_child=True,
+    )
+
+    # Assert
+    assert "node_a@1" in results
+    assert results["node_a@1"].resolved_responses["interrupt-1"] == {
+        "count": 42
+    }

tests/unittests/workflow/utils/test_replay_interceptor.py

@@ -0,0 +1,175 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for ReplayInterceptor.
+
+Verifies that ReplayInterceptor correctly checks and manages workflow resumption
+replay interception.
+"""
+
+from google.adk.workflow._base_node import BaseNode
+from google.adk.workflow._dynamic_node_scheduler import DynamicNodeRun
+from google.adk.workflow._node_state import NodeState
+from google.adk.workflow._node_status import NodeStatus
+from google.adk.workflow.utils._rehydration_utils import _ChildScanState
+from google.adk.workflow.utils._replay_interceptor import check_interception
+import pytest
+
+
+def test_same_turn_completed():
+  """Same-turn completed run intercepts and returns cached output."""
+  # Given a same-turn completed run
+  run = DynamicNodeRun(
+      state=NodeState(status=NodeStatus.COMPLETED),
+      output='cached-out',
+      transfer_to_agent='target-agent',
+  )
+
+  # When checked
+  result = check_interception(
+      node=BaseNode(name='node'),
+      current_run=run,
+  )
+
+  # Then it intercepts with cached results
+  assert not result.should_run
+  assert result.output == 'cached-out'
+  assert result.transfer_to_agent == 'target-agent'
+
+
+def test_same_turn_waiting():
+  """Same-turn waiting run intercepts and returns unresolved interrupts."""
+  # Given a same-turn waiting run
+  run = DynamicNodeRun(
+      state=NodeState(status=NodeStatus.WAITING, interrupts=['fc-1']),
+  )
+
+  # When checked
+  result = check_interception(
+      node=BaseNode(name='node'),
+      current_run=run,
+  )
+
+  # Then it intercepts and keeps waiting
+  assert not result.should_run
+  assert result.interrupts == {'fc-1'}
+
+
+def test_cross_turn_unresolved_interrupts_no_rerun():
+  """Cross-turn unresolved interrupts keep waiting without rerun."""
+  # Given unresolved interrupts and node without rerun_on_resume
+  recovered = _ChildScanState(
+      run_id='1',
+      interrupt_ids={'fc-1', 'fc-2'},
+      resolved_ids={'fc-1'},
+  )
+  node = BaseNode(name='node', rerun_on_resume=False)
+
+  # When checked
+  result = check_interception(
+      node=node,
+      recovered=recovered,
+  )
+
+  # Then it stays waiting on unresolved interrupts
+  assert not result.should_run
+  assert result.interrupts == {'fc-2'}
+
+
+def test_cross_turn_unresolved_interrupts_rerun():
+  """Cross-turn unresolved interrupts with rerun resolves progress and reruns."""
+  # Given unresolved interrupts and node with rerun_on_resume
+  recovered = _ChildScanState(
+      run_id='1',
+      interrupt_ids={'fc-1', 'fc-2'},
+      resolved_ids={'fc-1'},
+      resolved_responses={'fc-1': 'ans'},
+  )
+  node = BaseNode(name='node', rerun_on_resume=True)
+
+  # When checked
+  result = check_interception(
+      node=node,
+      recovered=recovered,
+  )
+
+  # Then it reruns with partial resolved inputs
+  assert result.should_run
+  assert result.resume_inputs == {'fc-1': 'ans'}
+
+
+def test_cross_turn_completed():
+  """Cross-turn completed run fast-forwards output and route."""
+  # Given a completed run from history
+  recovered = _ChildScanState(
+      run_id='1',
+      output='past-out',
+      route='route-a',
+  )
+  node = BaseNode(name='node')
+
+  # When checked
+  result = check_interception(
+      node=node,
+      recovered=recovered,
+  )
+
+  # Then it fast-forwards with cached output and route
+  assert not result.should_run
+  assert result.output == 'past-out'
+  assert result.route == 'route-a'
+
+
+def test_cross_turn_all_resolved_no_rerun():
+  """Cross-turn all resolved run without rerun auto-completes with responses."""
+  # Given all resolved interrupts and node without rerun_on_resume
+  recovered = _ChildScanState(
+      run_id='1',
+      interrupt_ids={'fc-1'},
+      resolved_ids={'fc-1'},
+      resolved_responses={'fc-1': 'ans'},
+  )
+  node = BaseNode(name='node', rerun_on_resume=False)
+
+  # When checked
+  result = check_interception(
+      node=node,
+      recovered=recovered,
+  )
+
+  # Then it auto-completes
+  assert not result.should_run
+  assert result.output == 'ans'
+
+
+def test_cross_turn_all_resolved_rerun():
+  """Cross-turn all resolved run with rerun triggers rerun with responses."""
+  # Given all resolved interrupts and node with rerun_on_resume
+  recovered = _ChildScanState(
+      run_id='1',
+      interrupt_ids={'fc-1'},
+      resolved_ids={'fc-1'},
+      resolved_responses={'fc-1': 'ans'},
+  )
+  node = BaseNode(name='node', rerun_on_resume=True)
+
+  # When checked
+  result = check_interception(
+      node=node,
+      recovered=recovered,
+  )
+
+  # Then it reruns
+  assert result.should_run
+  assert result.resume_inputs == {'fc-1': 'ans'}

tests/unittests/workflow/utils/test_replay_manager.py

@@ -0,0 +1,102 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for ReplayManager utility."""
+
+from unittest.mock import MagicMock
+
+from google.adk.events.event import Event
+from google.adk.events.event import NodeInfo
+from google.adk.workflow.utils._replay_manager import ReplayManager
+import pytest
+
+
+def test_replay_manager_init() -> None:
+  """Tests that ReplayManager initializes with empty state."""
+  mgr = ReplayManager()
+  assert mgr.recovered_executions == {}
+  assert mgr.sequence_barrier is None
+
+
+def _make_event(
+    path='', output=None, interrupt_ids=None, invocation_id='inv-1'
+):
+  """Create a minimal Event for session event lists."""
+  event = MagicMock(spec=Event)
+  event.invocation_id = invocation_id
+  event.author = 'node'
+  event.output = output
+  event.partial = False
+  event.node_info = MagicMock(spec=NodeInfo)
+  event.node_info.path = path
+  event.node_info.output_for = None
+  event.node_info.message_as_output = None
+  event.branch = None
+  event.isolation_scope = None
+  event.long_running_tool_ids = set(interrupt_ids) if interrupt_ids else None
+  event.content = None
+  event.actions = None
+  return event
+
+
+@pytest.mark.asyncio
+async def test_scan_workflow_events():
+  """Scan workflow events populates recovered_executions and sequence_barrier."""
+  mgr = ReplayManager()
+  events = [
+      _make_event(path='wf/child1@1', output='out1'),
+      _make_event(path='wf/child2@1', output='out2'),
+  ]
+  ctx = MagicMock()
+  ctx._invocation_context = MagicMock()
+  ctx._invocation_context.invocation_id = 'inv-1'
+  ctx._invocation_context.session = MagicMock()
+  ctx._invocation_context.session.events = events
+  ctx.node_path = 'wf'
+
+  recovered, sequence = mgr.scan_workflow_events(ctx)
+
+  assert 'child1@1' in recovered
+  assert 'child2@1' in recovered
+  assert sequence == ['child1@1', 'child2@1']
+  assert mgr.sequence_barrier is not None
+
+
+@pytest.mark.asyncio
+async def test_scan_child_events_ignores_descendant_run_id_resets():
+  """scan_workflow_events only resets run_id from direct child events."""
+  mgr = ReplayManager()
+
+  event1 = Event(
+      author='node',
+      node_info=NodeInfo(path='wf@1/child@1', run_id='1'),
+      invocation_id='test_inv',
+  )
+  event2 = Event(
+      author='node',
+      node_info=NodeInfo(path='wf@1/child@1/grandchild@2', run_id='2'),
+      invocation_id='test_inv',
+  )
+
+  ctx = MagicMock()
+  ctx._invocation_context = MagicMock()
+  ctx._invocation_context.invocation_id = 'test_inv'
+  ctx._invocation_context.session = MagicMock()
+  ctx._invocation_context.session.events = [event1, event2]
+  ctx.node_path = 'wf@1'
+
+  children, _ = mgr.scan_workflow_events(ctx)
+
+  # Assert child 'child' run_id remains '1' (not '2' from the descendant).
+  assert children['child@1'].run_id == '1'

tests/unittests/workflow/utils/test_replay_sequence_barrier.py

@@ -0,0 +1,108 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for ReplaySequenceBarrier."""
+
+from __future__ import annotations
+
+import asyncio
+
+from google.adk.workflow.utils._replay_sequence_barrier import ReplaySequenceBarrier
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_barrier_initialization():
+  """Verifies that barrier initializes sequence index and sets the first event."""
+  # Given a chronological sequence of completions
+  sequence = ['NodeA@1', 'NodeB@1']
+
+  # When barrier is created
+  barrier = ReplaySequenceBarrier(sequence)
+
+  # Then state is correctly set
+  assert barrier.sequence == sequence
+  assert barrier.current_index == 0
+  assert len(barrier.events) == 2
+  assert barrier.events['NodeA@1'].is_set()
+  assert not barrier.events['NodeB@1'].is_set()
+
+
+@pytest.mark.asyncio
+async def test_barrier_wait_blocks_and_unblocks():
+  """Verifies that wait blocks on subsequent keys and is unblocked by advance."""
+  sequence = ['NodeA@1', 'NodeB@1']
+  barrier = ReplaySequenceBarrier(sequence)
+
+  # When first key waits, it completes instantly
+  await barrier.wait('NodeA@1')
+
+  # When second key waits, it blocks
+  b_completed = False
+
+  async def wait_b():
+    nonlocal b_completed
+    await barrier.wait('NodeB@1')
+    b_completed = True
+
+  task = asyncio.create_task(wait_b())
+  await asyncio.sleep(0.05)
+  assert not b_completed  # Still blocked
+
+  # When first key advances the sequence
+  barrier.check_and_advance('NodeA@1')
+
+  # Then index progresses and second event is released
+  await task
+  assert b_completed
+  assert barrier.current_index == 1
+  assert barrier.events['NodeB@1'].is_set()
+
+
+def test_barrier_advance_out_of_order_ignored():
+  """Verifies that out-of-order advance calls are ignored and do not progress index."""
+  sequence = ['NodeA@1', 'NodeB@1']
+  barrier = ReplaySequenceBarrier(sequence)
+
+  # When second key tries to advance out of order
+  barrier.check_and_advance('NodeB@1')
+
+  # Then state remains unchanged
+  assert barrier.current_index == 0
+  assert not barrier.events['NodeB@1'].is_set()
+
+
+@pytest.mark.asyncio
+async def test_barrier_wait_non_existent_key():
+  """Verifies that waiting on a key not in sequence does not block."""
+  sequence = ['NodeA@1']
+  barrier = ReplaySequenceBarrier(sequence)
+
+  # When a key not in sequence waits, it passes instantly
+  await barrier.wait('NonExistent@1')
+
+  # No blocks, successfully completes!
+  assert True
+
+
+@pytest.mark.asyncio
+async def test_barrier_wait_timeout_on_divergence():
+  """Verifies that waiting on a blocked key raises RuntimeError due to timeout."""
+  # We use a short sequence where NodeB is never unblocked
+  sequence = ['NodeA@1', 'NodeB@1']
+  # Use a fast timeout to keep the test rapid without mocking standard library functions
+  barrier = ReplaySequenceBarrier(sequence, timeout_sec=0.01)
+
+  with pytest.raises(RuntimeError, match='Replay divergence detected'):
+    await barrier.wait('NodeB@1')

tests/unittests/workflow/utils/test_retry_utils.py

@@ -0,0 +1,70 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from google.adk.workflow._node_state import NodeState
+from google.adk.workflow._retry_config import RetryConfig
+from google.adk.workflow.utils._retry_utils import _get_retry_delay
+import pytest
+
+
+class TestGetRetryDelay:
+
+  def test_returns_default_delay_without_config(self):
+    """Returns default delay of 1.0 second when config is missing."""
+    state = NodeState(attempt_count=1)
+
+    result = _get_retry_delay(None, state)
+
+    assert result == 1.0
+
+  def test_returns_initial_delay_on_first_failure(self):
+    """Returns initial delay on the first failure attempt."""
+    config = RetryConfig(initial_delay=2.0, jitter=0.0)
+    state = NodeState(attempt_count=1)
+
+    result = _get_retry_delay(config, state)
+
+    assert result == 2.0
+
+  def test_applies_exponential_backoff(self):
+    """Applies exponential backoff for subsequent attempts."""
+    config = RetryConfig(initial_delay=2.0, backoff_factor=2.0, jitter=0.0)
+    state = NodeState(attempt_count=2)
+
+    result = _get_retry_delay(config, state)
+
+    assert result == 4.0
+
+  def test_caps_at_max_delay(self):
+    """Caps calculated delay at the specified maximum delay."""
+    config = RetryConfig(
+        initial_delay=2.0, backoff_factor=10.0, max_delay=15.0, jitter=0.0
+    )
+    state = NodeState(attempt_count=2)
+
+    result = _get_retry_delay(config, state)
+
+    assert result == 15.0
+
+  def test_adds_jitter_when_enabled(self):
+    """Adds random jitter to the calculated delay."""
+    config = RetryConfig(initial_delay=10.0, backoff_factor=1.0, jitter=0.5)
+    state = NodeState(attempt_count=1)
+
+    delays = [_get_retry_delay(config, state) for _ in range(10)]
+
+    assert all(5.0 <= d <= 15.0 for d in delays)
+    assert len(set(delays)) > 1

tests/unittests/workflow/utils/test_transfer_utils.py

@@ -0,0 +1,188 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for agent transfer utilities."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.workflow.utils._transfer_utils import resolve_and_derive_transfer_context
+import pytest
+
+
+def test_resolve_and_derive_transfer_context_raises_value_error_on_self_transfer():
+  """resolve_and_derive_transfer_context raises ValueError when target is the current agent."""
+  # Arrange
+  current = LlmAgent(name='current')
+  root = LlmAgent(name='root', sub_agents=[current])
+
+  # Act & Assert
+  with pytest.raises(ValueError, match='cannot transfer to itself'):
+    resolve_and_derive_transfer_context(
+        'current', current, root, MagicMock(), None
+    )
+
+
+def test_resolve_and_derive_transfer_context_returns_child_context():
+  """resolve_and_derive_transfer_context returns current context as parent context for CHILD transfers."""
+  # Arrange
+  target = LlmAgent(name='target')
+  current = LlmAgent(name='current', sub_agents=[target])
+  root = LlmAgent(name='root', sub_agents=[current])
+
+  curr_ctx = MagicMock()
+
+  # Act
+  resolved_agent, parent_ctx = resolve_and_derive_transfer_context(
+      'target', current, root, curr_ctx, None
+  )
+
+  # Assert
+  assert resolved_agent is target
+  assert parent_ctx is curr_ctx
+
+
+def test_resolve_and_derive_transfer_context_returns_sibling_context():
+  """resolve_and_derive_transfer_context returns parent context for SIBLING transfers."""
+  # Arrange
+  current = LlmAgent(name='current')
+  target = LlmAgent(name='target')
+  root = LlmAgent(name='root', sub_agents=[current, target])
+
+  parent_ctx = MagicMock()
+
+  # Act
+  resolved_agent, derived_ctx = resolve_and_derive_transfer_context(
+      'target', current, root, MagicMock(), parent_ctx
+  )
+
+  # Assert
+  assert resolved_agent is target
+  assert derived_ctx is parent_ctx
+
+
+def test_resolve_and_derive_transfer_context_climbs_parent_context():
+  """resolve_and_derive_transfer_context climbs context chain to find the target parent's parent context."""
+  # Arrange
+  root_ctx = MagicMock()
+  root_ctx.node = MagicMock()
+  root_ctx.node.name = 'root'
+  root_ctx.parent_ctx = MagicMock()
+  root_ctx.parent_ctx.node = None
+  root_ctx.parent_ctx.parent_ctx = None
+
+  child_ctx = MagicMock()
+  child_ctx.node = MagicMock()
+  child_ctx.node.name = 'child'
+  child_ctx.parent_ctx = root_ctx
+
+  # Target is 'root', current is 'child'
+  child = LlmAgent(name='child')
+  root = LlmAgent(name='root', sub_agents=[child])
+  child.parent_agent = root
+
+  # Act
+  resolved_agent, derived_ctx = resolve_and_derive_transfer_context(
+      'root', child, root, child_ctx, None
+  )
+
+  # Assert
+  assert resolved_agent is root
+  assert derived_ctx is root_ctx.parent_ctx
+
+
+def test_resolve_and_derive_transfer_context_returns_root_context_when_parent_bypassed():
+  """resolve_and_derive_transfer_context returns root context for PARENT transfers when parent was bypassed."""
+  # Arrange
+  root_ctx = MagicMock()
+  root_ctx.node = None
+  root_ctx.parent_ctx = None
+
+  child_ctx = MagicMock()
+  child_ctx.node = MagicMock()
+  child_ctx.node.name = 'child'
+  child_ctx.parent_ctx = root_ctx
+
+  # Target is 'root', current is 'child'
+  child = LlmAgent(name='child')
+  root = LlmAgent(name='root', sub_agents=[child])
+  child.parent_agent = root
+
+  # Act
+  resolved_agent, derived_ctx = resolve_and_derive_transfer_context(
+      'root', child, root, child_ctx, None
+  )
+
+  # Assert
+  assert resolved_agent is root
+  assert derived_ctx is root_ctx
+
+
+def test_resolve_and_derive_transfer_context_returns_none_when_agent_not_found():
+  """resolve_and_derive_transfer_context returns (None, None) when target agent is not found."""
+  # Arrange
+  current = LlmAgent(name='current')
+  root = LlmAgent(name='root', sub_agents=[current])
+
+  # Act
+  resolved_agent, derived_ctx = resolve_and_derive_transfer_context(
+      'target', current, root, MagicMock(), None
+  )
+
+  # Assert
+  assert resolved_agent is None
+  assert derived_ctx is None
+
+
+def test_resolve_and_derive_transfer_context_returns_target_and_none_when_no_relationship():
+  """resolve_and_derive_transfer_context returns (target_agent, None) for unrelated transfers."""
+  # Arrange
+  current = LlmAgent(name='current')
+  target = LlmAgent(name='target')
+  root1 = LlmAgent(name='root1', sub_agents=[current])
+  root2 = LlmAgent(name='root2', sub_agents=[target])
+
+  # Act
+  resolved_agent, derived_ctx = resolve_and_derive_transfer_context(
+      'target', current, root2, MagicMock(), None
+  )
+
+  # Assert
+  assert resolved_agent is target
+  assert derived_ctx is None
+
+
+def test_resolve_and_derive_transfer_context_works_with_cloned_agents():
+  """resolve_and_derive_transfer_context works correctly when the current agent is cloned (name-based matching)."""
+  # Arrange
+  target = LlmAgent(name='target')
+  current = LlmAgent(name='current', sub_agents=[target])
+  root = LlmAgent(name='root', sub_agents=[current])
+
+  cloned_current = current.clone()
+  assert cloned_current is not current
+  assert cloned_current.name == current.name
+
+  curr_ctx = MagicMock()
+
+  # Act
+  resolved_agent, derived_ctx = resolve_and_derive_transfer_context(
+      'target', cloned_current, root, curr_ctx, None
+  )
+
+  # Assert
+  assert resolved_agent is target
+  assert derived_ctx is curr_ctx

tests/unittests/workflow/utils/test_workflow_graph_utils.py

@@ -0,0 +1,136 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.tools.base_tool import BaseTool
+from google.adk.workflow._base_node import BaseNode
+from google.adk.workflow._base_node import START
+from google.adk.workflow._function_node import FunctionNode
+from google.adk.workflow._tool_node import _ToolNode
+from google.adk.workflow.utils._workflow_graph_utils import build_node
+from google.adk.workflow.utils._workflow_graph_utils import is_node_like
+import pytest
+
+
+class TestIsNodeLike:
+
+  def test_returns_true_for_base_node(self):
+    """is_node_like returns True for BaseNode instances."""
+
+    class DummyNode(BaseNode):
+
+      async def _run_impl(self, *, ctx, node_input):
+        yield node_input
+
+    node = DummyNode(name="test")
+
+    assert is_node_like(node) is True
+
+  def test_returns_true_for_base_tool(self):
+    """is_node_like returns True for BaseTool instances."""
+
+    class DummyTool(BaseTool):
+
+      def execute(self, **kwargs):
+        return "done"
+
+    tool = DummyTool(name="test", description="test")
+
+    assert is_node_like(tool) is True
+
+  def test_returns_true_for_callable(self):
+    """is_node_like returns True for callables."""
+
+    def my_func():
+      pass
+
+    assert is_node_like(my_func) is True
+
+  def test_returns_true_for_start_string(self):
+    """is_node_like returns True for 'START' string."""
+    assert is_node_like("START") is True
+
+  def test_returns_false_for_invalid_types(self):
+    """is_node_like returns False for invalid types."""
+    assert is_node_like(123) is False
+    assert is_node_like("NOT_START") is False
+
+
+class TestBuildNode:
+
+  def test_returns_start_when_node_like_is_start(self):
+    """build_node returns START sentinel when input is 'START'."""
+    assert build_node("START") == START
+
+  def test_returns_copy_of_base_node_with_overrides(self):
+    """build_node returns a copy of BaseNode with provided overrides."""
+
+    class DummyNode(BaseNode):
+
+      async def _run_impl(self, *, ctx, node_input):
+        yield node_input
+
+    node = DummyNode(name="original")
+
+    built = build_node(node, name="new_name")
+
+    assert built != node
+    assert built.name == "new_name"
+
+  def test_returns_tool_node_for_base_tool(self):
+    """build_node wraps BaseTool in a _ToolNode."""
+
+    class DummyTool(BaseTool):
+
+      def execute(self, **kwargs):
+        return "done"
+
+    tool = DummyTool(name="test", description="test")
+
+    built = build_node(tool)
+
+    assert isinstance(built, _ToolNode)
+
+  def test_returns_function_node_for_callable(self):
+    """build_node wraps callable in a FunctionNode."""
+
+    def my_func(x):
+      return x
+
+    built = build_node(my_func)
+
+    assert isinstance(built, FunctionNode)
+
+  def test_raises_value_error_for_invalid_type(self):
+    """build_node raises ValueError for invalid types."""
+    with pytest.raises(ValueError, match="Invalid node type"):
+      build_node(123)
+
+  def test_llm_agent_mode_defaults(self):
+    """build_node sets correct default mode for LlmAgent."""
+    root_agent = LlmAgent(name="root", instruction="test")
+    sub_agent = LlmAgent(name="sub", description="test")
+    # Dynamic subagent attachment without model_post_init normalization
+    sub_agent.parent_agent = root_agent
+
+    # Subagent with parent_agent should default to chat mode
+    built_sub = build_node(sub_agent)
+    assert built_sub.mode == "chat"
+
+    # Standalone agent without parent_agent should default to single_turn
+    standalone = LlmAgent(name="standalone", instruction="test")
+    built_standalone = build_node(standalone)
+    assert built_standalone.mode == "single_turn"

tests/unittests/workflow/utils/test_workflow_hitl_utils.py

@@ -0,0 +1,219 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import json
+
+from google.adk.events.event import Event
+from google.adk.events.event import NodeInfo
+from google.adk.events.request_input import RequestInput
+from google.adk.workflow.utils._rehydration_utils import _ChildScanState
+from google.adk.workflow.utils._workflow_hitl_utils import create_auth_request_event
+from google.adk.workflow.utils._workflow_hitl_utils import create_request_input_event
+from google.adk.workflow.utils._workflow_hitl_utils import create_request_input_response
+from google.adk.workflow.utils._workflow_hitl_utils import get_request_input_interrupt_ids
+from google.adk.workflow.utils._workflow_hitl_utils import has_request_input_function_call
+from google.adk.workflow.utils._workflow_hitl_utils import REQUEST_CREDENTIAL_FUNCTION_CALL_NAME
+from google.genai import types
+
+# --- create_request_input_event ---
+
+
+class TestCreateRequestInputEvent:
+
+  def test_basic_event(self):
+    ri = RequestInput(
+        interrupt_id="test-id",
+        message="Please approve",
+    )
+    event = create_request_input_event(ri)
+
+    assert event.long_running_tool_ids == {"test-id"}
+    assert event.content is not None
+    assert event.content.role == "model"
+    fc = event.content.parts[0].function_call
+    assert fc.name == "adk_request_input"
+    assert fc.id == "test-id"
+    assert fc.args["message"] == "Please approve"
+
+  def test_with_payload(self):
+    ri = RequestInput(
+        interrupt_id="id-1",
+        payload={"key": "value"},
+    )
+    event = create_request_input_event(ri)
+    fc = event.content.parts[0].function_call
+    assert fc.args["payload"] == {"key": "value"}
+
+  def test_with_response_schema(self):
+    from pydantic import BaseModel
+
+    class MySchema(BaseModel):
+      approved: bool
+
+    ri = RequestInput(
+        interrupt_id="id-2",
+        response_schema=MySchema,
+    )
+    event = create_request_input_event(ri)
+    fc = event.content.parts[0].function_call
+    schema = fc.args["response_schema"]
+    assert "approved" in schema["properties"]
+    assert schema["properties"]["approved"]["type"] == "boolean"
+
+
+# --- has_request_input_function_call ---
+
+
+class TestHasRequestInputFunctionCall:
+
+  def test_true_for_request_input_event(self):
+    event = create_request_input_event(
+        RequestInput(interrupt_id="id-1", message="test")
+    )
+    assert has_request_input_function_call(event) is True
+
+  def test_false_for_empty_event(self):
+    assert has_request_input_function_call(Event()) is False
+
+  def test_false_for_non_request_input(self):
+    from google.genai import types
+
+    event = Event(
+        content=types.Content(
+            parts=[
+                types.Part(
+                    function_call=types.FunctionCall(name="other_tool", args={})
+                )
+            ]
+        )
+    )
+    assert has_request_input_function_call(event) is False
+
+
+# --- create_request_input_response ---
+
+
+class TestCreateRequestInputResponse:
+
+  def test_creates_function_response_part(self):
+    part = create_request_input_response("id-1", {"approved": True})
+    assert part.function_response.id == "id-1"
+    assert part.function_response.name == "adk_request_input"
+    assert part.function_response.response == {"approved": True}
+
+
+# --- get_request_input_interrupt_ids ---
+
+
+class TestGetRequestInputInterruptIds:
+
+  def test_extracts_ids(self):
+    event = create_request_input_event(
+        RequestInput(interrupt_id="id-1", message="test")
+    )
+    assert get_request_input_interrupt_ids(event) == ["id-1"]
+
+  def test_empty_for_no_function_calls(self):
+    assert get_request_input_interrupt_ids(Event()) == []
+
+  def test_empty_for_non_request_input(self):
+    from google.genai import types
+
+    event = Event(
+        content=types.Content(
+            parts=[
+                types.Part(
+                    function_call=types.FunctionCall(
+                        name="other_tool", args={}, id="id-1"
+                    )
+                )
+            ]
+        )
+    )
+    assert get_request_input_interrupt_ids(event) == []
+
+
+# --- create_auth_request_event ---
+
+
+class TestCreateAuthRequestEvent:
+
+  def test_creates_credential_request(self):
+    from fastapi.openapi.models import APIKey
+    from fastapi.openapi.models import APIKeyIn
+    from google.adk.auth.auth_credential import AuthCredential
+    from google.adk.auth.auth_credential import AuthCredentialTypes
+    from google.adk.auth.auth_tool import AuthConfig
+
+    auth_config = AuthConfig(
+        auth_scheme=APIKey(**{"in": APIKeyIn.header, "name": "X-Api-Key"}),
+        raw_auth_credential=AuthCredential(
+            auth_type=AuthCredentialTypes.API_KEY,
+            api_key="test_key",
+        ),
+        credential_key="test_cred",
+    )
+    event = create_auth_request_event(auth_config, "auth-id-1")
+
+    assert event.long_running_tool_ids is not None
+    fc = event.content.parts[0].function_call
+    assert fc.name == REQUEST_CREDENTIAL_FUNCTION_CALL_NAME
+    assert fc.id == "auth-id-1"
+    assert "authConfig" in fc.args
+
+  def test_args_are_json_serializable(self):
+    from fastapi.openapi.models import OAuth2
+    from fastapi.openapi.models import OAuthFlowAuthorizationCode
+    from fastapi.openapi.models import OAuthFlows
+    from google.adk.auth.auth_credential import AuthCredential
+    from google.adk.auth.auth_credential import AuthCredentialTypes
+    from google.adk.auth.auth_credential import OAuth2Auth
+    from google.adk.auth.auth_tool import AuthConfig
+
+    auth_config = AuthConfig(
+        auth_scheme=OAuth2(
+            flows=OAuthFlows(
+                authorizationCode=OAuthFlowAuthorizationCode(
+                    authorizationUrl=(
+                        "https://accounts.google.com/o/oauth2/auth"
+                    ),
+                    tokenUrl="https://oauth2.googleapis.com/token",
+                    scopes={
+                        "https://www.googleapis.com/auth/calendar": (
+                            "See calendars"
+                        )
+                    },
+                )
+            )
+        ),
+        raw_auth_credential=AuthCredential(
+            auth_type=AuthCredentialTypes.OAUTH2,
+            oauth2=OAuth2Auth(
+                client_id="oauth_client_id",
+                client_secret="oauth_client_secret",
+            ),
+        ),
+    )
+    event = create_auth_request_event(auth_config, "auth-id-1")
+
+    fc = event.content.parts[0].function_call
+
+    # python-mode dump leaves auth_scheme.type a live enum, breaking json.dumps
+    json.dumps(fc.args)
+    assert fc.args["authConfig"]["authScheme"]["type"] == "oauth2"
+
+
+#

tests/unittests/workflow/workflow_testing_utils.py

@@ -0,0 +1,399 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Testing utils for the Workflow."""
+
+import copy
+import inspect
+from typing import Any
+from typing import AsyncGenerator
+from typing import Callable
+from typing import List
+from typing import Optional
+
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.context import Context
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.invocation_context import InvocationContext as BaseInvocationContext
+from google.adk.apps.app import ResumabilityConfig
+from google.adk.events.event import Event
+from google.adk.events.request_input import RequestInput
+from google.adk.runners import Runner
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.workflow import BaseNode
+from google.adk.workflow._graph import RouteValue
+from google.adk.workflow.utils._workflow_hitl_utils import has_auth_request_function_call
+from google.adk.workflow.utils._workflow_hitl_utils import has_request_input_function_call
+from google.genai import types
+from pydantic import ConfigDict
+from pydantic import Field
+from typing_extensions import override
+
+from .testing_utils import END_OF_AGENT
+from .testing_utils import simplify_content
+
+
+async def run_workflow(wf, message='start'):
+  """Run a Workflow through Runner, return collected events."""
+  ss = InMemorySessionService()
+  runner = Runner(app_name=wf.name, node=wf, session_service=ss)
+  session = await ss.create_session(app_name=wf.name, user_id='u')
+  msg = types.Content(parts=[types.Part(text=message)], role='user')
+  events = []
+  async for event in runner.run_async(
+      user_id='u', session_id=session.id, new_message=msg
+  ):
+    events.append(event)
+  return events, ss, session
+
+
+# Emulates a node that outputs an Event & a route.
+# If output is not None, the output is set as the data field in the event.
+# If route is not None, the route is set in the node output event.
+# The route can be set without the output. This means didn't produce any output
+# but wants to signal a route to take.
+class TestingNode(BaseNode):
+  model_config = ConfigDict(arbitrary_types_allowed=True)
+
+  output: Optional[Any] = None
+  route: (
+      RouteValue | list[RouteValue] | Callable[[Context, Any], Any] | None
+  ) = None
+  received_inputs: List[Any] = Field(default_factory=list)
+
+  @override
+  async def _run_impl(
+      self,
+      *,
+      ctx: Context,
+      node_input: Any,
+  ) -> AsyncGenerator[Any, None]:
+    if self.output is not None or self.route is not None:
+      route = None
+      if callable(self.route):
+        if inspect.iscoroutinefunction(self.route):
+          route = await self.route(ctx, node_input)
+        else:
+          route = self.route(ctx, node_input)
+      else:
+        route = self.route
+
+      self.received_inputs.append(node_input)
+      yield Event(
+          output=self.output,
+          route=route,
+      )
+
+
+class TestingNodeWithIntermediateContent(BaseNode):
+  model_config = ConfigDict(arbitrary_types_allowed=True)
+
+  intermediate_content: list[types.Content] = Field(default_factory=list)
+  output: Optional[Any] = None
+  route: Optional[str] = None
+
+  @override
+  async def _run_impl(
+      self,
+      *,
+      ctx: Context,
+      node_input: Any,
+  ) -> AsyncGenerator[Any, None]:
+    for content in self.intermediate_content:
+      yield Event(
+          author=self.name,
+          invocation_id=ctx.invocation_id,
+          content=content,
+      )
+
+    if self.output is not None:
+      yield Event(
+          output=self.output,
+          route=self.route,
+      )
+
+
+class InputCapturingNode(BaseNode):
+  """A node that captures the inputs it receives."""
+
+  model_config = ConfigDict(arbitrary_types_allowed=True)
+
+  received_inputs: List[Any] = Field(default_factory=list)
+
+  @override
+  async def _run_impl(
+      self,
+      *,
+      ctx: Context,
+      node_input: Any,
+  ) -> AsyncGenerator[Any, None]:
+    self.received_inputs.append(node_input)
+    yield Event(
+        output={'received': node_input},
+    )
+
+
+class RequestInputNode(BaseNode):
+  """A simple node that requests input from the user."""
+
+  model_config = ConfigDict(arbitrary_types_allowed=True)
+
+  message: str = Field(default='')
+  response_schema: Optional[dict[str, Any]] = None
+
+  @override
+  async def _run_impl(
+      self,
+      *,
+      ctx: Context,
+      node_input: Any,
+  ) -> AsyncGenerator[Any, None]:
+    yield RequestInput(
+        message=self.message,
+        response_schema=self.response_schema,
+    )
+
+
+async def create_parent_invocation_context(
+    test_name: str, agent: BaseAgent, resumable: bool = False
+) -> InvocationContext:
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user'
+  )
+  return InvocationContext(
+      invocation_id=f'{test_name}_invocation_id',
+      agent=agent,
+      session=session,
+      session_service=session_service,
+      resumability_config=ResumabilityConfig(is_resumable=resumable),
+  )
+
+
+def simplify_event_with_node(
+    event: Event,
+    include_state_delta: bool = False,
+) -> Any | None:
+  if isinstance(event, Event):
+    if (
+        'output' not in event.model_fields_set
+        and not (include_state_delta and event.actions.state_delta)
+        and not event.content
+    ):
+      return None
+
+    # If the event has content, return the simplified content.
+    if event.content:
+      return simplify_content(event.content)
+
+    simplified_event = {}
+
+    # Also simplify event.output if it contains Content.
+    # The tests assume that Content found in event data should be simplified
+    # (IDs stripped) just like event.content. This ensures consistent
+    # assertion behavior.
+    output = event.output
+    if isinstance(output, types.Content):
+      output = copy.deepcopy(output)
+      for part in output.parts:
+        if part.function_call and part.function_call.id:
+          part.function_call.id = None
+        if part.function_response and part.function_response.id:
+          part.function_response.id = None
+    simplified_event['output'] = output
+
+    if include_state_delta and event.actions.state_delta:
+      simplified_event['state_delta'] = event.actions.state_delta
+
+    return simplified_event
+  elif event.content:
+    return simplify_content(event.content)
+
+
+def simplify_events_with_node(
+    events: list[Event],
+    *,
+    include_state_delta: bool = False,
+    include_workflow_output: bool = False,
+) -> list[tuple[str, Any]]:
+  results = []
+
+  # Second pass: Simplify events
+  for event in events:
+    # Optionally skip top-level workflow output events (events emitted
+    # by the Workflow in _finalize_workflow). These events have a
+    # top-level node_path (no '/' separator) and carry output data.
+    if (
+        not include_workflow_output
+        and isinstance(event, Event)
+        and event.output is not None
+        and '/' not in (event.node_info.path or '')
+    ):
+      continue
+
+    simplified_event = simplify_event_with_node(event, include_state_delta)
+    if simplified_event:
+      # Map the author to the source node name if it exists.
+      if hasattr(event, 'node_info') and event.node_info.path:
+        author = event.node_info.path
+      else:
+        author = event.author
+      results.append((author, simplified_event))
+  return results
+
+
+def simplify_events_with_node_and_agent_state(
+    events: list[Event],
+    *,
+    include_state_delta: bool = False,
+    include_inputs_and_triggers: bool = False,
+    include_resume_inputs: bool = False,
+    include_workflow_output: bool = False,
+):
+  fields_to_exclude = {'run_id'}
+  if not include_inputs_and_triggers:
+    fields_to_exclude.add('input')
+  if not include_resume_inputs:
+    fields_to_exclude.add('resume_inputs')
+
+  results = []
+
+  for event in events:
+    # Optionally skip top-level workflow output events.
+    if (
+        not include_workflow_output
+        and isinstance(event, Event)
+        and event.output is not None
+        and '/' not in (event.node_info.path or '')
+    ):
+      continue
+    simplified_event = simplify_event_with_node(event, include_state_delta)
+
+    # Map the author to the source node name if it exists.
+    if hasattr(event, 'node_info') and event.node_info.path:
+      author = event.node_info.path
+    else:
+      author = event.author
+
+    if simplified_event:
+      results.append((author, simplified_event))
+    elif event.actions.end_of_agent:
+      results.append((author, END_OF_AGENT))
+    elif event.actions.agent_state is not None:
+      agent_state = event.actions.agent_state
+      nodes = agent_state.get('nodes', {})
+      simplified_nodes = {}
+      for node_name, node_state in nodes.items():
+        simplified_nodes[node_name] = {
+            k: v
+            for k, v in node_state.items()
+            if k not in fields_to_exclude
+            and (k != 'interrupts' or v)  # Exclude empty interrupts
+            and (k != 'resume_inputs' or v)  # Exclude empty resume_inputs
+        }
+      results.append((author, {'nodes': simplified_nodes}))
+  return results
+
+
+def get_request_input_events(events: list[Any]) -> list[Any]:
+  """Returns a list of request input events from the given list of events."""
+  return [e for e in events if has_request_input_function_call(e)]
+
+
+def get_auth_request_events(events: list[Any]) -> list[Any]:
+  """Returns a list of auth credential request events from the given list."""
+  return [e for e in events if has_auth_request_function_call(e)]
+
+
+def get_output_events(events: list[Any], output: Any = None) -> list[Any]:
+  """Returns a list of events that have output populated."""
+  return [
+      e
+      for e in events
+      if isinstance(e, Event)
+      and e.output is not None
+      and (output is None or e.output == output)
+  ]
+
+
+def get_outputs(events: list[Event]) -> list[Any]:
+  """Extracts output values from events, skipping non-output events."""
+  return [
+      e.output for e in events if isinstance(e, Event) and e.output is not None
+  ]
+
+
+def strip_checkpoint_events(
+    simplified_events: list[tuple[str, Any]],
+) -> list[tuple[str, Any]]:
+  """Strips agent_state checkpoint and end_of_agent events.
+
+  In non-resumable mode, the workflow does not emit checkpoint events
+  or end_of_agent events.  Use this to derive the expected simplified
+  output for non-resumable tests from the resumable expected output.
+  """
+  return [
+      (author, data)
+      for author, data in simplified_events
+      if not (isinstance(data, dict) and 'nodes' in data)
+      and data != END_OF_AGENT
+  ]
+
+
+def find_function_call_event(
+    events: list[Any], name: str | None = None
+) -> Any | None:
+  """Finds the first event containing a function call."""
+  for e in events:
+    if hasattr(e, 'content') and e.content and e.content.parts:
+      for part in e.content.parts:
+        if part.function_call:
+          if name is None or part.function_call.name == name:
+            return e
+  return None
+
+
+class CustomRetryableError(Exception):
+  """A custom error meant to be retried."""
+
+
+class CustomNonRetryableError(Exception):
+  """A custom error not meant to be retried."""
+
+
+class _FlakyNode(BaseNode):
+  model_config = ConfigDict(arbitrary_types_allowed=True)
+
+  message: str = Field(default='')
+  succeed_on_iteration: int = Field(default=0)
+  tracker: dict[str, Any] = Field(default_factory=dict)
+  exception_to_raise: Exception = Field(...)
+
+  @override
+  async def run(
+      self,
+      *,
+      ctx: Context,
+      node_input: Any,
+  ) -> AsyncGenerator[Any, None]:
+    iteration_count = self.tracker.get('iteration_count', 0) + 1
+    self.tracker['iteration_count'] = iteration_count
+    self.tracker.setdefault('attempt_counts', []).append(ctx.attempt_count)
+
+    if iteration_count < self.succeed_on_iteration:
+      raise self.exception_to_raise
+
+    yield Event(
+        output=self.message,
+    )