diff --git a/contributing/samples/models/interactions_api/tests/basic_text.json b/contributing/samples/models/interactions_api/tests/basic_text.json new file mode 100644 index 00000000000..f200e01e94b --- /dev/null +++ b/contributing/samples/models/interactions_api/tests/basic_text.json @@ -0,0 +1,37 @@ +{ + "events": [ + { + "author": "user", + "content": { + "parts": [ + { + "text": "Hello! What can you help me with?" + } + ], + "role": "user" + }, + "id": "e-1", + "invocationId": "i-1", + "nodeInfo": { + "path": "" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "text": "Hello! I am \"interactions_test_agent,\" an agent designed for testing the Interactions API integration.\n\nI can help you with two main tasks:\n1. **Searching the web:** If you have questions that require up-to-date information, I can search the web for you.\n2. **Checking the weather:** I can provide the current weather for a specific city.\n\nHow can I assist you today?" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-2", + "invocationId": "i-1", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + } + ] +} diff --git a/contributing/samples/models/interactions_api/tests/custom_function_weather.json b/contributing/samples/models/interactions_api/tests/custom_function_weather.json new file mode 100644 index 00000000000..f5f5bcad550 --- /dev/null +++ b/contributing/samples/models/interactions_api/tests/custom_function_weather.json @@ -0,0 +1,86 @@ +{ + "events": [ + { + "author": "user", + "content": { + "parts": [ + { + "text": "What's the weather like in Tokyo?" + } + ], + "role": "user" + }, + "id": "e-1", + "invocationId": "i-1", + "nodeInfo": { + "path": "" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "functionCall": { + "args": { + "city": "Tokyo" + }, + "id": "fc-1", + "name": "get_current_weather" + } + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-2", + "invocationId": "i-1", + "longRunningToolIds": [], + "nodeInfo": { + "path": "interactions_test_agent@1" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "functionResponse": { + "id": "fc-1", + "name": "get_current_weather", + "response": { + "city": "Tokyo", + "condition": "Partly Cloudy", + "humidity": 60, + "temperature_f": 68 + } + } + } + ], + "role": "user" + }, + "id": "e-3", + "invocationId": "i-1", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "text": "I have retrieved the weather information for Tokyo. It is currently 68\u00b0F and partly cloudy with 60% humidity." + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-4", + "invocationId": "i-1", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + } + ] +} diff --git a/contributing/samples/models/interactions_api/tests/google_search_1984.json b/contributing/samples/models/interactions_api/tests/google_search_1984.json new file mode 100644 index 00000000000..317116c7e1d --- /dev/null +++ b/contributing/samples/models/interactions_api/tests/google_search_1984.json @@ -0,0 +1,40 @@ +{ + "events": [ + { + "author": "user", + "content": { + "parts": [ + { + "text": "Use google search to find out who wrote the novel '1984'." + } + ], + "role": "user" + }, + "id": "e-1", + "invocationId": "i-1", + "nodeInfo": { + "path": "" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "text": "search_suggestions='\\n
\\n
\\n \\n \\n \\n \\n \\n \\n \\n \\n \\n \\n \\n \\n \\n
\\n
\\n
\\nwho wrote the novel 1984\\n
\\n
'" + }, + { + "text": "The novel *1984* was written by the English author **George Orwell** (whose real name was Eric Arthur Blair). It was published in 1949." + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-2", + "invocationId": "i-1", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + } + ] +} diff --git a/contributing/samples/models/interactions_api/tests/google_search_france.json b/contributing/samples/models/interactions_api/tests/google_search_france.json new file mode 100644 index 00000000000..f2ca266d82a --- /dev/null +++ b/contributing/samples/models/interactions_api/tests/google_search_france.json @@ -0,0 +1,40 @@ +{ + "events": [ + { + "author": "user", + "content": { + "parts": [ + { + "text": "Search for the capital of France." + } + ], + "role": "user" + }, + "id": "e-1", + "invocationId": "i-1", + "nodeInfo": { + "path": "" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "text": "search_suggestions='\\n
\\n
\\n \\n \\n \\n \\n \\n \\n \\n \\n \\n \\n \\n \\n \\n
\\n
\\n
\\ncapital of France\\n
\\n
'" + }, + { + "text": "I have searched for the capital of France and confirmed that it is Paris." + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-2", + "invocationId": "i-1", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + } + ] +} diff --git a/contributing/samples/models/interactions_api/tests/multi_turn.json b/contributing/samples/models/interactions_api/tests/multi_turn.json new file mode 100644 index 00000000000..aeb4d0494b6 --- /dev/null +++ b/contributing/samples/models/interactions_api/tests/multi_turn.json @@ -0,0 +1,152 @@ +{ + "events": [ + { + "author": "user", + "content": { + "parts": [ + { + "text": "My favorite color is blue. Just acknowledge this, don't use any tools." + } + ], + "role": "user" + }, + "id": "e-1", + "invocationId": "i-1", + "nodeInfo": { + "path": "" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "text": "I acknowledge that your favorite color is blue." + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-2", + "invocationId": "i-1", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + }, + { + "author": "user", + "content": { + "parts": [ + { + "text": "What's the weather like in London?" + } + ], + "role": "user" + }, + "id": "e-3", + "invocationId": "i-2", + "nodeInfo": { + "path": "" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "functionCall": { + "args": { + "city": "London" + }, + "id": "fc-1", + "name": "get_current_weather" + } + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-4", + "invocationId": "i-2", + "longRunningToolIds": [], + "nodeInfo": { + "path": "interactions_test_agent@1" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "functionResponse": { + "id": "fc-1", + "name": "get_current_weather", + "response": { + "city": "London", + "condition": "Cloudy", + "humidity": 78, + "temperature_f": 59 + } + } + } + ], + "role": "user" + }, + "id": "e-5", + "invocationId": "i-2", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "text": "I have checked the weather in London for you. It is currently cloudy with a temperature of 59\u00b0F and 78% humidity." + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-6", + "invocationId": "i-2", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + }, + { + "author": "user", + "content": { + "parts": [ + { + "text": "What is my favorite color that I mentioned earlier in our conversation?" + } + ], + "role": "user" + }, + "id": "e-7", + "invocationId": "i-3", + "nodeInfo": { + "path": "" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "text": "Your favorite color, which you mentioned earlier, is blue." + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-8", + "invocationId": "i-3", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + } + ] +} diff --git a/src/google/adk/cli/agent_test_runner.py b/src/google/adk/cli/agent_test_runner.py index 60fa5b9c49a..5312b94d50e 100644 --- a/src/google/adk/cli/agent_test_runner.py +++ b/src/google/adk/cli/agent_test_runner.py @@ -51,6 +51,17 @@ "citation_metadata", } +# The Interactions API stamps these volatile, non-reproducible fields onto every +# model response (interaction_id is a server-issued token; turn_complete is only +# emitted by the live API), so they are excluded from fixture comparison too. +# They are kept in a separate, private constant so the value of the public +# EXCLUDED_EVENT_FIELDS stays stable for the API breaking-change detector. +_EXTRA_EXCLUDED_EVENT_FIELDS = frozenset({"interaction_id", "turn_complete"}) + +_ALL_EXCLUDED_EVENT_FIELDS = ( + EXCLUDED_EVENT_FIELDS | _EXTRA_EXCLUDED_EVENT_FIELDS +) + # Read target folder from environment def get_test_files( @@ -175,7 +186,7 @@ def normalize_events(events, is_json=False): for e in events: if is_json: d = dict(e) - for k in EXCLUDED_EVENT_FIELDS: + for k in _ALL_EXCLUDED_EVENT_FIELDS: d.pop(k, None) d.pop(alias_generators.to_camel(k), None) d = {k: v for k, v in d.items() if v is not None} @@ -183,7 +194,7 @@ def normalize_events(events, is_json=False): d = e.model_dump( mode="json", by_alias=True, - exclude=EXCLUDED_EVENT_FIELDS, + exclude=_ALL_EXCLUDED_EVENT_FIELDS, exclude_none=True, ) @@ -706,6 +717,7 @@ def mock_randint(a, b): def rebuild_tests(path: str): """Discovers test files and rebuilds them by running the agent live.""" + import asyncio import json import sys @@ -734,6 +746,7 @@ def rebuild_tests(path: str): # Add agent_dir.parent to sys.path so relative imports work sys_path_saved = list(sys.path) sys.path.insert(0, str(agent_dir.parent)) + rebuild_loop = None try: import random @@ -774,6 +787,29 @@ def rebuild_tests(path: str): else InMemoryRunner(root_agent=agent_or_app) ) + # Drive every turn of this fixture on a single, persistent event loop. + # The sync Runner.run() spins up a fresh loop per call via asyncio.run() + # and closes it afterwards. For multi-turn fixtures that closes the loop + # the model's cached async api_client was bound to, so subsequent turns + # raise "Event loop is closed" (e.g. with the Interactions API). Reusing + # one loop for all turns keeps the client valid across the conversation. + rebuild_loop = asyncio.new_event_loop() + + def run_turn(content): + session = runner.session + + async def _collect(): + events = [] + async for event in runner.runner.run_async( + user_id=session.user_id, + session_id=session.id, + new_message=content, + ): + events.append(event) + return events + + return rebuild_loop.run_until_complete(_collect()) + new_events = [] inv_counter = 1 @@ -857,7 +893,7 @@ def get_next_fc_id(): content=msg, ) - run_events = runner.run(msg) + run_events = run_turn(msg) # Build mapping from old IDs to new agent IDs for e in run_events: @@ -891,6 +927,10 @@ def get_next_fc_id(): "cache_metadata", "logprobs_result", "citation_metadata", + # Volatile/non-replayable Interactions API state; keeping + # these out keeps rebuilt fixtures deterministic. + "interaction_id", + "turn_complete", }, ) for e in new_events @@ -931,6 +971,10 @@ def get_next_fc_id(): except Exception as e: print(f"Error rebuilding {test_file}: {e}") finally: + # Always close the per-fixture event loop, even if a turn raised, so we + # don't leak unclosed loops and emit resource warnings. + if rebuild_loop is not None: + rebuild_loop.close() sys.path = sys_path_saved diff --git a/src/google/adk/tools/mcp_tool/mcp_session_manager.py b/src/google/adk/tools/mcp_tool/mcp_session_manager.py index 59b0dab478a..8407ee7ab80 100644 --- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py +++ b/src/google/adk/tools/mcp_tool/mcp_session_manager.py @@ -391,7 +391,7 @@ async def before_request( ) return - if 'Authorization' in headers: + if any(k.lower() == 'authorization' for k in headers): logger.debug('Authorization header already present, not overwriting') return diff --git a/src/google/adk/workflow/_dynamic_node_scheduler.py b/src/google/adk/workflow/_dynamic_node_scheduler.py index 110117c015f..6cdd45e6c14 100644 --- a/src/google/adk/workflow/_dynamic_node_scheduler.py +++ b/src/google/adk/workflow/_dynamic_node_scheduler.py @@ -153,138 +153,77 @@ async def __call__( Returns: Child Context with output, route, and interrupt_ids set. """ - curr_node = node - curr_name = node_name or node.name - curr_run_id = run_id - curr_input = node_input - curr_parent_ctx: Context | None = ctx - - while True: - curr_parent_path = curr_parent_ctx.node_path if curr_parent_ctx else None - base_path_builder = ( - _NodePathBuilder.from_string(curr_parent_path) - if curr_parent_path - else _NodePathBuilder([]) - ) - node_path = str(base_path_builder.append(curr_name, curr_run_id)) - - # Rehydration chronological sequence barrier setup for the parent path - parent_path = curr_parent_ctx.node_path if curr_parent_ctx else '' - if parent_path and parent_path not in self._parent_sequence_barriers: - seq = self._scan_parent_child_sequence(curr_parent_ctx, parent_path) - self._parent_sequence_barriers[parent_path] = ReplaySequenceBarrier(seq) - - # Runtime schema validation. - if curr_input is not None: - try: - curr_input = curr_node._validate_input_data(curr_input) - except ValidationError as e: - raise ValueError( - 'Runtime schema validation failed for dynamic node' - f" '{curr_name}'. Input does not match input_schema: {e}" - ) from e - - logger.debug('node %s schedule start.', node_path) - - # Phase 1: Lazy rehydration from session events. - if node_path not in self._state.runs: - self._rehydrate_from_events(curr_parent_ctx, node_path) + curr_parent_path = ctx.node_path if ctx else None + base_path_builder = ( + _NodePathBuilder.from_string(curr_parent_path) + if curr_parent_path + else _NodePathBuilder([]) + ) + node_path = str(base_path_builder.append(node_name or node.name, run_id)) + + # Rehydration chronological sequence barrier setup for the parent path + parent_path = ctx.node_path if ctx else '' + if parent_path and parent_path not in self._parent_sequence_barriers: + seq = self._scan_parent_child_sequence(ctx, parent_path) + self._parent_sequence_barriers[parent_path] = ReplaySequenceBarrier(seq) + + # Runtime schema validation. + if node_input is not None: + try: + node_input = node._validate_input_data(node_input) + except ValidationError as e: + raise ValueError( + 'Runtime schema validation failed for dynamic node' + f" '{node_name or node.name}'. Input does not match" + f' input_schema: {e}' + ) from e + + logger.debug('node %s schedule start.', node_path) + + # Phase 1: Lazy rehydration from session events. + if node_path not in self._state.runs: + self._rehydrate_from_events(ctx, node_path) + + # Check existing run and determine if fresh execution is needed. + child_ctx, run_completed = await self._check_existing_run( + ctx, + node, + node_name or node.name, + node_path, + run_id, + node_input, + use_as_output, + use_sub_branch, + override_branch, + override_isolation_scope=override_isolation_scope, + ) - # Check existing run and determine if fresh execution is needed. - child_ctx, run_completed = await self._check_existing_run( - curr_parent_ctx, - curr_node, - curr_name, + if not run_completed: + # Phase 3: Fresh execution. + logger.debug('node %s schedule: Fresh execution.', node_path) + child_ctx = await self._run_node_internal( + ctx, + node, + node_name or node.name, node_path, - curr_run_id, - curr_input, + run_id, + node_input, use_as_output, - use_sub_branch, - override_branch, + is_fresh=True, + use_sub_branch=use_sub_branch, + override_branch=override_branch, override_isolation_scope=override_isolation_scope, ) - if not run_completed: - # Phase 3: Fresh execution. - logger.debug('node %s schedule: Fresh execution.', node_path) - child_ctx = await self._run_node_internal( - curr_parent_ctx, - curr_node, - curr_name, - node_path, - curr_run_id, - curr_input, - use_as_output, - is_fresh=True, - use_sub_branch=use_sub_branch, - override_branch=override_branch, - override_isolation_scope=override_isolation_scope, - ) - - logger.debug('node %s schedule end.', node_path) - - # Advance chronological sequence for this parent path and key - parent_path = curr_parent_ctx.node_path if curr_parent_ctx else '' - key = f'{curr_name}@{curr_run_id}' - if parent_path in self._parent_sequence_barriers: - self._parent_sequence_barriers[parent_path].check_and_advance(key) - - # Check for transfer_to_agent signal. - transfer_to_agent = ( - child_ctx.actions.transfer_to_agent if child_ctx else None - ) - if isinstance(transfer_to_agent, str): - target_name = transfer_to_agent - root_agent = getattr(curr_node, 'root_agent', None) - if not root_agent: - raise ValueError(f'Cannot find root_agent on node {curr_node.name}') - - # Local import to avoid runtime circular dependencies with Context - from .utils._transfer_utils import resolve_and_derive_transfer_context - - target_agent, next_parent_ctx = resolve_and_derive_transfer_context( - target_name=target_name, - current_agent=curr_node, - root_agent=root_agent, - curr_ctx=child_ctx, - curr_parent_ctx=curr_parent_ctx, - ) - if not target_agent: - raise ValueError(f"Transfer target agent '{target_name}' not found.") - if not next_parent_ctx: - available = [] - if hasattr(curr_node, '_get_available_agent_names'): - available = curr_node._get_available_agent_names() - available_str = ( - f"\nAvailable agents: {', '.join(available)}" if available else '' - ) - raise ValueError( - f"Cannot transfer from '{curr_name}' to unrelated agent" - f" '{target_name}'.{available_str}" - ) - curr_parent_ctx = next_parent_ctx - - # Set up parameters for next iteration. - curr_node = target_agent - curr_name = target_agent.name - - if not curr_parent_ctx: - raise AssertionError( - 'curr_parent_ctx cannot be None during active workflow execution' - ) - - curr_parent_ctx._child_run_counters[target_agent.name] = ( - curr_parent_ctx._child_run_counters.get(target_agent.name, 0) + 1 - ) - curr_run_id = str( - curr_parent_ctx._child_run_counters[target_agent.name] - ) - curr_input = None # Input for transfer target is usually empty. + logger.debug('node %s schedule end.', node_path) - # Loop continues to execute the next agent - continue + # Advance chronological sequence for this parent path and key + parent_path = ctx.node_path if ctx else '' + key = f'{node_name or node.name}@{run_id}' + if parent_path in self._parent_sequence_barriers: + self._parent_sequence_barriers[parent_path].check_and_advance(key) - return child_ctx + return child_ctx async def _check_existing_run( self, diff --git a/src/google/adk/workflow/_llm_agent_wrapper.py b/src/google/adk/workflow/_llm_agent_wrapper.py index 62920db2912..a4f0dc900ed 100644 --- a/src/google/adk/workflow/_llm_agent_wrapper.py +++ b/src/google/adk/workflow/_llm_agent_wrapper.py @@ -394,19 +394,23 @@ async def run_llm_agent_as_node( break # close this run_iter; outer loop re-enters if event.actions.transfer_to_agent: target_name = event.actions.transfer_to_agent - if target_name != agent.name: - from ..agents.llm_agent import LlmAgent - - if ( - isinstance(agent, LlmAgent) - and ctx._invocation_context.is_resumable - ): - ctx._invocation_context.set_agent_state( - agent.name, end_of_agent=True - ) - yield agent._create_agent_state_event(ctx._invocation_context) - transferred = True - break + if target_name == agent.name: + raise ValueError( + f"Agent '{target_name}' cannot transfer to itself." + ) + + from ..agents.llm_agent import LlmAgent + + if ( + isinstance(agent, LlmAgent) + and ctx._invocation_context.is_resumable + ): + ctx._invocation_context.set_agent_state( + agent.name, end_of_agent=True + ) + yield agent._create_agent_state_event(ctx._invocation_context) + transferred = True + break if not had_task_fc or transferred: # LLM finished without delegating (or transferred away); # nothing more for this wrapper to do. diff --git a/src/google/adk/workflow/_workflow.py b/src/google/adk/workflow/_workflow.py index dc9527cff0e..d022dcd8528 100644 --- a/src/google/adk/workflow/_workflow.py +++ b/src/google/adk/workflow/_workflow.py @@ -37,7 +37,6 @@ from ._dynamic_node_scheduler import DynamicNodeState from ._graph import EdgeItem from ._graph import Graph -from ._node_runner import NodeRunner from ._node_state import NodeState from ._node_status import NodeStatus from ._trigger import Trigger @@ -536,7 +535,7 @@ def _start_node_task( node_name: str, trigger: Trigger, ) -> None: - """Create NodeRunner and start asyncio task for a node.""" + """Start asyncio task for scheduling and executing a node.""" assert self.graph is not None @@ -602,22 +601,24 @@ async def return_ctx(): key ].isolation_scope - runner = NodeRunner( - node=node, - parent_ctx=ctx, - run_id=run_id, - use_as_output=is_terminal, - use_sub_branch=trigger.use_sub_branch, - override_branch=trigger.branch, - override_isolation_scope=self._compute_isolation_scope_for_node( - node, trigger, ctx, run_id - ), - ) resume_inputs = ( dict(node_state.resume_inputs) if node_state.resume_inputs else None ) loop_state.pending_tasks[node_name] = asyncio.create_task( - runner.run(node_input=trigger.input, resume_inputs=resume_inputs) + ctx._run_node_internal( + node, + node_input=trigger.input, + use_sub_branch=trigger.use_sub_branch, + override_branch=trigger.branch, + override_isolation_scope=self._compute_isolation_scope_for_node( + node, trigger, ctx, run_id + ), + return_ctx=True, + resume_inputs=resume_inputs, + run_id=run_id, + use_as_output=is_terminal, + skip_run_id_validation=True, + ) ) def _make_schedule_dynamic_node( diff --git a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py index e37fdfe67ce..b769eca55e1 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py @@ -1298,6 +1298,28 @@ def mock_refresh(req): assert headers["Authorization"] == "Bearer refreshed_token" + @pytest.mark.skipif(not AIO_SUPPORTED, reason="google.auth.aio not supported") + @pytest.mark.parametrize( + "existing_header_key", + ["Authorization", "authorization", "AUTHORIZATION", "authORIZATION"], + ) + @pytest.mark.asyncio + async def test_before_request_skips_refresh_if_authorization_header_exists_case_insensitive( + self, existing_header_key + ): + mock_creds = Mock() + mock_creds.expired = True + mock_creds.token = "new_token" + mock_creds.refresh = Mock() + + credentials = _RefreshableAsyncCredentials(mock_creds) + headers = {existing_header_key: "Bearer existing_token"} + + await credentials.before_request(None, "GET", "http://example.com", headers) + + mock_creds.refresh.assert_not_called() + assert headers == {existing_header_key: "Bearer existing_token"} + class TestGoogleAuthAsyncByteStream: diff --git a/tests/unittests/workflow/test_dynamic_node_scheduler.py b/tests/unittests/workflow/test_dynamic_node_scheduler.py index 6ef1b8f77f2..27244caeafb 100644 --- a/tests/unittests/workflow/test_dynamic_node_scheduler.py +++ b/tests/unittests/workflow/test_dynamic_node_scheduler.py @@ -698,456 +698,6 @@ async def test_runtime_schema_validation_content_fallback(): # Should not raise -# ========================================================================= -# __call__ — Agent Transfer logic -# ========================================================================= - - -@pytest.mark.asyncio -async def test_scheduler_handles_child_transfer(): - """Scheduler processes CHILD relationship by nesting next context.""" - # Arrange - from google.adk.agents.llm_agent import LlmAgent - from google.adk.events.event_actions import EventActions - from google.adk.workflow._dynamic_node_scheduler import DynamicNodeScheduler - from google.adk.workflow._workflow import _LoopState - - target = LlmAgent(name='target') - current = LlmAgent(name='current', sub_agents=[target]) - root = LlmAgent(name='root', sub_agents=[current]) - - ctx, _ = _make_parent_ctx() - ls = _LoopState() - scheduler = DynamicNodeScheduler(state=ls) - - child_ctx1 = MagicMock() - child_ctx1.node_path = 'wf/parent/current@1' - child_ctx1.parent_ctx = ctx - child_ctx1.actions = EventActions(transfer_to_agent='target') - child_ctx1.error = None - child_ctx1.interrupt_ids = set() - child_ctx1._invocation_context = ctx._invocation_context - child_ctx1._child_run_counters = {} - - child_ctx2 = MagicMock() - child_ctx2.node_path = 'wf/parent/current@1/target@1' - child_ctx2.parent_ctx = child_ctx1 - child_ctx2.actions = EventActions() - child_ctx2.error = None - child_ctx2.interrupt_ids = set() - child_ctx2._invocation_context = ctx._invocation_context - child_ctx2._child_run_counters = {} - - scheduler._run_node_internal = AsyncMock(side_effect=[child_ctx1, child_ctx2]) - - # Act - final_ctx = await scheduler( - ctx, - current, - 'input', - node_name='current', - run_id='1', - ) - - # Assert - assert final_ctx is child_ctx2 - calls = scheduler._run_node_internal.call_args_list - assert len(calls) == 2 - - # First call: current node - args1, kwargs1 = calls[0] - assert args1[0] is ctx - assert args1[1] is current - assert args1[2] == 'current' - assert args1[4] == '1' - - # Second call: target node (transferred CHILD) - args2, kwargs2 = calls[1] - assert args2[0] is child_ctx1 - assert args2[1] is target - assert args2[2] == 'target' - assert args2[4] == '1' - - -@pytest.mark.asyncio -async def test_scheduler_handles_sibling_transfer(): - """Scheduler processes SIBLING relationship by sharing parent context.""" - # Arrange - from google.adk.agents.llm_agent import LlmAgent - from google.adk.events.event_actions import EventActions - from google.adk.workflow._dynamic_node_scheduler import DynamicNodeScheduler - from google.adk.workflow._workflow import _LoopState - - current = LlmAgent(name='current') - target = LlmAgent(name='target') - root = LlmAgent(name='root', sub_agents=[current, target]) - - ctx, _ = _make_parent_ctx() - ls = _LoopState() - scheduler = DynamicNodeScheduler(state=ls) - - child_ctx1 = MagicMock() - child_ctx1.node_path = 'wf/parent/current@1' - child_ctx1.parent_ctx = ctx - child_ctx1.actions = EventActions(transfer_to_agent='target') - child_ctx1.error = None - child_ctx1.interrupt_ids = set() - child_ctx1._invocation_context = ctx._invocation_context - child_ctx1._child_run_counters = {} - - child_ctx2 = MagicMock() - child_ctx2.node_path = 'wf/parent/target@1' - child_ctx2.parent_ctx = ctx - child_ctx2.actions = EventActions() - child_ctx2.error = None - child_ctx2.interrupt_ids = set() - child_ctx2._invocation_context = ctx._invocation_context - child_ctx2._child_run_counters = {} - - scheduler._run_node_internal = AsyncMock(side_effect=[child_ctx1, child_ctx2]) - - # Act - final_ctx = await scheduler( - ctx, - current, - 'input', - node_name='current', - run_id='1', - ) - - # Assert - assert final_ctx is child_ctx2 - calls = scheduler._run_node_internal.call_args_list - assert len(calls) == 2 - - # Second call: target node (transferred SIBLING) - args2, kwargs2 = calls[1] - assert args2[0] is ctx - assert args2[1] is target - assert args2[2] == 'target' - assert args2[4] == '1' - - -@pytest.mark.asyncio -async def test_scheduler_handles_parent_transfer(): - """Scheduler processes PARENT relationship by truncating parent context.""" - # Arrange - from google.adk.agents.llm_agent import LlmAgent - from google.adk.events.event_actions import EventActions - from google.adk.workflow._dynamic_node_scheduler import DynamicNodeScheduler - from google.adk.workflow._workflow import _LoopState - - current = LlmAgent(name='current') - target = LlmAgent(name='target', sub_agents=[current]) - root = LlmAgent(name='root', sub_agents=[target]) - - ctx, _ = _make_parent_ctx() - ctx._child_run_counters = {'target': 1} - ls = _LoopState() - scheduler = DynamicNodeScheduler(state=ls) - - target_ctx = MagicMock() - target_ctx.node_path = 'wf/parent/target@1' - target_ctx.node = target - target_ctx.parent_ctx = ctx - target_ctx._invocation_context = ctx._invocation_context - target_ctx._child_run_counters = {} - - child_ctx1 = MagicMock() - child_ctx1.node_path = 'wf/parent/target@1/current@1' - child_ctx1.node = current - child_ctx1.parent_ctx = target_ctx - child_ctx1.actions = EventActions(transfer_to_agent='target') - child_ctx1.error = None - child_ctx1.interrupt_ids = set() - child_ctx1._invocation_context = ctx._invocation_context - child_ctx1._child_run_counters = {} - - child_ctx2 = MagicMock() - child_ctx2.node_path = 'wf/parent/target@2' - child_ctx2.parent_ctx = ctx - child_ctx2.actions = EventActions() - child_ctx2.error = None - child_ctx2.interrupt_ids = set() - child_ctx2._invocation_context = ctx._invocation_context - child_ctx2._child_run_counters = {} - - scheduler._run_node_internal = AsyncMock(side_effect=[child_ctx1, child_ctx2]) - - # Act - final_ctx = await scheduler( - target_ctx, - current, - 'input', - node_name='current', - run_id='1', - ) - - # Assert - assert final_ctx is child_ctx2 - calls = scheduler._run_node_internal.call_args_list - assert len(calls) == 2 - - # Second call: target node (transferred ANCESTOR) - args2, kwargs2 = calls[1] - assert args2[0] is ctx - assert args2[1] is target - assert args2[2] == 'target' - assert args2[4] == '2' - - -@pytest.mark.asyncio -async def test_scheduler_raises_value_error_on_self_transfer(): - """Scheduler raises ValueError when agent attempts to transfer to itself.""" - # Arrange - from google.adk.agents.llm_agent import LlmAgent - from google.adk.events.event_actions import EventActions - from google.adk.workflow._dynamic_node_scheduler import DynamicNodeScheduler - from google.adk.workflow._workflow import _LoopState - - current = LlmAgent(name='current') - root = LlmAgent(name='root', sub_agents=[current]) - - ctx, _ = _make_parent_ctx() - ls = _LoopState() - scheduler = DynamicNodeScheduler(state=ls) - - child_ctx = MagicMock() - child_ctx.node_path = 'wf/parent/current@1' - child_ctx.parent_ctx = ctx - child_ctx.actions = EventActions(transfer_to_agent='current') - child_ctx.error = None - child_ctx.interrupt_ids = set() - child_ctx._invocation_context = ctx._invocation_context - child_ctx._child_run_counters = {} - - scheduler._run_node_internal = AsyncMock(return_value=child_ctx) - - # Act & Assert - with pytest.raises(ValueError, match='cannot transfer to itself'): - await scheduler( - ctx, - current, - 'input', - node_name='current', - run_id='1', - ) - - -@pytest.mark.asyncio -async def test_scheduler_handles_parent_transfer_bypassed_on_resume(): - """Scheduler processes PARENT relationship when parent was bypassed on resume.""" - # Arrange - from google.adk.agents.llm_agent import LlmAgent - from google.adk.events.event_actions import EventActions - from google.adk.workflow._dynamic_node_scheduler import DynamicNodeScheduler - from google.adk.workflow._workflow import _LoopState - - current = LlmAgent(name='current') - target = LlmAgent(name='target', sub_agents=[current]) - root = LlmAgent(name='root', sub_agents=[target]) - - current.parent_agent = target - target.parent_agent = root - - ctx, _ = _make_parent_ctx() - ctx.node = None - ctx._child_run_counters = {'target': 1} - ls = _LoopState() - scheduler = DynamicNodeScheduler(state=ls) - - target_ctx = MagicMock() - target_ctx.node_path = 'current@1' - target_ctx.node = current - target_ctx.parent_ctx = ctx - target_ctx._invocation_context = ctx._invocation_context - target_ctx._child_run_counters = {} - - child_ctx1 = MagicMock() - child_ctx1.node_path = 'current@1' - child_ctx1.node = current - child_ctx1.parent_ctx = ctx - child_ctx1.actions = EventActions(transfer_to_agent='target') - child_ctx1.error = None - child_ctx1.interrupt_ids = set() - child_ctx1._invocation_context = ctx._invocation_context - child_ctx1._child_run_counters = {} - - child_ctx2 = MagicMock() - child_ctx2.node_path = 'target@2' - child_ctx2.parent_ctx = ctx - child_ctx2.actions = EventActions() - child_ctx2.error = None - child_ctx2.interrupt_ids = set() - child_ctx2._invocation_context = ctx._invocation_context - child_ctx2._child_run_counters = {} - - scheduler._run_node_internal = AsyncMock(side_effect=[child_ctx1, child_ctx2]) - - # Act - final_ctx = await scheduler( - ctx, - current, - 'input', - node_name='current', - run_id='1', - ) - - # Assert - assert final_ctx is child_ctx2 - calls = scheduler._run_node_internal.call_args_list - assert len(calls) == 2 - - # First call: current node - args1, kwargs1 = calls[0] - assert args1[0] is ctx - assert args1[1] is current - assert args1[2] == 'current' - assert args1[4] == '1' - - # Second call: target node (transferred PARENT) - args2, kwargs2 = calls[1] - assert args2[0] is ctx - assert args2[1] is target - assert args2[2] == 'target' - assert args2[4] == '2' - - -@pytest.mark.asyncio -async def test_scheduler_handles_three_layer_agent_transfers_round_trip(): - """Verify 3-layer agent transfers (Root -> Child -> Grandchild -> Child -> Root).""" - # Arrange - from google.adk.agents.llm_agent import LlmAgent - from google.adk.events.event_actions import EventActions - from google.adk.workflow._dynamic_node_scheduler import DynamicNodeScheduler - from google.adk.workflow._workflow import _LoopState - - grandchild = LlmAgent(name='grandchild') - child = LlmAgent(name='child', sub_agents=[grandchild]) - root = LlmAgent(name='root', sub_agents=[child]) - - grandchild.parent_agent = child - child.parent_agent = root - - # ctx is the root context (node = None, parent_ctx = None) - ctx = MagicMock() - ctx.node = None - ctx.parent_ctx = None - ctx.node_path = '' - ctx._child_run_counters = {'child': 1} - ctx._output_for_ancestors = [] - - ic = MagicMock() - ic.invocation_id = 'inv-1' - ic.session = MagicMock() - ic.session.state = {} - ic.session.events = [] - ic.run_config = None - ctx._invocation_context = ic - - # root_ctx is the parent context of child (node = root, parent_ctx = ctx) - root_ctx = MagicMock() - root_ctx.node = root - root_ctx.parent_ctx = ctx - root_ctx.node_path = 'root@1' - root_ctx._child_run_counters = {'child': 1} - root_ctx._invocation_context = ic - - ls = _LoopState() - scheduler = DynamicNodeScheduler(state=ls) - - # Step 1: root delegates to child - child_ctx1 = MagicMock() - child_ctx1.node_path = 'root@1/child@1' - child_ctx1.node = child - child_ctx1.parent_ctx = root_ctx - child_ctx1.actions = EventActions(transfer_to_agent='grandchild') - child_ctx1.error = None - child_ctx1.interrupt_ids = set() - child_ctx1._invocation_context = ic - child_ctx1._child_run_counters = {} - - # Step 2: child delegates to grandchild - grandchild_ctx1 = MagicMock() - grandchild_ctx1.node_path = 'root@1/child@1/grandchild@1' - grandchild_ctx1.node = grandchild - grandchild_ctx1.parent_ctx = child_ctx1 - grandchild_ctx1.actions = EventActions(transfer_to_agent='child') - grandchild_ctx1.error = None - grandchild_ctx1.interrupt_ids = set() - grandchild_ctx1._invocation_context = ic - grandchild_ctx1._child_run_counters = {} - - # Step 3: grandchild transfers back to child - child_ctx2 = MagicMock() - child_ctx2.node_path = 'root@1/child@2' - child_ctx2.node = child - child_ctx2.parent_ctx = root_ctx - child_ctx2.actions = EventActions(transfer_to_agent='root') - child_ctx2.error = None - child_ctx2.interrupt_ids = set() - child_ctx2._invocation_context = ic - child_ctx2._child_run_counters = {} - - # Step 4: child transfers back to root - ctx._child_run_counters = {'root': 1} - root_ctx2 = MagicMock() - root_ctx2.node_path = 'root@2' - root_ctx2.node = root - root_ctx2.parent_ctx = ctx - root_ctx2.actions = EventActions() - root_ctx2.error = None - root_ctx2.interrupt_ids = set() - root_ctx2._invocation_context = ic - root_ctx2._child_run_counters = {} - - scheduler._run_node_internal = AsyncMock( - side_effect=[child_ctx1, grandchild_ctx1, child_ctx2, root_ctx2] - ) - - # Act - final_ctx = await scheduler( - root_ctx, - child, - 'input', - node_name='child', - run_id='1', - ) - - # Assert - assert final_ctx is root_ctx2 - calls = scheduler._run_node_internal.call_args_list - assert len(calls) == 4 - - # 1st call: child (scheduled by root) - args1, kwargs1 = calls[0] - assert args1[0] is root_ctx - assert args1[1] is child - assert args1[2] == 'child' - assert args1[4] == '1' - - # 2nd call: grandchild (transferred CHILD from child) - args2, kwargs2 = calls[1] - assert args2[0] is child_ctx1 - assert args2[1] is grandchild - assert args2[2] == 'grandchild' - assert args2[4] == '1' - - # 3rd call: child (transferred PARENT from grandchild) - args3, kwargs3 = calls[2] - assert args3[0] is root_ctx - assert args3[1] is child - assert args3[2] == 'child' - assert args3[4] == '2' - - # 4th call: root (transferred PARENT from child) - args4, kwargs4 = calls[3] - assert args4[0] is ctx - assert args4[1] is root - assert args4[2] == 'root' - assert args4[4] == '2' - - # ========================================================================= # Replay Sequence Ordering preservation for Dynamic Nodes # ========================================================================= diff --git a/tests/unittests/workflow/test_workflow_hitl.py b/tests/unittests/workflow/test_workflow_hitl.py index 8a2558b1533..dbe485aca54 100644 --- a/tests/unittests/workflow/test_workflow_hitl.py +++ b/tests/unittests/workflow/test_workflow_hitl.py @@ -1858,15 +1858,16 @@ async def _run_impl( async def test_multiple_invocations_isolation(request: pytest.FixtureRequest): """Verify that a new invocation ignores events from a previous invocation.""" + run_counts = [] + class CounterNode(BaseNode): name: str = Field(default='counter_node') - run_count: int = Field(default=0) async def _run_impl( self, *, ctx: Context, node_input: Any ) -> AsyncGenerator[Any, None]: - self.run_count += 1 - yield f'Run {self.run_count}' + run_counts.append(1) + yield f'Run {len(run_counts)}' node_a = CounterNode() wf = Workflow(name='wf', edges=[(START, node_a)]) @@ -1883,7 +1884,7 @@ async def _run_impl( ): events1.append(event) - assert node_a.run_count == 1 + assert len(run_counts) == 1 # Invocation 2 (New invocation in SAME session) msg2 = types.Content(parts=[types.Part(text='go 2')], role='user') @@ -1894,7 +1895,7 @@ async def _run_impl( events2.append(event) # If isolation works, CounterNode should run AGAIN! - assert node_a.run_count == 2 + assert len(run_counts) == 2 @pytest.mark.asyncio diff --git a/tests/unittests/workflow/test_workflow_live.py b/tests/unittests/workflow/test_workflow_live.py index 697b5775efb..d5207d666dc 100644 --- a/tests/unittests/workflow/test_workflow_live.py +++ b/tests/unittests/workflow/test_workflow_live.py @@ -27,6 +27,7 @@ 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 @@ -39,6 +40,7 @@ class _MockNonLiveNode(BaseNode): 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) @@ -51,6 +53,8 @@ async def _run_impl( ) -> 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") @@ -72,29 +76,6 @@ async def _run_impl( yield Event(output=self.output_value) -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 - - 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: - self.child_output = await ctx.run_node( - self.child_node, node_input=node_input - ) - yield Event(output=f"{self.name}_output") - - # --- Live Workflow Unit Tests (TDD) --- @@ -468,6 +449,30 @@ async def test_nested_live_node_and_outer_live_node(): @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( @@ -531,7 +536,9 @@ async def test_dynamic_node_scheduling_of_live_node(): {"result": "DynamicLiveNode_output"}, "SchedulerNode_output", ] - assert scheduler_node.child_output == {"result": "DynamicLiveNode_output"} + assert scheduler_node.shared_state.get("child_output") == { + "result": "DynamicLiveNode_output" + } # Assert content events content_texts = [ @@ -663,7 +670,7 @@ async def test_single_turn_agent_runs_as_non_live_in_live_session(): outputs = [e.output for e in events if e.output is not None] assert outputs == ["initial_text_input", "capture_output"] - assert capture.actual_input == "SingleTurn_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 ( diff --git a/tests/unittests/workflow/test_workflow_schema.py b/tests/unittests/workflow/test_workflow_schema.py index 5fdfc706ffd..0574aad6d09 100644 --- a/tests/unittests/workflow/test_workflow_schema.py +++ b/tests/unittests/workflow/test_workflow_schema.py @@ -26,6 +26,7 @@ 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 @@ -630,9 +631,10 @@ async def _run_impl( msg = types.Content(parts=[types.Part(text='hello')], role='user') # We expect it to raise ValidationError - from pydantic import ValidationError - - with pytest.raises(ValidationError): + with pytest.raises( + ValueError, + match=r"Runtime schema validation failed for dynamic node 'node'", + ): async for event in runner.run_async( user_id='u', session_id=session.id, new_message=msg ):