ec2b666284
Continuous Integration / Pre-commit Linter (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.10) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.11) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.12) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.10) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.11) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.12) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.14) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Waiting to run
Copybara PR Handler / close-imported-pr (push) Waiting to run
589 lines
18 KiB
Python
589 lines
18 KiB
Python
# 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
|