From ade8577745bda23588bc0ac71f93a9d3343ec80a Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 6 Jul 2026 09:59:24 -0700 Subject: [PATCH 01/15] fix: drop nonexistent log_query tool from session_state_agent sample The instruction told the model to use a `log_query` tool that is not registered on the agent, so when the model followed it the run failed with "Tool 'log_query' not found". That made the integration_test presubmit, which runs this sample, flaky. Reword the instruction to just reply, matching the sample's documented output. Co-authored-by: George Weale PiperOrigin-RevId: 943350141 --- .../samples/context_management/session_state_agent/agent.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/contributing/samples/context_management/session_state_agent/agent.py b/contributing/samples/context_management/session_state_agent/agent.py index b4d665179ca..0c641b24b95 100644 --- a/contributing/samples/context_management/session_state_agent/agent.py +++ b/contributing/samples/context_management/session_state_agent/agent.py @@ -168,8 +168,8 @@ async def after_agent_callback(callback_context: CallbackContext): name='root_agent', description='a verification agent.', instruction=( - 'Log all users query with `log_query` tool. Must always remind user you' - ' cannot answer second query because your setup.' + 'Reply to the user. Must always remind user you cannot answer a second' + ' query because your setup.' ), model='gemini-3.5-flash', before_agent_callback=before_agent_callback, From 8a2a948bfb4c531d0d6d10592a6c5b56c47afd7e Mon Sep 17 00:00:00 2001 From: Cyrus M <42863660+cyrus-mz@users.noreply.github.com> Date: Mon, 6 Jul 2026 10:29:15 -0700 Subject: [PATCH 02/15] fix: Windows file artifact URI path normalization Merge https://github.com/google/adk-python/pull/5523 Closes #5486 ## Summary This PR fixes Windows handling of local `file://` artifact URIs. Both `google/adk/cli/service_registry.py` and `google/adk/artifacts/file_artifact_service.py` were passing URI-style paths into `Path(...)` without Windows URI-to-filesystem normalization. On Windows, this can produce malformed drive-relative paths such as `C:foo` instead of absolute paths like `C:\foo`. This change applies `url2pathname()` under `os.name == "nt"` in both locations so canonical file URIs like `file:///C:/...` are converted to proper native Windows paths before constructing `Path(...)`. ## Validation Manually tested on Windows with: 1. current working directory on `C:` and artifact target on `C:` - before the patch, the bug reproduced and the artifact root was created in the wrong location - after the patch, the artifact root was created in the intended location 2. current working directory on `C:` and artifact target on `D:` - before the patch, the current code could appear to work - after the patch, it still worked, now through correct URI-to-path normalization 3. paths containing spaces - after the patch, the artifact root was created correctly in the intended directory I also verified the drive-relative behavior separately with: ```python import ntpath print(ntpath.abspath("C:foo")) print(ntpath.abspath("D:foo")) ``` With current working directory `C:\Users\user1\projects\agent1`, this produced: ```text C:\Users\user1\projects\agent1\foo D:\foo ``` This matches the observed issue: without Windows-specific normalization, the malformed URI-derived path can behave like a drive-relative path such as `C:foo` instead of a proper absolute path like `C:\foo`. ## Notes I have not added a regression test in this PR yet. I can add one in a follow-up if maintainers prefer a specific test location for this path-conversion logic. Co-authored-by: George Weale COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/5523 from cyrus-mz:fix/windows-file-artifact-uri e93847b22ef3dd480836a851b2773e9e8640f17e PiperOrigin-RevId: 943367471 --- .../adk/artifacts/file_artifact_service.py | 6 +++- src/google/adk/cli/service_registry.py | 8 ++++- .../artifacts/test_artifact_service.py | 24 ++++++++++++++ tests/unittests/cli/test_service_registry.py | 31 +++++++++++++++++++ 4 files changed, 67 insertions(+), 2 deletions(-) diff --git a/src/google/adk/artifacts/file_artifact_service.py b/src/google/adk/artifacts/file_artifact_service.py index 6c74c53572d..417266f7fee 100644 --- a/src/google/adk/artifacts/file_artifact_service.py +++ b/src/google/adk/artifacts/file_artifact_service.py @@ -25,6 +25,7 @@ from typing import Union from urllib.parse import unquote from urllib.parse import urlparse +from urllib.request import url2pathname from google.genai import types from pydantic import alias_generators @@ -59,7 +60,10 @@ def _file_uri_to_path(uri: str) -> Optional[Path]: parsed = urlparse(uri) if parsed.scheme != "file": return None - return Path(unquote(parsed.path)) + path_str = unquote(parsed.path) + if os.name == "nt": + path_str = url2pathname(path_str) + return Path(path_str) _USER_NAMESPACE_PREFIX = "user:" diff --git a/src/google/adk/cli/service_registry.py b/src/google/adk/cli/service_registry.py index 517222d9324..f4f29d165ba 100644 --- a/src/google/adk/cli/service_registry.py +++ b/src/google/adk/cli/service_registry.py @@ -72,6 +72,7 @@ def my_session_factory(uri: str, **kwargs): from typing import Protocol from urllib.parse import unquote from urllib.parse import urlparse +from urllib.request import url2pathname from ..artifacts.base_artifact_service import BaseArtifactService from ..memory.base_memory_service import BaseMemoryService @@ -310,7 +311,12 @@ def file_artifact_factory(uri: str, **_): ) if not parsed_uri.path: raise ValueError("file:// artifact URIs must include a path component.") - artifact_path = Path(unquote(parsed_uri.path)) + + artifact_path_str = unquote(parsed_uri.path) + if os.name == "nt": + artifact_path_str = url2pathname(artifact_path_str) + + artifact_path = Path(artifact_path_str) return FileArtifactService(root_dir=artifact_path) registry.register_artifact_service("memory", memory_artifact_factory) diff --git a/tests/unittests/artifacts/test_artifact_service.py b/tests/unittests/artifacts/test_artifact_service.py index 8dfa29aeab1..0e04c8adef8 100644 --- a/tests/unittests/artifacts/test_artifact_service.py +++ b/tests/unittests/artifacts/test_artifact_service.py @@ -20,6 +20,7 @@ import enum import json from pathlib import Path +from types import SimpleNamespace from typing import Any from typing import Optional from typing import Union @@ -28,6 +29,7 @@ from urllib.parse import unquote from urllib.parse import urlparse +from google.adk.artifacts import file_artifact_service from google.adk.artifacts.base_artifact_service import ArtifactVersion from google.adk.artifacts.base_artifact_service import ensure_part from google.adk.artifacts.file_artifact_service import FileArtifactService @@ -1591,3 +1593,25 @@ async def test_save_load_empty_text_artifact( assert loaded is not None assert loaded.text == "" assert loaded.inline_data is None + + +def test_file_uri_to_path_normalizes_windows_file_uri(monkeypatch): + monkeypatch.setattr(file_artifact_service, "os", SimpleNamespace(name="nt")) + mocked_url2pathname = mock.Mock(return_value=r"C:\tmp\adk artifacts") + monkeypatch.setattr( + file_artifact_service, "url2pathname", mocked_url2pathname + ) + + result = file_artifact_service._file_uri_to_path( + "file:///C:/tmp/adk%20artifacts" + ) + + mocked_url2pathname.assert_called_once_with("/C:/tmp/adk artifacts") + assert result == Path(r"C:\tmp\adk artifacts") + + +def test_file_uri_to_path_returns_none_for_non_file_uri(): + assert ( + file_artifact_service._file_uri_to_path("gs://bucket/adk_artifacts") + is None + ) diff --git a/tests/unittests/cli/test_service_registry.py b/tests/unittests/cli/test_service_registry.py index 094c4ea4284..4af657ac28b 100644 --- a/tests/unittests/cli/test_service_registry.py +++ b/tests/unittests/cli/test_service_registry.py @@ -12,8 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from pathlib import Path +from types import SimpleNamespace +from unittest import mock from unittest.mock import patch +from google.adk.cli import service_registry import pytest @@ -124,6 +128,33 @@ def test_create_artifact_service_gcs(registry, mock_services): ) +def test_file_artifact_factory_normalizes_windows_file_uri(monkeypatch): + monkeypatch.setattr(service_registry, "os", SimpleNamespace(name="nt")) + mocked_url2pathname = mock.Mock(return_value=r"C:\tmp\adk artifacts") + monkeypatch.setattr(service_registry, "url2pathname", mocked_url2pathname) + + registry = service_registry.ServiceRegistry() + service_registry._register_builtin_services(registry) + + with mock.patch( + "google.adk.artifacts.file_artifact_service.FileArtifactService" + ) as mock_file_artifact_service: + registry.create_artifact_service("file:///C:/tmp/adk%20artifacts") + + mocked_url2pathname.assert_called_once_with("/C:/tmp/adk artifacts") + mock_file_artifact_service.assert_called_once_with( + root_dir=Path(r"C:\tmp\adk artifacts") + ) + + +def test_file_artifact_factory_rejects_non_local_authority(): + registry = service_registry.ServiceRegistry() + service_registry._register_builtin_services(registry) + + with pytest.raises(ValueError, match="local filesystem"): + registry.create_artifact_service("file://example.com/tmp/adk_artifacts") + + # Memory Service Tests @patch("google.adk.cli.utils.envs.load_dotenv_for_agent") def test_create_memory_service_rag( From 555171aa547ebec086a74d9a53dc11d699e03393 Mon Sep 17 00:00:00 2001 From: Shangjie Chen Date: Mon, 6 Jul 2026 10:45:22 -0700 Subject: [PATCH 03/15] refactor: Extract ReplayManager for workflow event rehydration and replay synchronization Extract `ReplayManager` to unify session event scanning, replay interception, and sequence barrier management across static and dynamic workflow nodes. Co-authored-by: Shangjie Chen PiperOrigin-RevId: 943376672 --- .../adk/workflow/_dynamic_node_scheduler.py | 46 +------ src/google/adk/workflow/_workflow.py | 61 +-------- .../adk/workflow/utils/_replay_manager.py | 117 ++++++++++++++++++ .../workflow/test_workflow_nested.py | 37 ------ .../workflow/utils/test_replay_manager.py | 102 +++++++++++++++ 5 files changed, 229 insertions(+), 134 deletions(-) create mode 100644 src/google/adk/workflow/utils/_replay_manager.py create mode 100644 tests/unittests/workflow/utils/test_replay_manager.py diff --git a/src/google/adk/workflow/_dynamic_node_scheduler.py b/src/google/adk/workflow/_dynamic_node_scheduler.py index e59e1ac624c..7841823e538 100644 --- a/src/google/adk/workflow/_dynamic_node_scheduler.py +++ b/src/google/adk/workflow/_dynamic_node_scheduler.py @@ -36,10 +36,9 @@ from ._schedule_dynamic_node import ScheduleDynamicNode from .utils._rehydration_utils import _ChildScanState from .utils._rehydration_utils import _reconstruct_node_states -from .utils._rehydration_utils import is_terminal_event from .utils._replay_interceptor import check_interception from .utils._replay_interceptor import create_mock_context -from .utils._replay_sequence_barrier import ReplaySequenceBarrier +from .utils._replay_manager import ReplayManager if TYPE_CHECKING: from ..agents.context import Context @@ -121,7 +120,7 @@ class DynamicNodeScheduler(ScheduleDynamicNode): def __init__(self, *, state: DynamicNodeState) -> None: self._state = state - self._parent_sequence_barriers: dict[str, ReplaySequenceBarrier] = {} + self._replay_manager = ReplayManager() async def __call__( self, @@ -163,9 +162,8 @@ async def __call__( # 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) + if parent_path: + self._replay_manager.prepare_parent_sequence_barrier(ctx, parent_path) # Runtime schema validation. if node_input is not None: @@ -220,8 +218,7 @@ async def __call__( # 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) + await self._replay_manager.advance_sequence(parent_path, key) return child_ctx @@ -299,8 +296,7 @@ async def _check_existing_run( # Chronological sequence barrier wait for replayed dynamic nodes 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: - await self._parent_sequence_barriers[parent_path].wait(key) + await self._replay_manager.wait_sequence(parent_path, key) return mock_ctx, True @@ -349,36 +345,6 @@ def _rehydrate_from_events(self, ctx: Context, node_path: str) -> None: logger.debug('node %s rehydrate end.', node_path) - def _scan_parent_child_sequence( - self, ctx: Context, parent_path: str - ) -> list[str]: - """Scan historical events and extract direct dynamic child completion sequence.""" - ic = ctx._invocation_context - base_path_builder = _NodePathBuilder.from_string(parent_path) - sequence: list[str] = [] - - for event in ic.session.events: - if event.invocation_id != ic.invocation_id: - continue - event_node_path = event.node_info.path or '' - event_path_builder = _NodePathBuilder.from_string(event_node_path) - - if not event_path_builder.is_descendant_of(base_path_builder): - continue - - child_path = base_path_builder.get_direct_child(event_path_builder) - if event_path_builder != child_path: - continue - - segment = child_path.leaf_segment - - if is_terminal_event(event): - if segment in sequence: - sequence.remove(segment) - sequence.append(segment) - - return sequence - # --- Execution --- async def _run_node_internal( diff --git a/src/google/adk/workflow/_workflow.py b/src/google/adk/workflow/_workflow.py index b6c92ab673c..ff952392e7b 100644 --- a/src/google/adk/workflow/_workflow.py +++ b/src/google/adk/workflow/_workflow.py @@ -41,10 +41,9 @@ from ._node_status import NodeStatus from ._trigger import Trigger from .utils._rehydration_utils import _ChildScanState -from .utils._rehydration_utils import _reconstruct_node_states -from .utils._rehydration_utils import is_terminal_event from .utils._replay_interceptor import check_interception from .utils._replay_interceptor import create_mock_context +from .utils._replay_manager import ReplayManager from .utils._replay_sequence_barrier import ReplaySequenceBarrier if TYPE_CHECKING: @@ -253,10 +252,9 @@ async def _run_impl( # --- SETUP: resume from events or start fresh --- # TODO: resume from checkpoint event. loop_state = _LoopState() - loop_state.recovered_executions, recovered_sequence = ( - self._scan_child_events(ctx) - ) - loop_state.sequence_barrier = ReplaySequenceBarrier(recovered_sequence) + replay_mgr = ReplayManager() + loop_state.recovered_executions, _ = replay_mgr.scan_workflow_events(ctx) + loop_state.sequence_barrier = replay_mgr.sequence_barrier if ctx.resume_inputs and not loop_state.recovered_executions: logger.warning( @@ -729,57 +727,6 @@ def _collect_remaining_interrupts(self, loop_state: _LoopState) -> None: # --- Resume --- - def _scan_child_events( - self, ctx: Context - ) -> tuple[dict[str, _ChildScanState], list[str]]: - """Scan session events and collect per-child state and completion sequence. - - Forward pass through events for this invocation. For each direct - child, tracks the latest run_id and accumulates output, - interrupt IDs, and resolved interrupt IDs. - - Returns: - Tuple of: - - dict of child_name → _ChildScanState. - - list of child_name@run_id in chronological order of completion. - """ - ic = ctx._invocation_context - raw_results = _reconstruct_node_states( - events=ic.session.events, - base_path=ctx.node_path, - group_by_direct_child=True, - invocation_id=ic.invocation_id, - ) - - from ..events._node_path_builder import _NodePathBuilder - - # Build chronological sequence of completions - sequence: list[str] = [] - base_path_builder = _NodePathBuilder.from_string(ctx.node_path) - - for event in ic.session.events: - if event.invocation_id != ic.invocation_id: - continue - - event_node_path = event.node_info.path or '' - event_path_builder = _NodePathBuilder.from_string(event_node_path) - - if not event_path_builder.is_descendant_of(base_path_builder): - continue - - child_path = base_path_builder.get_direct_child(event_path_builder) - segment: str = child_path.leaf_segment - - if is_terminal_event(event): - # Maintain unique segments ordered by their LAST terminal event. - # If a node interrupts in turn 1 and completes in turn 2, moving it - # to the end ensures we record the order of its final completion. - if segment in sequence: - sequence.remove(segment) - sequence.append(segment) - - return raw_results, sequence - # --- FINALIZE --- def _finalize(self, loop_state: _LoopState, ctx: Context) -> None: diff --git a/src/google/adk/workflow/utils/_replay_manager.py b/src/google/adk/workflow/utils/_replay_manager.py new file mode 100644 index 00000000000..7446009fb20 --- /dev/null +++ b/src/google/adk/workflow/utils/_replay_manager.py @@ -0,0 +1,117 @@ +# 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. + +"""ReplayManager — unified orchestrator for event rehydration, interception, and sequence barriers.""" + +from __future__ import annotations + +import logging + +from ...agents.context import Context +from ...events._node_path_builder import _NodePathBuilder +from ._rehydration_utils import _ChildScanState +from ._rehydration_utils import _reconstruct_node_states +from ._rehydration_utils import is_terminal_event +from ._replay_sequence_barrier import ReplaySequenceBarrier + +logger = logging.getLogger("google_adk." + __name__) + + +class ReplayManager: + """Unifies rehydration, replay interception, and sequence barrier synchronization across static and dynamic nodes.""" + + def __init__(self) -> None: + self._recovered_executions: dict[str, _ChildScanState] = {} + self._sequence_barrier: ReplaySequenceBarrier | None = None + self._parent_sequence_barriers: dict[str, ReplaySequenceBarrier] = {} + + @property + def recovered_executions(self) -> dict[str, _ChildScanState]: + """Recovered child states from event scan.""" + return self._recovered_executions + + @property + def sequence_barrier(self) -> ReplaySequenceBarrier | None: + """Sequence barrier for deterministic replay ordering.""" + return self._sequence_barrier + + def _scan_sequence( + self, ctx: Context, base_path: str, strict_direct_child: bool = False + ) -> list[str]: + """Extract chronological child completion sequence under base_path.""" + ic = ctx._invocation_context + base_path_builder = _NodePathBuilder.from_string(base_path) + sequence: list[str] = [] + + for event in ic.session.events: + if event.invocation_id != ic.invocation_id: + continue + + event_node_path = event.node_info.path or "" + event_path_builder = _NodePathBuilder.from_string(event_node_path) + + if not event_path_builder.is_descendant_of(base_path_builder): + continue + + child_path = base_path_builder.get_direct_child(event_path_builder) + if strict_direct_child and event_path_builder != child_path: + continue + + segment: str = child_path.leaf_segment + + if is_terminal_event(event): + if segment in sequence: + sequence.remove(segment) + sequence.append(segment) + + return sequence + + def scan_workflow_events( + self, ctx: Context + ) -> tuple[dict[str, _ChildScanState], list[str]]: + """Scan session events for direct child workflow nodes and initialize sequence barrier.""" + ic = ctx._invocation_context + raw_results = _reconstruct_node_states( + events=ic.session.events, + base_path=ctx.node_path, + group_by_direct_child=True, + invocation_id=ic.invocation_id, + ) + + sequence = self._scan_sequence( + ctx, ctx.node_path, strict_direct_child=False + ) + + self._recovered_executions = raw_results + self._sequence_barrier = ReplaySequenceBarrier(sequence) + return raw_results, sequence + + def prepare_parent_sequence_barrier( + self, ctx: Context, parent_path: str + ) -> ReplaySequenceBarrier: + """Ensure a sequence barrier is set up for dynamic nodes under parent_path.""" + if parent_path not in self._parent_sequence_barriers: + seq = self._scan_sequence(ctx, parent_path, strict_direct_child=True) + self._parent_sequence_barriers[parent_path] = ReplaySequenceBarrier(seq) + return self._parent_sequence_barriers[parent_path] + + async def advance_sequence(self, parent_path: str, key: str) -> None: + """Advance sequence barrier if initialized for parent_path.""" + if parent_path in self._parent_sequence_barriers: + self._parent_sequence_barriers[parent_path].check_and_advance(key) + + async def wait_sequence(self, parent_path: str, key: str) -> None: + """Wait for sequence barrier if initialized for parent_path.""" + if parent_path in self._parent_sequence_barriers: + await self._parent_sequence_barriers[parent_path].wait(key) diff --git a/tests/unittests/workflow/test_workflow_nested.py b/tests/unittests/workflow/test_workflow_nested.py index 8be4b48420d..9017200669b 100644 --- a/tests/unittests/workflow/test_workflow_nested.py +++ b/tests/unittests/workflow/test_workflow_nested.py @@ -1122,40 +1122,3 @@ async def _run_impl( if e.long_running_tool_ids: final_interrupts.update(e.long_running_tool_ids) assert not final_interrupts - - -@pytest.mark.asyncio -async def test_scan_child_events_ignores_descendant_run_id_resets(): - """_scan_child_events only resets run_id from direct child events.""" - from unittest.mock import MagicMock - - from google.adk.events.event import Event - from google.adk.events.event import NodeInfo - - # We create a Workflow instance to test its private method _scan_child_events. - wf = Workflow(name='wf', edges=[]) - - # Given a direct child event and a descendant event. - 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] - # _scan_child_events reads ctx.node_path to determine the base workflow path. - ctx.node_path = 'wf@1' - - children = wf._scan_child_events(ctx) - - # Assert child 'child' run_id remains '1' (not '2' from the descendant). - assert children[0]['child@1'].run_id == '1' diff --git a/tests/unittests/workflow/utils/test_replay_manager.py b/tests/unittests/workflow/utils/test_replay_manager.py new file mode 100644 index 00000000000..7e3e0de98f4 --- /dev/null +++ b/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' From 6b385e47ffebcd694078d2856fda7a2b8e69e2ce Mon Sep 17 00:00:00 2001 From: VectorPeak <73048950+VectorPeak@users.noreply.github.com> Date: Mon, 6 Jul 2026 11:10:48 -0700 Subject: [PATCH 04/15] fix: avoid mutating event validation input Merge https://github.com/google/adk-python/pull/6314 PiperOrigin-RevId: 943391685 --- src/google/adk/events/event.py | 1 + tests/unittests/events/test_event.py | 20 ++++++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/src/google/adk/events/event.py b/src/google/adk/events/event.py index b4304d30fc8..6445bcf1eff 100644 --- a/src/google/adk/events/event.py +++ b/src/google/adk/events/event.py @@ -171,6 +171,7 @@ def _accept_convenience_kwargs(cls, data: Any) -> Any: if not isinstance(data, dict): return data + data = dict(data) field_names: set[str] = set(cls.model_fields.keys()) for f in cls.model_fields.values(): if f.alias: diff --git a/tests/unittests/events/test_event.py b/tests/unittests/events/test_event.py index afcc64db7ed..8c1fb8794e0 100644 --- a/tests/unittests/events/test_event.py +++ b/tests/unittests/events/test_event.py @@ -16,6 +16,8 @@ """Unit tests for the helper methods on the Event class.""" +import copy + from google.adk.events.event import Event from google.adk.events.event import NodeInfo from google.adk.events.event_actions import EventActions @@ -372,6 +374,24 @@ def test_round_trip_via_content(self): assert restored.message is not None assert restored.message.parts[0].text == 'Hello!' + def test_model_validate_does_not_mutate_input_dict(self): + data = { + 'message': 'Hello!', + 'state': {'key': 'value'}, + 'route': 'next', + 'node_path': 'root.node', + } + original = copy.deepcopy(data) + + event = Event.model_validate(data) + + assert data == original + assert event.content is not None + assert event.content.parts[0].text == 'Hello!' + assert event.actions.state_delta == {'key': 'value'} + assert event.actions.route == 'next' + assert event.node_info.path == 'root.node' + class TestMessageWithOtherKwargs: """Tests message combined with other convenience kwargs.""" From 1509dcf3f4b9d8383fe3f596984623ddfb51b2cb Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 6 Jul 2026 11:19:19 -0700 Subject: [PATCH 05/15] fix: remove unused userinfo_endpoint from GoogleApiToolset OIDC config The userinfo_endpoint field is never read anywhere in ADK, so this hardcoded URL had no runtime effect. The OAuth flow only uses the authorization and token endpoints. Removing it. Co-authored-by: George Weale PiperOrigin-RevId: 943396380 --- src/google/adk/tools/google_api_tool/google_api_toolset.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/google/adk/tools/google_api_tool/google_api_toolset.py b/src/google/adk/tools/google_api_tool/google_api_toolset.py index 66e6dad88e4..35d130c7ab5 100644 --- a/src/google/adk/tools/google_api_tool/google_api_toolset.py +++ b/src/google/adk/tools/google_api_tool/google_api_toolset.py @@ -143,9 +143,6 @@ def _load_toolset_with_oidc_auth(self) -> OpenAPIToolset: 'https://accounts.google.com/o/oauth2/v2/auth' ), token_endpoint='https://oauth2.googleapis.com/token', - userinfo_endpoint=( - 'https://openidconnect.googleapis.com/v1/userinfo' - ), revocation_endpoint='https://oauth2.googleapis.com/revoke', token_endpoint_auth_methods_supported=[ 'client_secret_post', From cca8c5678d1066221f881debe7751ae08b84ab02 Mon Sep 17 00:00:00 2001 From: Dinesh Thumma <160909147+DineshThumma9@users.noreply.github.com> Date: Mon, 6 Jul 2026 11:20:18 -0700 Subject: [PATCH 06/15] feat(tools): exposed configurable parameter as property in McpToolset Merge https://github.com/google/adk-python/pull/4271 **Please ensure you have read the [contribution guide](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) before creating a pull request.** ### Link to Issue or Description of Change - Closes: #4270 **Problem:** When using MCP tools with McpToolset, configuration information for the object is not exposed as object properties or attributes meaning that 'private' attributes need to be accessed, e.g. toolset._connection_params rather than toolset.connection_params **Solution:** Exposed them as properties ### Testing Plan **Unit Tests:** - [x] I have added or updated unit tests for my change. - [x] All unit tests pass locally. **Manual End-to-End (E2E) Tests:** I have added new test cases also checked locally ### Checklist - [x] I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) document. - [x] I have performed a self-review of my own code. - [x] I have commented my code, particularly in hard-to-understand areas. - [x] I have added tests that prove my fix is effective or that my feature works. - [x] New and existing unit tests pass locally with my changes. - [x] I have manually tested my changes end-to-end. - [x] Any dependent changes have been merged and published in downstream modules. ### Additional context Co-authored-by: Haran Rajkumar COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/4271 from DineshThumma9:mcp-missing-attr 3d48401e1715096cd1fcbacac2ac6acca0751d36 PiperOrigin-RevId: 943396972 --- src/google/adk/tools/mcp_tool/mcp_toolset.py | 35 ++++++++++ .../tools/mcp_tool/test_mcp_toolset.py | 65 +++++++++++++++++++ 2 files changed, 100 insertions(+) diff --git a/src/google/adk/tools/mcp_tool/mcp_toolset.py b/src/google/adk/tools/mcp_tool/mcp_toolset.py index d736b00c5fa..c4a3d24f286 100644 --- a/src/google/adk/tools/mcp_tool/mcp_toolset.py +++ b/src/google/adk/tools/mcp_tool/mcp_toolset.py @@ -286,6 +286,41 @@ def _get_auth_headers( return headers + @property + def connection_params(self) -> Union[ + StdioServerParameters, + StdioConnectionParams, + SseConnectionParams, + StreamableHTTPConnectionParams, + ]: + return self._connection_params + + @property + def auth_scheme(self) -> Optional[AuthScheme]: + return self._auth_scheme + + @property + def auth_credential(self) -> Optional[AuthCredential]: + return self._auth_credential + + @property + def require_confirmation(self) -> Union[bool, Callable[..., bool]]: + return self._require_confirmation + + @property + def header_provider( + self, + ) -> Optional[ + Callable[ + [ReadonlyContext], Union[Dict[str, str], Awaitable[Dict[str, str]]] + ] + ]: + return self._header_provider + + @property + def errlog(self) -> TextIO: + return self._errlog + async def _execute_with_session( self, coroutine_func: Callable[[Any], Awaitable[T]], diff --git a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py index 0c1700ae5f6..0cdb72c96b7 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py @@ -22,6 +22,8 @@ from unittest.mock import Mock from fastapi.openapi.models import OAuth2 +from fastapi.openapi.models import OAuthFlowAuthorizationCode +from fastapi.openapi.models import OAuthFlows from google.adk.agents.readonly_context import ReadonlyContext from google.adk.auth.auth_credential import AuthCredential from google.adk.auth.auth_credential import AuthCredentialTypes @@ -96,6 +98,69 @@ def test_init_with_use_mcp_resources(self): ) assert toolset._use_mcp_resources is True + def test_connection_params(self): + """Test getting connection params.""" + toolset = McpToolset(connection_params=self.mock_stdio_params) + assert toolset.connection_params == self.mock_stdio_params + + def test_auth_scheme(self): + """Test getting auth scheme.""" + toolset = McpToolset(connection_params=self.mock_stdio_params) + assert toolset.auth_scheme is None + + def test_auth_credential(self): + """Test getting auth credential.""" + toolset = McpToolset(connection_params=self.mock_stdio_params) + assert toolset.auth_credential is None + + def test_error_log(self): + """Test getting error log.""" + toolset = McpToolset(connection_params=self.mock_stdio_params) + assert toolset.errlog == sys.stderr + + def test_auth_scheme_with_value(self): + """Test getting auth scheme when provided at initialization.""" + auth_scheme = OAuth2( + flows=OAuthFlows( + authorizationCode=OAuthFlowAuthorizationCode( + authorizationUrl="https://example.com/auth", + tokenUrl="https://example.com/token", + scopes={"read": "Read access"}, + ) + ) + ) + toolset = McpToolset( + connection_params=self.mock_stdio_params, + auth_scheme=auth_scheme, + ) + assert toolset.auth_scheme == auth_scheme + + def test_require_confirmation(self): + """Test getting require_confirmation flag.""" + toolset = McpToolset( + connection_params=self.mock_stdio_params, + require_confirmation=True, + ) + assert toolset.require_confirmation is True + + def test_header_provider(self): + """Test getting header_provider.""" + mock_header_provider = Mock() + toolset = McpToolset( + connection_params=self.mock_stdio_params, + header_provider=mock_header_provider, + ) + assert toolset.header_provider == mock_header_provider + + def test_auth_credential_with_value(self): + """Test getting auth credential when provided at initialization.""" + mock_credential = Mock(spec=AuthCredential) + toolset = McpToolset( + connection_params=self.mock_stdio_params, + auth_credential=mock_credential, + ) + assert toolset.auth_credential == mock_credential + def test_init_with_stdio_connection_params(self): """Test initialization with StdioConnectionParams.""" stdio_params = StdioConnectionParams( From cf91b8443fffca4561668a424f80e9c3feee2a78 Mon Sep 17 00:00:00 2001 From: Haran Rajkumar Date: Mon, 6 Jul 2026 11:37:27 -0700 Subject: [PATCH 07/15] feat(agents): add ManagedAgent backed by the Managed Agents API Add ManagedAgent(BaseAgent), which drives the Managed Agents interactions.create API directly. This first cut supports server-side tools only (ADK built-in tools and raw types.Tool configs); client-executed tools (FunctionTool/callables) and MCP are rejected with NotImplementedError. Multi-turn chaining reuses previous_interaction_id and recovers the sandbox environment across turns. Co-authored-by: Haran Rajkumar PiperOrigin-RevId: 943406685 --- src/google/adk/agents/__init__.py | 2 + src/google/adk/agents/_managed_agent.py | 367 ++++++++++ tests/integration/test_managed_agent.py | 88 +++ tests/unittests/agents/test_managed_agent.py | 698 +++++++++++++++++++ 4 files changed, 1155 insertions(+) create mode 100644 src/google/adk/agents/_managed_agent.py create mode 100644 tests/integration/test_managed_agent.py create mode 100644 tests/unittests/agents/test_managed_agent.py diff --git a/src/google/adk/agents/__init__.py b/src/google/adk/agents/__init__.py index 9d5749f50fc..e44aab18e82 100644 --- a/src/google/adk/agents/__init__.py +++ b/src/google/adk/agents/__init__.py @@ -16,6 +16,7 @@ from typing import Any from typing import TYPE_CHECKING +from ._managed_agent import ManagedAgent from .base_agent import BaseAgent from .base_agent_config import BaseAgentConfig from .context import Context @@ -42,6 +43,7 @@ 'Context', 'LlmAgent', 'LoopAgent', + 'ManagedAgent', 'McpInstructionProvider', 'ParallelAgent', 'SequentialAgent', diff --git a/src/google/adk/agents/_managed_agent.py b/src/google/adk/agents/_managed_agent.py new file mode 100644 index 00000000000..a7ea227128c --- /dev/null +++ b/src/google/adk/agents/_managed_agent.py @@ -0,0 +1,367 @@ +# 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 logging +from typing import Any +from typing import AsyncGenerator +from typing import Callable +from typing import Optional +from typing import TYPE_CHECKING +from typing import Union + +from google.genai import types +from google.genai.interactions import CreateAgentInteractionAgentConfigParam +from google.genai.interactions import CreateAgentInteractionEnvironmentParam +from google.genai.interactions import ToolParam +from pydantic import ConfigDict +from pydantic import Field +from pydantic import PrivateAttr + +from ..events.event import Event +from ..flows.llm_flows.interactions_processor import _find_previous_interaction_state +from ..models.interactions_utils import _convert_content_to_step +from ..models.interactions_utils import _create_interactions +from ..models.interactions_utils import build_interactions_request_log +from ..models.interactions_utils import convert_tools_config_to_interactions_format +from ..models.llm_request import LlmRequest +from ..models.llm_response import LlmResponse +from ..telemetry import tracer +from ..tools.base_tool import BaseTool +from ..tools.tool_context import ToolContext +from ..utils.context_utils import Aclosing +from ..utils.env_utils import is_enterprise_mode_enabled +from .base_agent import BaseAgent +from .invocation_context import InvocationContext +from .run_config import StreamingMode + +if TYPE_CHECKING: + from google.genai import Client + +logger = logging.getLogger('google_adk.' + __name__) + +# The Managed Agents / Interactions API is only served from the `global` +# location; regional endpoints reject these calls (e.g. "Resource setup has +# just started"). We pin it here so the agent works regardless of +# GOOGLE_CLOUD_LOCATION in the caller's environment. The project is still +# resolved from the environment / ADC as usual. +_MANAGED_AGENT_LOCATION = 'global' + + +def _resolve_client_location(api_client: Client) -> Optional[str]: + """Return the client's resolved location, or ``None`` if unavailable. + + google-genai 2.9.0 exposes no public accessor for a ``Client``'s location, so + we read the genai-internal ``client._api_client.location``. This is the single + remaining private dependency; the enterprise backend flag uses the public + ``Client.vertexai`` property. A missing value (e.g. test doubles) yields + ``None`` and is treated as acceptable. + """ + try: + # google-genai 2.9.0 has no public accessor for a Client's location. + return api_client._api_client.location # pylint: disable=protected-access + except AttributeError: + return None + + +def _validate_client_location(api_client: Client) -> None: + """Reject an injected enterprise client not targeting the `global` location. + + The Managed Agents API is only served from `global`. This check applies only + to enterprise (Vertex) clients: the Gemini Developer API has no location + concept, yet google-genai still stamps `GOOGLE_CLOUD_LOCATION` onto every + client's `_api_client.location`, so a Developer-API client must not be + rejected for it. We do not override a caller-supplied client, but a + non-`global` enterprise client cannot work, so we reject it loudly. The + backend is read from the public `Client.vertexai` property; the resolved + location has no public accessor in google-genai 2.9.0, so it is read from the + genai-internal `client._api_client.location` via `_resolve_client_location` + (an unresolvable location is treated as acceptable). + """ + # `Client.vertexai` is the public accessor (it returns False for the Gemini + # Developer API, which has no location concept); only enterprise (Vertex) + # clients have a meaningful location. + if not api_client.vertexai: + return + location = _resolve_client_location(api_client) + if isinstance(location, str) and location != _MANAGED_AGENT_LOCATION: + raise ValueError( + 'ManagedAgent requires an enterprise client configured for the' + f" '{_MANAGED_AGENT_LOCATION}' location; got location='{location}'." + ' The Managed Agents API is only served from' + f" '{_MANAGED_AGENT_LOCATION}'." + ) + + +class ManagedAgent(BaseAgent): + """An agent backed by the Managed Agents API (interactions.create). + + This agent calls the Managed Agents API directly from its execution loop. + In this version only server-side tools are supported: ADK built-in tools and + raw ``google.genai.types.Tool`` configs (the kinds the interactions converter + understands). Client-executed tools (FunctionTool/callables) and MCP are not + yet supported. + + ManagedAgent supports streaming interactions only. Interactions are always + created with ``background=True`` (required by the Managed Agents workflow) and + consumed over the streaming connection; non-streaming / background-polling + execution is not yet supported. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True, extra='forbid') + + agent_id: str + """The Managed Agent id (e.g. 'antigravity-preview-05-2026' or 'agents/ID').""" + + environment: Optional[CreateAgentInteractionEnvironmentParam] = None + """A sandbox environment spec (e.g. ``{'type': 'remote'}``) or an existing + environment id string to reuse across turns.""" + + agent_config: Optional[CreateAgentInteractionAgentConfigParam] = None + """Runtime configuration passed to interactions.create.""" + + tools: list[Union[types.Tool, BaseTool, Callable[..., Any]]] = Field( + default_factory=list + ) + """Server-side tools: ADK built-in tools or raw types.Tool configs.""" + + _api_client: Optional[Client] = PrivateAttr(default=None) + + def __init__( + self, *, api_client: Optional[Client] = None, **kwargs: Any + ) -> None: + super().__init__(**kwargs) + if api_client is not None: + _validate_client_location(api_client) + self._api_client = api_client + + @property + def api_client(self) -> Client: + """The genai client, lazily created if none was injected. + + The backend is resolved from the environment + (``GOOGLE_GENAI_USE_ENTERPRISE`` or the legacy + ``GOOGLE_GENAI_USE_VERTEXAI``), matching google-genai semantics; the + no-env default is the Gemini Developer API. The enterprise backend is + pinned to the ``global`` location (the Managed Agents API is only served + from ``global``); the Developer API takes no ``location`` (it is + meaningless there). + """ + if self._api_client is None: + from google.genai import Client + + if is_enterprise_mode_enabled(): + self._api_client = Client( + enterprise=True, location=_MANAGED_AGENT_LOCATION + ) + else: + self._api_client = Client(enterprise=False) + return self._api_client + + async def _resolve_backend_tools( + self, ctx: InvocationContext + ) -> list[ToolParam]: + """Resolve self.tools into interaction ToolParams (server-side only). + + Raw types.Tool configs are passed through; ADK built-in tools are processed + into native tool configs. Client-executed tools (FunctionTool/callables) and + MCP tools are rejected. + """ + # Built-in tools are resolved in "managed agent" mode: the request carries + # the internal _is_managed_agent flag (and no model), so tools that normally + # gate on a Gemini model still resolve. Nothing here is sent to the API; the + # real call uses ``agent=self.agent_id``. + llm_request = LlmRequest(config=types.GenerateContentConfig()) + llm_request._is_managed_agent = True + tool_context = ToolContext(ctx) + + for tool in self.tools: + if isinstance(tool, types.Tool): + if tool.mcp_servers: + raise NotImplementedError( + 'Raw mcp_servers tools are not yet supported by ManagedAgent ' + '(MCP is deferred).' + ) + if tool.function_declarations: + raise NotImplementedError( + 'client-executed tools are not yet supported by ManagedAgent: ' + f'{tool!r}' + ) + if not ( + tool.google_search + or tool.code_execution + or tool.url_context + or tool.computer_use + ): + raise NotImplementedError( + 'Unsupported raw types.Tool for ManagedAgent; supported ' + 'server-side fields are google_search, code_execution, ' + f'url_context, computer_use: {tool!r}' + ) + llm_request.config.tools = (llm_request.config.tools or []) + [tool] + continue + + if not isinstance(tool, BaseTool): + raise NotImplementedError( + 'client-executed tools are not yet supported by ManagedAgent: ' + f'{tool!r}' + ) + + # Built-in (server-side) tools mutate config.tools directly; tools that + # register a function declaration via append_tools grow tools_dict and are + # therefore client-executed. + before = len(llm_request.tools_dict) + await tool.process_llm_request( + tool_context=tool_context, llm_request=llm_request + ) + if len(llm_request.tools_dict) > before: + # The tool registered a function declaration -> client-executed. + raise NotImplementedError( + 'client-executed tools are not yet supported by ManagedAgent: ' + f'{tool.name}' + ) + + return convert_tools_config_to_interactions_format(llm_request.config) + + def _response_to_event( + self, ctx: InvocationContext, llm_response: LlmResponse + ) -> Event: + """Map a streamed LlmResponse to an ADK Event authored by this agent.""" + base_event = Event( + invocation_id=ctx.invocation_id, + author=self.name, + branch=ctx.branch, + ) + return Event.model_validate({ + **base_event.model_dump(exclude_none=True), + **llm_response.model_dump(exclude_none=True), + }) + + def _error_event( + self, + ctx: InvocationContext, + *, + error_code: str, + error_message: str, + ) -> Event: + """Build a terminal error event authored by this agent. + + Always sets ``turn_complete=True`` so the Runner receives a terminal event + even when the interactions call/stream fails. + """ + return Event( + invocation_id=ctx.invocation_id, + author=self.name, + branch=ctx.branch, + error_code=error_code, + error_message=error_message, + turn_complete=True, + ) + + async def _run_async_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + # Lazy import: google.genai is heavy, so only `types` is imported at module + # level (see CheckGoogleGenaiLazyImport / base_llm_flow.run_live). + from google.genai import errors + + # Recovery and tool resolution run outside the try so config errors (e.g. + # unsupported tools) surface loudly rather than becoming an error event. + prev_interaction_id, prev_environment_id = _find_previous_interaction_state( + ctx.session.events, + agent_name=self.name, + current_branch=ctx.branch, + ) + + environment = prev_environment_id or self.environment + + input_steps = ( + _convert_content_to_step(ctx.user_content) if ctx.user_content else [] + ) + interaction_tools = await self._resolve_backend_tools(ctx) + + create_kwargs: dict[str, Any] = { + 'agent': self.agent_id, + 'input': input_steps, + # The Managed Agents interactions workflow (server-side tools + remote + # environment) requires background execution. ManagedAgent supports + # streaming only, so the background result is consumed via the open SSE + # stream (stream=True at the _create_interactions call site below). + 'background': True, + } + if interaction_tools: + create_kwargs['tools'] = interaction_tools + if environment is not None: + create_kwargs['environment'] = environment + if self.agent_config is not None: + create_kwargs['agent_config'] = self.agent_config + if prev_interaction_id: + create_kwargs['previous_interaction_id'] = prev_interaction_id + + logger.info( + 'Sending request via interactions API, agent: %s, stream: %s, ' + 'previous_interaction_id: %s, environment: %s', + self.agent_id, + True, + prev_interaction_id, + environment, + ) + logger.debug( + build_interactions_request_log( + model=self.agent_id, + input_steps=input_steps, + system_instruction=None, + tools=interaction_tools if interaction_tools else None, + generation_config=None, + previous_interaction_id=prev_interaction_id, + stream=True, + ) + ) + + try: + with tracer.start_as_current_span('managed_agent_interaction'): + async with Aclosing( + _create_interactions( + self.api_client, create_kwargs=create_kwargs, stream=True + ) + ) as agen: + async for llm_response in agen: + # ManagedAgent always streams from the server, but only surface + # intermediate partials to the caller in SSE mode. In non-streaming + # mode (the default) emit just the non-partial events (the + # aggregated final event, plus any error event), mirroring + # base_llm_flow's behavior for LlmAgent. + if ( + ctx.run_config is not None + and ctx.run_config.streaming_mode == StreamingMode.SSE + ) or not llm_response.partial: + yield self._response_to_event(ctx, llm_response) + except errors.APIError as e: + # Surface the backend's real status/code (e.g. RESOURCE_EXHAUSTED) instead + # of a blanket UNKNOWN_ERROR, mirroring the status=='failed' interaction + # path and base_llm_flow's APIError handling. + logger.exception('ManagedAgent interaction failed with backend API error') + yield self._error_event( + ctx, + error_code=e.status or 'UNKNOWN_ERROR', + error_message=e.message or str(e), + ) + except Exception as e: # pylint: disable=broad-except + # Top-level safety net: any other failure still becomes a terminal error + # event so the Runner never hangs. + logger.exception('ManagedAgent interaction failed') + yield self._error_event( + ctx, error_code='UNKNOWN_ERROR', error_message=str(e) + ) diff --git a/tests/integration/test_managed_agent.py b/tests/integration/test_managed_agent.py new file mode 100644 index 00000000000..befc57ab5a9 --- /dev/null +++ b/tests/integration/test_managed_agent.py @@ -0,0 +1,88 @@ +# 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. + +"""Live integration tests for ManagedAgent. + +Assumes tests/integration/.env is present (auto-loaded by conftest.py) and that +auth (ADC) is configured. Run explicitly: + + pytest tests/integration/test_managed_agent.py -v -s +""" + +from __future__ import annotations + +from google.adk.agents import ManagedAgent +from google.adk.runners import Runner +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.tools import google_search +from google.adk.utils.context_utils import Aclosing +from google.genai import types +import pytest + +_AGENT_ID = 'antigravity-preview-05-2026' + + +async def _run_turn(runner, session, text: str) -> list: + events = [] + async with Aclosing( + runner.run_async( + user_id='test_user', + session_id=session.id, + new_message=types.Content( + role='user', parts=[types.Part.from_text(text=text)] + ), + ) + ) as agen: + async for event in agen: + events.append(event) + return events + + +def _joined_text(events) -> str: + return ' '.join( + part.text + for e in events + if e.content and e.content.parts + for part in e.content.parts + if part.text + ) + + +@pytest.mark.asyncio +async def test_google_search_project_hail_mary(): + agent = ManagedAgent( + name='managed_search_agent', + agent_id=_AGENT_ID, + environment={'type': 'remote'}, + tools=[google_search], + ) + session_service = InMemorySessionService() + runner = Runner( + app_name='managed_agent_it', + agent=agent, + session_service=session_service, + ) + session = await session_service.create_session( + app_name='managed_agent_it', user_id='test_user' + ) + + events = await _run_turn( + runner, session, 'Who plays Rocky in the movie Project Hail Mary?' + ) + + answer = _joined_text(events) + print('\n=== ManagedAgent answer ===\n', answer) + assert ( + 'james ortiz' in answer.lower() + ), f'expected the grounded answer to contain "James Ortiz"; got: {answer!r}' diff --git a/tests/unittests/agents/test_managed_agent.py b/tests/unittests/agents/test_managed_agent.py new file mode 100644 index 00000000000..7bb2841a1c3 --- /dev/null +++ b/tests/unittests/agents/test_managed_agent.py @@ -0,0 +1,698 @@ +# 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 unittest.mock import MagicMock + +from google.adk.agents._managed_agent import ManagedAgent +from google.adk.agents.run_config import RunConfig +from google.adk.agents.run_config import StreamingMode +from google.adk.events.event import Event +from google.adk.models.llm_response import LlmResponse +from google.adk.tools import google_search +from google.adk.tools.function_tool import FunctionTool +from google.genai import types +from google.genai import types as genai_types +import pytest + + +class _FakeClient: + vertexai = False + + +class _FakeApiClient: + + def __init__(self, location, vertexai=None): + self.location = location + self.vertexai = vertexai + + +class _FakeClientWithLocation: + """Mimics a genai Client: public vertexai, private _api_client.location.""" + + def __init__(self, location, vertexai=None): + self.vertexai = bool(vertexai) + self._api_client = _FakeApiClient(location, vertexai) + + +def test_construction_sets_fields_and_injectable_client(): + client = _FakeClient() + agent = ManagedAgent( + name='mgr', + agent_id='antigravity-preview-05-2026', + environment={'type': 'remote'}, + api_client=client, + ) + + assert agent.name == 'mgr' + assert agent.agent_id == 'antigravity-preview-05-2026' + assert agent.environment == {'type': 'remote'} + assert agent.tools == [] + # Injected client is returned without constructing a real genai client. + assert agent.api_client is client + + +def test_lazy_client_enterprise_uses_global_location(monkeypatch): + import google.genai as genai + + monkeypatch.setenv('GOOGLE_GENAI_USE_ENTERPRISE', '1') + captured = {} + + def _fake_client(**kwargs): + captured.update(kwargs) + return _FakeClient() + + monkeypatch.setattr(genai, 'Client', _fake_client) + + agent = ManagedAgent(name='mgr', agent_id='agents/a') + _ = agent.api_client # triggers lazy construction + + assert captured['enterprise'] is True + assert captured['location'] == 'global' + + +def test_lazy_client_dev_api_omits_location(monkeypatch): + import google.genai as genai + + monkeypatch.setenv('GOOGLE_GENAI_USE_ENTERPRISE', '0') + captured = {} + + def _fake_client(**kwargs): + captured.update(kwargs) + return _FakeClient() + + monkeypatch.setattr(genai, 'Client', _fake_client) + + agent = ManagedAgent(name='mgr', agent_id='agents/a') + _ = agent.api_client # triggers lazy construction + + assert captured['enterprise'] is False + assert 'location' not in captured + + +def test_injected_non_global_enterprise_client_raises(): + client = _FakeClientWithLocation('us-central1', vertexai=True) + + with pytest.raises(ValueError, match='global'): + ManagedAgent(name='mgr', agent_id='agents/a', api_client=client) + + +def test_injected_global_enterprise_client_is_accepted(): + client = _FakeClientWithLocation('global', vertexai=True) + + agent = ManagedAgent(name='mgr', agent_id='agents/a', api_client=client) + + assert agent.api_client is client + + +def test_injected_regional_dev_api_client_is_accepted(): + # Developer API clients have no meaningful location; genai still stamps + # GOOGLE_CLOUD_LOCATION onto _api_client.location, so a regional value must + # NOT be rejected for a non-enterprise client. + client = _FakeClientWithLocation('us-central1', vertexai=False) + + agent = ManagedAgent(name='mgr', agent_id='agents/a', api_client=client) + + assert agent.api_client is client + + +def test_injected_client_without_location_is_accepted(): + client = _FakeClient() + + agent = ManagedAgent(name='mgr', agent_id='agents/a', api_client=client) + + assert agent.api_client is client + + +def test_validate_uses_public_vertexai_property(): + # The enterprise decision must come from the PUBLIC `Client.vertexai` + # property, not the private `_api_client.vertexai`. This client reports + # enterprise via the public property while its private `_api_client.vertexai` + # is unset; a non-global location must therefore be rejected. + class _PublicVertexClient: + vertexai = True # public property says enterprise + + def __init__(self): + self._api_client = _FakeApiClient('us-central1', vertexai=None) + + with pytest.raises(ValueError, match='global'): + ManagedAgent( + name='mgr', agent_id='agents/a', api_client=_PublicVertexClient() + ) + + +def _ctx() -> Any: + # _resolve_backend_tools only needs an InvocationContext to build a + # ToolContext; a MagicMock satisfies the built-in tools used here. + return MagicMock() + + +def test_resolve_builtin_google_search(): + agent = ManagedAgent( + name='mgr', + agent_id='agents/a', + tools=[google_search], + api_client=_FakeClient(), + ) + + tool_params = asyncio.run(agent._resolve_backend_tools(_ctx())) + + assert {'type': 'google_search'} in tool_params + + +def test_resolve_raw_tool_passthrough(): + raw = types.Tool(url_context=types.UrlContext()) + agent = ManagedAgent( + name='mgr', agent_id='agents/a', tools=[raw], api_client=_FakeClient() + ) + + tool_params = asyncio.run(agent._resolve_backend_tools(_ctx())) + + assert {'type': 'url_context'} in tool_params + + +def test_resolve_rejects_raw_mcp_server(): + raw = types.Tool( + mcp_servers=[ + types.McpServer( + name='db', + streamable_http_transport=types.StreamableHttpTransport( + url='https://x' + ), + ) + ] + ) + agent = ManagedAgent( + name='mgr', agent_id='agents/a', tools=[raw], api_client=_FakeClient() + ) + + with pytest.raises(NotImplementedError, match='mcp'): + asyncio.run(agent._resolve_backend_tools(_ctx())) + + +def test_resolve_rejects_function_tool(): + def my_fn(x: str) -> str: + return x + + agent = ManagedAgent( + name='mgr', + agent_id='agents/a', + tools=[FunctionTool(func=my_fn)], + api_client=_FakeClient(), + ) + + with pytest.raises(NotImplementedError, match='client-executed'): + asyncio.run(agent._resolve_backend_tools(_ctx())) + + +def test_resolve_rejects_plain_callable(): + def my_fn(x: str) -> str: + return x + + agent = ManagedAgent( + name='mgr', agent_id='agents/a', tools=[my_fn], api_client=_FakeClient() + ) + + with pytest.raises(NotImplementedError, match='client-executed'): + asyncio.run(agent._resolve_backend_tools(_ctx())) + + +def test_resolve_rejects_raw_tool_with_function_declarations(): + raw = types.Tool( + function_declarations=[ + types.FunctionDeclaration(name='my_fn', description='d') + ] + ) + agent = ManagedAgent( + name='mgr', agent_id='agents/a', tools=[raw], api_client=_FakeClient() + ) + + with pytest.raises(NotImplementedError, match='client-executed'): + asyncio.run(agent._resolve_backend_tools(_ctx())) + + +def test_resolve_rejects_unsupported_raw_tool(): + raw = types.Tool(google_search_retrieval=types.GoogleSearchRetrieval()) + agent = ManagedAgent( + name='mgr', agent_id='agents/a', tools=[raw], api_client=_FakeClient() + ) + + with pytest.raises(NotImplementedError, match='Unsupported raw'): + asyncio.run(agent._resolve_backend_tools(_ctx())) + + +def test_resolve_combines_multiple_tools(): + agent = ManagedAgent( + name='mgr', + agent_id='agents/a', + tools=[google_search, types.Tool(url_context=types.UrlContext())], + api_client=_FakeClient(), + ) + + tool_params = asyncio.run(agent._resolve_backend_tools(_ctx())) + + assert {'type': 'google_search'} in tool_params + assert {'type': 'url_context'} in tool_params + + +def test_resolve_empty_tools_returns_empty(): + agent = ManagedAgent( + name='mgr', agent_id='agents/a', api_client=_FakeClient() + ) + + assert asyncio.run(agent._resolve_backend_tools(_ctx())) == [] + + +def test_resolve_passes_managed_agent_flag_and_no_model(): + from google.adk.tools.base_tool import BaseTool + + class _RecordingTool(BaseTool): + + def __init__(self): + super().__init__(name='rec', description='rec') + self.captured = {} + + async def process_llm_request(self, *, tool_context, llm_request): + self.captured['is_managed_agent'] = llm_request._is_managed_agent + self.captured['model'] = llm_request.model + + rec = _RecordingTool() + agent = ManagedAgent( + name='mgr', agent_id='agents/a', tools=[rec], api_client=_FakeClient() + ) + + asyncio.run(agent._resolve_backend_tools(_ctx())) + + assert rec.captured['is_managed_agent'] is True + assert rec.captured['model'] is None + + +class _RecordingInteractions: + + def __init__(self, responses_per_call): + self._responses_per_call = list(responses_per_call) + self.calls = [] + + async def create(self, **kwargs): + self.calls.append(kwargs) + responses = self._responses_per_call.pop(0) + + class _Iter: + + def __init__(self, items): + self._it = iter(items) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self._it) + except StopIteration as exc: + raise StopAsyncIteration from exc + + return _Iter(responses) + + +class _RecordingClient: + vertexai = False + + def __init__(self, responses_per_call): + self.aio = MagicMock() + self.aio.interactions = _RecordingInteractions(responses_per_call) + + +def _user_ctx(text, *, session_events=None, invocation_id='inv1', branch=None): + ctx = MagicMock() + ctx.user_content = genai_types.Content( + role='user', parts=[genai_types.Part(text=text)] + ) + ctx.invocation_id = invocation_id + ctx.branch = branch + ctx.session.events = session_events or [] + return ctx + + +def _make_llm_response(text, interaction_id, environment_id): + return LlmResponse( + content=genai_types.Content( + role='model', parts=[genai_types.Part(text=text)] + ), + interaction_id=interaction_id, + environment_id=environment_id, + ) + + +def _partial_text_response(text): + return LlmResponse( + content=genai_types.Content( + role='model', parts=[genai_types.Part(text=text)] + ), + partial=True, + ) + + +def _final_text_response(text): + return LlmResponse( + content=genai_types.Content( + role='model', parts=[genai_types.Part(text=text)] + ), + partial=False, + turn_complete=True, + ) + + +def test_run_async_yields_events_with_ids(monkeypatch): + from google.adk.agents import _managed_agent as mod + + async def _fake_stream(api_client, *, create_kwargs, stream): + yield _make_llm_response('Hello!', 'int_1', 'env_1') + + monkeypatch.setattr(mod, '_create_interactions', _fake_stream) + + agent = ManagedAgent( + name='mgr', agent_id='agents/a', api_client=_FakeClient() + ) + + async def _collect(): + out = [] + async for e in agent._run_async_impl(_user_ctx('hi')): + out.append(e) + return out + + events = asyncio.run(_collect()) + assert len(events) == 1 + assert events[0].author == 'mgr' + assert events[0].content.parts[0].text == 'Hello!' + assert events[0].interaction_id == 'int_1' + assert events[0].environment_id == 'env_1' + + +def test_run_async_recovers_previous_state(): + prior = Event( + author='mgr', interaction_id='int_prev', environment_id='env_prev' + ) + ctx = _user_ctx('again', session_events=[prior]) + + client = _RecordingClient([[]]) + agent = ManagedAgent(name='mgr', agent_id='agents/a', api_client=client) + + asyncio.run(_drain(agent._run_async_impl(ctx))) + + create_kwargs = client.aio.interactions.calls[0] + assert create_kwargs['previous_interaction_id'] == 'int_prev' + assert create_kwargs['environment'] == 'env_prev' + assert create_kwargs['agent'] == 'agents/a' + assert create_kwargs['stream'] is True + assert create_kwargs['background'] is True + + +def test_run_async_forwards_tools_and_agent_config(): + from google.adk.tools import google_search + + client = _RecordingClient([[]]) + agent = ManagedAgent( + name='mgr', + agent_id='agents/a', + tools=[google_search], + agent_config={'type': 'dynamic'}, + api_client=client, + ) + + asyncio.run(_drain(agent._run_async_impl(_user_ctx('hi')))) + + create_kwargs = client.aio.interactions.calls[0] + assert {'type': 'google_search'} in create_kwargs['tools'] + assert create_kwargs['agent_config'] == {'type': 'dynamic'} + + +def test_run_async_sets_background_true(): + client = _RecordingClient([[]]) + agent = ManagedAgent(name='mgr', agent_id='agents/a', api_client=client) + + asyncio.run(_drain(agent._run_async_impl(_user_ctx('hi')))) + + create_kwargs = client.aio.interactions.calls[0] + assert create_kwargs['background'] is True + + +def test_run_async_yields_multiple_events_in_order(monkeypatch): + from google.adk.agents import _managed_agent as mod + + async def _fake_stream(api_client, *, create_kwargs, stream): + yield _make_llm_response('one', 'int_1', 'env_1') + yield _make_llm_response('two', 'int_1', 'env_1') + + monkeypatch.setattr(mod, '_create_interactions', _fake_stream) + + agent = ManagedAgent( + name='mgr', agent_id='agents/a', api_client=_FakeClient() + ) + + events = asyncio.run(_drain_collect(agent._run_async_impl(_user_ctx('hi')))) + assert [e.content.parts[0].text for e in events] == ['one', 'two'] + + +def test_run_async_error_yields_error_event(monkeypatch): + from google.adk.agents import _managed_agent as mod + + async def _boom(api_client, *, create_kwargs, stream): + raise RuntimeError('api exploded') + yield # pragma: no cover + + monkeypatch.setattr(mod, '_create_interactions', _boom) + + agent = ManagedAgent( + name='mgr', agent_id='agents/a', api_client=_FakeClient() + ) + + events = asyncio.run(_drain_collect(agent._run_async_impl(_user_ctx('hi')))) + assert len(events) == 1 + assert events[0].author == 'mgr' + assert 'api exploded' in (events[0].error_message or '') + assert events[0].error_code == 'UNKNOWN_ERROR' + assert events[0].turn_complete is True + + +def test_run_async_api_error_surfaces_backend_status_and_message(monkeypatch): + from google.adk.agents import _managed_agent as mod + from google.genai import errors + + async def _boom(api_client, *, create_kwargs, stream): + raise errors.ClientError( + 429, + { + 'error': { + 'code': 429, + 'status': 'RESOURCE_EXHAUSTED', + 'message': 'Quota exceeded.', + } + }, + ) + yield # pragma: no cover + + monkeypatch.setattr(mod, '_create_interactions', _boom) + agent = ManagedAgent( + name='mgr', agent_id='agents/a', api_client=_FakeClient() + ) + + events = asyncio.run(_drain_collect(agent._run_async_impl(_user_ctx('hi')))) + assert len(events) == 1 + assert events[0].author == 'mgr' + assert events[0].error_code == 'RESOURCE_EXHAUSTED' + assert events[0].error_message == 'Quota exceeded.' + assert events[0].turn_complete is True + + +def test_run_async_uses_self_environment_when_no_prior(): + client = _RecordingClient([[]]) + agent = ManagedAgent( + name='mgr', + agent_id='agents/a', + environment={'type': 'remote'}, + api_client=client, + ) + + asyncio.run(_drain(agent._run_async_impl(_user_ctx('hi')))) + + create_kwargs = client.aio.interactions.calls[0] + assert create_kwargs['environment'] == {'type': 'remote'} + assert 'previous_interaction_id' not in create_kwargs + + +def test_run_async_raises_on_unsupported_tool(): + def my_fn(x: str) -> str: + return x + + agent = ManagedAgent( + name='mgr', agent_id='agents/a', tools=[my_fn], api_client=_FakeClient() + ) + + with pytest.raises(NotImplementedError, match='client-executed'): + asyncio.run(_drain(agent._run_async_impl(_user_ctx('hi')))) + + +def test_managed_agent_exported_from_package(): + import google.adk.agents as agents_pkg + + assert agents_pkg.ManagedAgent is ManagedAgent + assert 'ManagedAgent' in agents_pkg.__all__ + + +def test_run_async_non_streaming_suppresses_partials(monkeypatch): + from google.adk.agents import _managed_agent as mod + + async def _fake_stream(api_client, *, create_kwargs, stream): + yield _partial_text_response('thinking') + yield _partial_text_response('searching') + yield _final_text_response('Final answer.') + + monkeypatch.setattr(mod, '_create_interactions', _fake_stream) + agent = ManagedAgent( + name='mgr', agent_id='agents/a', api_client=_FakeClient() + ) + ctx = _user_ctx('hi') + ctx.run_config.streaming_mode = StreamingMode.NONE + + events = asyncio.run(_drain_collect(agent._run_async_impl(ctx))) + + assert len(events) == 1 + assert events[0].content.parts[0].text == 'Final answer.' + assert not events[0].partial + + +def test_run_async_sse_yields_all_partials(monkeypatch): + from google.adk.agents import _managed_agent as mod + + async def _fake_stream(api_client, *, create_kwargs, stream): + yield _partial_text_response('thinking') + yield _partial_text_response('searching') + yield _final_text_response('Final answer.') + + monkeypatch.setattr(mod, '_create_interactions', _fake_stream) + agent = ManagedAgent( + name='mgr', agent_id='agents/a', api_client=_FakeClient() + ) + ctx = _user_ctx('hi') + ctx.run_config.streaming_mode = StreamingMode.SSE + + events = asyncio.run(_drain_collect(agent._run_async_impl(ctx))) + + assert [e.content.parts[0].text for e in events] == [ + 'thinking', + 'searching', + 'Final answer.', + ] + + +def test_run_async_non_streaming_surfaces_error_event(monkeypatch): + from google.adk.agents import _managed_agent as mod + + async def _fake_stream(api_client, *, create_kwargs, stream): + yield _partial_text_response('thinking') + yield LlmResponse( + error_code='UNKNOWN_ERROR', error_message='boom', turn_complete=True + ) + + monkeypatch.setattr(mod, '_create_interactions', _fake_stream) + agent = ManagedAgent( + name='mgr', agent_id='agents/a', api_client=_FakeClient() + ) + ctx = _user_ctx('hi') + ctx.run_config.streaming_mode = StreamingMode.NONE + + events = asyncio.run(_drain_collect(agent._run_async_impl(ctx))) + + assert len(events) == 1 + assert events[0].error_code == 'UNKNOWN_ERROR' + assert events[0].error_message == 'boom' + + +def test_run_async_default_run_config_suppresses_partials(monkeypatch): + from google.adk.agents import _managed_agent as mod + + async def _fake_stream(api_client, *, create_kwargs, stream): + yield _partial_text_response('thinking') + yield _final_text_response('Final answer.') + + monkeypatch.setattr(mod, '_create_interactions', _fake_stream) + agent = ManagedAgent( + name='mgr', agent_id='agents/a', api_client=_FakeClient() + ) + ctx = _user_ctx('hi') + ctx.run_config = RunConfig() # default streaming_mode == StreamingMode.NONE + + events = asyncio.run(_drain_collect(agent._run_async_impl(ctx))) + + assert len(events) == 1 + assert events[0].content.parts[0].text == 'Final answer.' + + +def test_run_async_non_streaming_final_event_carries_grounding_and_usage(): + from google.genai.interactions import InteractionCompletedEvent + from google.genai.interactions import InteractionSseEventInteraction + from google.genai.interactions import StepDelta + from google.genai.interactions import Usage + + sse_events = [ + StepDelta( + event_type='step.delta', + index=0, + delta={ + 'type': 'google_search_call', + 'arguments': {'queries': ['q1']}, + }, + ), + StepDelta( + event_type='step.delta', + index=0, + delta={'type': 'text', 'text': 'Final answer.'}, + ), + InteractionCompletedEvent( + event_type='interaction.completed', + interaction=InteractionSseEventInteraction( + id='int_e2e', + status='completed', + steps=[], + usage=Usage(total_input_tokens=12, total_output_tokens=7), + ), + ), + ] + client = _RecordingClient([sse_events]) + agent = ManagedAgent(name='mgr', agent_id='agents/a', api_client=client) + ctx = _user_ctx('hi') + ctx.run_config.streaming_mode = StreamingMode.NONE + + events = asyncio.run(_drain_collect(agent._run_async_impl(ctx))) + + assert len(events) == 1 + final = events[0] + assert not final.partial + assert final.content.parts[-1].text == 'Final answer.' + assert final.grounding_metadata.web_search_queries == ['q1'] + assert final.usage_metadata.prompt_token_count == 12 + assert final.usage_metadata.candidates_token_count == 7 + + +async def _drain(agen): + async for _ in agen: + pass + + +async def _drain_collect(agen): + out = [] + async for e in agen: + out.append(e) + return out From 07aa1e09e2c2b3aebd9c3457f30e57b9f9b167b7 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Mon, 6 Jul 2026 11:58:48 -0700 Subject: [PATCH 08/15] fix(live): keep streaming tool yields from completing turns Fixes #5947. Co-authored-by: George Weale COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/6009 from he-yufeng:fix/live-stream-turn-complete f72e0b51b392340423a7e59972e2d467fd1ac295 PiperOrigin-RevId: 943418613 --- src/google/adk/agents/live_request_queue.py | 7 ++-- .../adk/flows/llm_flows/base_llm_flow.py | 10 +++--- src/google/adk/flows/llm_flows/functions.py | 4 ++- src/google/adk/models/base_llm_connection.py | 15 +++++++++ .../adk/models/gemini_llm_connection.py | 17 ++++++++-- .../agents/test_live_request_queue.py | 11 +++++++ .../llm_flows/test_base_llm_flow_realtime.py | 33 ++++++++++++++++++- .../models/test_gemini_llm_connection.py | 16 +++++++++ 8 files changed, 103 insertions(+), 10 deletions(-) diff --git a/src/google/adk/agents/live_request_queue.py b/src/google/adk/agents/live_request_queue.py index 8de2108acf5..9e6d20bd3da 100644 --- a/src/google/adk/agents/live_request_queue.py +++ b/src/google/adk/agents/live_request_queue.py @@ -55,6 +55,9 @@ class LiveRequest(BaseModel): close: bool = False """If set, close the queue. queue.shutdown() is only supported in Python 3.13+.""" + partial: bool = False + """If set, the content is a partial turn update that does not complete the current model turn.""" + class LiveRequestQueue: """Queue used to send LiveRequest in a live(bidirectional streaming) way.""" @@ -65,8 +68,8 @@ def __init__(self): def close(self): self._queue.put_nowait(LiveRequest(close=True)) - def send_content(self, content: types.Content): - self._queue.put_nowait(LiveRequest(content=content)) + def send_content(self, content: types.Content, partial: bool = False): + self._queue.put_nowait(LiveRequest(content=content, partial=partial)) def send_realtime(self, blob: types.Blob): self._queue.put_nowait(LiveRequest(blob=blob)) diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index 1a48b66c41a..7d5a487755a 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -811,9 +811,9 @@ async def _send_to_model( is_function_response = content.parts and any( part.function_response for part in content.parts ) - if not is_function_response: - if not content.role: - content.role = 'user' + if not is_function_response and not content.role: + content.role = 'user' + if not is_function_response and not live_request.partial: user_content_event = Event( id=Event.new_id(), invocation_id=invocation_context.invocation_id, @@ -824,7 +824,9 @@ async def _send_to_model( session=invocation_context.session, event=user_content_event, ) - await llm_connection.send_content(live_request.content) + await llm_connection._send_content( + live_request.content, partial=live_request.partial + ) async def _receive_from_model( self, diff --git a/src/google/adk/flows/llm_flows/functions.py b/src/google/adk/flows/llm_flows/functions.py index 6e40439d607..5a1bd3fee7d 100644 --- a/src/google/adk/flows/llm_flows/functions.py +++ b/src/google/adk/flows/llm_flows/functions.py @@ -995,7 +995,9 @@ async def run_tool_and_update_queue(tool, function_args, tool_context): updated_content = _build_function_response_content( tool, result, tool_context.function_call_id ) - invocation_context.live_request_queue.send_content(updated_content) + invocation_context.live_request_queue.send_content( + updated_content, partial=True + ) except asyncio.CancelledError: raise # Re-raise to properly propagate the cancellation diff --git a/src/google/adk/models/base_llm_connection.py b/src/google/adk/models/base_llm_connection.py index 80c041f5841..4b129717d97 100644 --- a/src/google/adk/models/base_llm_connection.py +++ b/src/google/adk/models/base_llm_connection.py @@ -51,6 +51,21 @@ async def send_content(self, content: types.Content): """ pass + async def _send_content( + self, content: types.Content, *, partial: bool = False + ) -> None: + """Sends content, optionally as a partial (non-turn-completing) update. + + The default implementation ignores ``partial`` and completes the turn. + Connections that support turn-based partial updates override this. + + Args: + content: The content to send to the model. + partial: Whether this content is a partial turn update that does not + complete the model turn. + """ + await self.send_content(content) + @abstractmethod async def send_realtime(self, blob: types.Blob): """Sends a chunk of audio or a frame of video to the model in realtime. diff --git a/src/google/adk/models/gemini_llm_connection.py b/src/google/adk/models/gemini_llm_connection.py index 1bf01825993..5e632285bae 100644 --- a/src/google/adk/models/gemini_llm_connection.py +++ b/src/google/adk/models/gemini_llm_connection.py @@ -104,6 +104,18 @@ async def send_content(self, content: types.Content): Args: content: The content to send to the model. """ + await self._send_content(content) + + async def _send_content( + self, content: types.Content, *, partial: bool = False + ) -> None: + """Sends content, optionally as a partial (non-turn-completing) update. + + Args: + content: The content to send to the model. + partial: Whether this content is a partial turn update that does not + complete the model turn. + """ assert content.parts if content.parts[0].function_response: # All parts have to be function responses. @@ -115,7 +127,8 @@ async def send_content(self, content: types.Content): else: logger.debug('Sending LLM new content %s', content) if ( - self._is_gemini_3_x_live + not partial + and self._is_gemini_3_x_live and len(content.parts) == 1 and content.parts[0].text ): @@ -127,7 +140,7 @@ async def send_content(self, content: types.Content): await self._gemini_session.send( input=types.LiveClientContent( turns=[content], - turn_complete=True, + turn_complete=not partial, ) ) diff --git a/tests/unittests/agents/test_live_request_queue.py b/tests/unittests/agents/test_live_request_queue.py index 1bcf92574b3..5334c8d68b2 100644 --- a/tests/unittests/agents/test_live_request_queue.py +++ b/tests/unittests/agents/test_live_request_queue.py @@ -40,6 +40,17 @@ def test_send_content(): mock_put_nowait.assert_called_once_with(LiveRequest(content=content)) +def test_send_content_sets_partial(): + queue = LiveRequestQueue() + content = MagicMock(spec=types.Content) + + with patch.object(queue._queue, "put_nowait") as mock_put_nowait: + queue.send_content(content, partial=True) + mock_put_nowait.assert_called_once_with( + LiveRequest(content=content, partial=True) + ) + + def test_send_realtime(): queue = LiveRequestQueue() blob = MagicMock(spec=types.Blob) diff --git a/tests/unittests/flows/llm_flows/test_base_llm_flow_realtime.py b/tests/unittests/flows/llm_flows/test_base_llm_flow_realtime.py index 054e06d5420..17f5aa59ab2 100644 --- a/tests/unittests/flows/llm_flows/test_base_llm_flow_realtime.py +++ b/tests/unittests/flows/llm_flows/test_base_llm_flow_realtime.py @@ -197,5 +197,36 @@ async def test_send_to_model_with_text_content(mock_llm_connection): await flow._send_to_model(mock_llm_connection, invocation_context) # Verify send_content was called instead of send_realtime - mock_llm_connection.send_content.assert_called_once_with(content) + mock_llm_connection._send_content.assert_called_once_with( + content, partial=False + ) mock_llm_connection.send_realtime.assert_not_called() + + +@pytest.mark.asyncio +async def test_send_to_model_with_intermediate_text_content( + mock_llm_connection, +): + agent = Agent(name='test_agent', model='mock') + invocation_context = await testing_utils.create_invocation_context( + agent=agent, user_content='' + ) + invocation_context.live_request_queue = LiveRequestQueue() + invocation_context.session_service.append_event = mock.AsyncMock() + + flow = TestBaseLlmFlow() + + content = types.Content( + role='user', parts=[types.Part.from_text(text='progress')] + ) + invocation_context.live_request_queue.send( + LiveRequest(content=content, partial=True) + ) + invocation_context.live_request_queue.close() + + await flow._send_to_model(mock_llm_connection, invocation_context) + + mock_llm_connection._send_content.assert_called_once_with( + content, partial=True + ) + invocation_context.session_service.append_event.assert_not_called() diff --git a/tests/unittests/models/test_gemini_llm_connection.py b/tests/unittests/models/test_gemini_llm_connection.py index 4d1b07530b4..e594f617740 100644 --- a/tests/unittests/models/test_gemini_llm_connection.py +++ b/tests/unittests/models/test_gemini_llm_connection.py @@ -124,6 +124,22 @@ async def test_send_content_text(gemini_connection, mock_gemini_session): assert call_args['input'].turn_complete is True +@pytest.mark.asyncio +async def test_send_content_text_can_keep_turn_open( + gemini_connection, mock_gemini_session +): + content = types.Content( + role='user', parts=[types.Part.from_text(text='progress')] + ) + + await gemini_connection._send_content(content, partial=True) + + mock_gemini_session.send.assert_called_once() + call_args = mock_gemini_session.send.call_args[1] + assert call_args['input'].turns == [content] + assert call_args['input'].turn_complete is False + + @pytest.mark.asyncio async def test_send_content_function_response( gemini_connection, mock_gemini_session From 3009ec1da7e5c8a3b4e6ab39d956eeb006eb2d39 Mon Sep 17 00:00:00 2001 From: Shangjie Chen Date: Mon, 6 Jul 2026 12:21:34 -0700 Subject: [PATCH 09/15] feat: Add sub-branch creation and segment appending to _BranchPath Introduce `append` method and `create_sub_branch` classmethod on `_BranchPath` to standardize dynamic branch path creation across workflow node runners, agent tools, and parallel agent execution. Co-authored-by: Shangjie Chen PiperOrigin-RevId: 943430872 --- src/google/adk/agents/parallel_agent.py | 7 ++- src/google/adk/events/_branch_path.py | 56 +++++++++++++++++++-- src/google/adk/tools/agent_tool.py | 6 ++- src/google/adk/workflow/_node_runner.py | 6 ++- tests/unittests/events/test_branch_path.py | 58 ++++++++++++++++++++++ 5 files changed, 122 insertions(+), 11 deletions(-) diff --git a/src/google/adk/agents/parallel_agent.py b/src/google/adk/agents/parallel_agent.py index cfa4788de42..8e98a30d867 100644 --- a/src/google/adk/agents/parallel_agent.py +++ b/src/google/adk/agents/parallel_agent.py @@ -25,6 +25,7 @@ from typing_extensions import deprecated from typing_extensions import override +from ..events._branch_path import _BranchPath from ..events.event import Event from ..utils.context_utils import Aclosing from .base_agent import BaseAgent @@ -44,10 +45,8 @@ def _create_branch_ctx_for_sub_agent( """Create isolated branch for every sub-agent.""" invocation_context = invocation_context.model_copy() branch_suffix = f'{agent.name}.{sub_agent.name}' - invocation_context.branch = ( - f'{invocation_context.branch}.{branch_suffix}' - if invocation_context.branch - else branch_suffix + invocation_context.branch = _BranchPath.create_sub_branch( + invocation_context.branch, name=branch_suffix ) return invocation_context diff --git a/src/google/adk/events/_branch_path.py b/src/google/adk/events/_branch_path.py index cc392e2faba..9847fd21f55 100644 --- a/src/google/adk/events/_branch_path.py +++ b/src/google/adk/events/_branch_path.py @@ -37,11 +37,11 @@ def from_string(cls, path_str: str | None) -> _BranchPath: """Parses a _BranchPath from a dot-separated string representation.""" if not path_str: return cls([]) - return cls(path_str.split('.')) + return cls(path_str.split(".")) def __str__(self) -> str: """Returns the dot-separated string representation of the path.""" - return '.'.join(self._segments) + return ".".join(self._segments) def __eq__(self, other: object) -> bool: """Returns True if segments are equal.""" @@ -64,7 +64,7 @@ def run_ids(self) -> set[str]: """ ids = set() for segment in self._segments: - parts = segment.rsplit('@', 1) + parts = segment.rsplit("@", 1) if len(parts) > 1 and parts[1]: ids.add(parts[1]) return ids @@ -99,3 +99,53 @@ def common_prefix(paths: list[_BranchPath]) -> _BranchPath: else: break return _BranchPath(common_segments) + + def append( + self, + segment_or_path: str | _BranchPath, + run_id: str | None = None, + ) -> _BranchPath: + """Returns a new _BranchPath with segment(s) appended. + + Args: + segment_or_path: A segment name (str), dot-separated path (str), or another + _BranchPath instance to append. + run_id: Optional run ID (or function_call_id) to format segment as + 'name@run_id'. + """ + if isinstance(segment_or_path, _BranchPath): + if run_id is not None: + raise ValueError( + "run_id cannot be provided when segment_or_path is a _BranchPath" + " instance." + ) + return _BranchPath(self._segments + segment_or_path.segments) + + if run_id is not None: + if "." in segment_or_path: + raise ValueError( + "run_id cannot be provided when segment_or_path is a dot-separated" + " path." + ) + segment = f"{segment_or_path}@{run_id}" + return _BranchPath(self._segments + [segment]) + + new_segments = [s for s in segment_or_path.split(".") if s] + return _BranchPath(self._segments + new_segments) + + @classmethod + def create_sub_branch( + cls, + base_branch: str | None, + *, + name: str, + run_id: str | None = None, + ) -> str: + """Creates a new dot-separated branch path string by appending a segment. + + Example: + _BranchPath.create_sub_branch('parent', name='child', run_id='1') -> + 'parent.child@1' + _BranchPath.create_sub_branch(None, name='agent') -> 'agent' + """ + return str(cls.from_string(base_branch).append(name, run_id=run_id)) diff --git a/src/google/adk/tools/agent_tool.py b/src/google/adk/tools/agent_tool.py index 55fe7ff6038..28a4db0eee6 100644 --- a/src/google/adk/tools/agent_tool.py +++ b/src/google/adk/tools/agent_tool.py @@ -27,6 +27,7 @@ from . import _automatic_function_calling_util from ..agents.common_configs import AgentRefConfig +from ..events._branch_path import _BranchPath from ..features import FeatureName from ..features import is_feature_enabled from ..memory.in_memory_memory_service import InMemoryMemoryService @@ -364,8 +365,9 @@ async def run_async( # Align subagent branch scoping with node execution (Node as Tool) using function_call_id. fc_id = tool_context.function_call_id base_branch = tool_context.get_invocation_context().branch - segment = f'{self.agent.name}@{fc_id}' - tool_branch = f'{base_branch}.{segment}' if base_branch else segment + tool_branch = _BranchPath.create_sub_branch( + base_branch, name=self.agent.name, run_id=fc_id + ) try: return await tool_context.run_node( diff --git a/src/google/adk/workflow/_node_runner.py b/src/google/adk/workflow/_node_runner.py index e2d3a3dbe70..d5bf2cbe185 100644 --- a/src/google/adk/workflow/_node_runner.py +++ b/src/google/adk/workflow/_node_runner.py @@ -28,6 +28,7 @@ from typing import Any from typing import TYPE_CHECKING +from ..events._branch_path import _BranchPath from ..telemetry import node_tracing if TYPE_CHECKING: @@ -206,8 +207,9 @@ def _create_child_context( ) if self._use_sub_branch: - segment = f"{self._node.name}@{self._run_id}" - branch = f"{base_branch}.{segment}" if base_branch else segment + branch = _BranchPath.create_sub_branch( + base_branch, name=self._node.name, run_id=self._run_id + ) ic = ic.model_copy(update={"branch": branch}) elif self._override_branch is not None: ic = ic.model_copy(update={"branch": self._override_branch}) diff --git a/tests/unittests/events/test_branch_path.py b/tests/unittests/events/test_branch_path.py index 62042e85d2a..0d244718b82 100644 --- a/tests/unittests/events/test_branch_path.py +++ b/tests/unittests/events/test_branch_path.py @@ -153,6 +153,64 @@ def test_constructor_copies_segments_list(): assert path.segments == ["parent", "child"] +def test_append_single_segment_returns_new_path(): + """append adds a single segment to an existing path.""" + path = _BranchPath.from_string("parent") + new_path = path.append("child") + + assert new_path == _BranchPath.from_string("parent.child") + assert path == _BranchPath.from_string("parent") # Immutability check + + +def test_append_with_run_id_formats_segment(): + """append formats the segment as 'name@run_id' when run_id is provided.""" + path = _BranchPath.from_string("parent") + new_path = path.append("child", run_id="call_123") + + assert new_path == _BranchPath.from_string("parent.child@call_123") + + +def test_append_another_branch_path(): + """append combines segments from another _BranchPath instance.""" + path1 = _BranchPath.from_string("parent") + path2 = _BranchPath.from_string("child.grandchild") + new_path = path1.append(path2) + + assert new_path == _BranchPath.from_string("parent.child.grandchild") + + +def test_create_sub_branch_formats_string_correctly(): + """create_sub_branch constructs sub-branch strings safely.""" + # With base branch and run ID + res1 = _BranchPath.create_sub_branch( + "parent.sub", name="child", run_id="run_1" + ) + assert res1 == "parent.sub.child@run_1" + + # Without base branch (None or empty) + res2 = _BranchPath.create_sub_branch(None, name="agent", run_id="fc_456") + assert res2 == "agent@fc_456" + + # Dot-separated sub-path without run ID + res3 = _BranchPath.create_sub_branch("parent", name="agent.sub_agent") + assert res3 == "parent.agent.sub_agent" + + +def test_append_with_run_id_and_branch_path_raises_value_error(): + """append raises ValueError when run_id is provided with a _BranchPath.""" + path1 = _BranchPath.from_string("parent") + path2 = _BranchPath.from_string("child") + with pytest.raises(ValueError, match="run_id cannot be provided"): + path1.append(path2, run_id="123") + + +def test_append_with_run_id_and_dot_separated_path_raises_value_error(): + """append raises ValueError when run_id is provided with a dot-separated path.""" + path = _BranchPath.from_string("parent") + with pytest.raises(ValueError, match="run_id cannot be provided"): + path.append("child.sub", run_id="123") + + def test_run_ids_filters_out_empty_run_ids(): """run_ids filters out segments with empty run IDs (e.g. ending with '@').""" path = _BranchPath.from_string("parent@.child@2.node@") From 2aeb1e1b49b69d7d9a5384047ac3610188dca6c3 Mon Sep 17 00:00:00 2001 From: Nicolas Mota Date: Mon, 6 Jul 2026 14:02:17 -0700 Subject: [PATCH 10/15] feat: add model_input_context for transient context in LLM requests Merge https://github.com/google/adk-python/pull/5991 # Add Transient Model Input Context to RunConfig ### Link to Issue or Description of Change **1. Link to an existing issue (if applicable):** - Closes: #5990 **Problem:** Host applications sometimes need to provide request-scoped context to an agent for a single invocation without persisting that context into `session.events`. Before this change, callers had to either append synthetic session events or merge application context into the user message. Both approaches blur the boundary between durable conversation history and transient model input. **Solution:** This change adds `RunConfig.model_input_context`, a list of `google.genai.types.Content` values that are injected into the LLM request for the current invocation only. The context is deep-copied before insertion, added before the invocation user content, and never appended to `session.events`. The insertion path is separate from instruction-related content so the transient context keeps a stable position across tool-call loops and multi-agent `include_contents="none"` flows. ### Testing Plan **Unit Tests:** - [x] I have added or updated unit tests for my change. - [x] All unit tests pass locally. Passed locally: ```text .venv/bin/python -m pytest tests/unittests/agents/test_llm_agent_include_contents.py tests/unittests/agents/test_run_config.py 14 passed ``` Additional checks: ```text .venv/bin/pyink --check --diff --config pyproject.toml src/google/adk/flows/llm_flows/contents.py tests/unittests/agents/test_llm_agent_include_contents.py tests/unittests/agents/test_run_config.py .venv/bin/isort --check-only src/google/adk/flows/llm_flows/contents.py tests/unittests/agents/test_llm_agent_include_contents.py tests/unittests/agents/test_run_config.py src/google/adk/agents/run_config.py git diff --check ``` **Manual End-to-End (E2E) Tests:** Not run. This change is covered by focused unit tests around LLM request construction and session persistence. ### Checklist - [x] I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) document. - [x] I have performed a self-review of my own code. - [x] I have commented my code, particularly in hard-to-understand areas. - [x] I have added tests that prove my fix is effective or that my feature works. - [x] New and existing unit tests pass locally with my changes. - [x] I have manually tested my changes end-to-end. - [ ] Any dependent changes have been merged and published in downstream modules. ### Additional context The unit coverage verifies that transient context: - is sent to the model without being persisted to the session; - stays before the invocation user message after a tool call; - works for a sub-agent using `include_contents="none"` in a sequential flow; - is accepted by `RunConfig` as `types.Content`. Co-authored-by: George Weale COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/5991 from nicolasmota:spike/model-input-context c8359720bb933d79a85a35f575358c4b497d86a8 PiperOrigin-RevId: 943484488 --- src/google/adk/agents/run_config.py | 8 + src/google/adk/flows/llm_flows/contents.py | 31 +++ .../agents/test_llm_agent_include_contents.py | 191 ++++++++++++++++++ tests/unittests/agents/test_run_config.py | 8 + 4 files changed, 238 insertions(+) diff --git a/src/google/adk/agents/run_config.py b/src/google/adk/agents/run_config.py index 7037a4619e3..8c0570163eb 100644 --- a/src/google/adk/agents/run_config.py +++ b/src/google/adk/agents/run_config.py @@ -371,6 +371,14 @@ class RunConfig(BaseModel): ) """ + model_input_context: list[types.Content] | None = None + """Transient context to include in the model input for this invocation. + + The Runner does not persist these contents to the session. They are only + added to the LLM request assembled for the current invocation, which lets + callers provide per-turn context without changing the conversation history. + """ + @model_validator(mode='before') @classmethod def check_for_deprecated_save_live_audio(cls, data: Any) -> Any: diff --git a/src/google/adk/flows/llm_flows/contents.py b/src/google/adk/flows/llm_flows/contents.py index 8d1395769f4..cd0cc5e64a4 100644 --- a/src/google/adk/flows/llm_flows/contents.py +++ b/src/google/adk/flows/llm_flows/contents.py @@ -14,6 +14,7 @@ from __future__ import annotations +import copy import logging from typing import AsyncGenerator from typing import Optional @@ -105,6 +106,16 @@ async def run_async( user_content=invocation_context.user_content, ) + if ( + invocation_context.run_config + and invocation_context.run_config.model_input_context + ): + _add_model_input_context_to_user_content( + invocation_context, + llm_request, + copy.deepcopy(invocation_context.run_config.model_input_context), + ) + # Add instruction-related contents to proper position in conversation await _add_instructions_to_user_content( invocation_context, llm_request, instruction_related_contents @@ -1039,6 +1050,26 @@ def _content_contains_function_response(content: types.Content) -> bool: return False +def _add_model_input_context_to_user_content( + invocation_context: InvocationContext, + llm_request: LlmRequest, + model_input_context: list[types.Content], +) -> None: + """Insert transient model input context before the invocation user content.""" + if not model_input_context: + return + + insert_index = 0 + user_content = invocation_context.user_content + if user_content: + for i in range(len(llm_request.contents) - 1, -1, -1): + if llm_request.contents[i] == user_content: + insert_index = i + break + + llm_request.contents[insert_index:insert_index] = model_input_context + + async def _add_instructions_to_user_content( invocation_context: InvocationContext, llm_request: LlmRequest, diff --git a/tests/unittests/agents/test_llm_agent_include_contents.py b/tests/unittests/agents/test_llm_agent_include_contents.py index c24aab4ef09..5b62c59011b 100644 --- a/tests/unittests/agents/test_llm_agent_include_contents.py +++ b/tests/unittests/agents/test_llm_agent_include_contents.py @@ -15,6 +15,7 @@ """Unit tests for LlmAgent include_contents field behavior.""" from google.adk.agents.llm_agent import LlmAgent +from google.adk.agents.run_config import RunConfig from google.adk.agents.sequential_agent import SequentialAgent from google.genai import types import pytest @@ -189,6 +190,196 @@ def simple_tool(message: str) -> dict: assert len(mock_model.requests[0].config.tools) > 0 +def test_model_input_context_is_sent_to_model_without_persisting_to_session(): + mock_model = testing_utils.MockModel.create(responses=["Answer"]) + agent = LlmAgent(name="test_agent", model=mock_model) + runner = testing_utils.InMemoryRunner(agent) + session = runner.session + + list( + runner.runner.run( + user_id=session.user_id, + session_id=session.id, + new_message=testing_utils.get_user_content("Question"), + run_config=RunConfig( + model_input_context=[ + types.UserContent("Relevant context for this turn") + ] + ), + ) + ) + + assert testing_utils.simplify_contents(mock_model.requests[0].contents) == [ + ("user", "Relevant context for this turn"), + ("user", "Question"), + ] + assert testing_utils.simplify_events(runner.session.events) == [ + ("user", "Question"), + ("test_agent", "Answer"), + ] + + +def test_model_input_context_stays_before_user_message_after_tool_call(): + def simple_tool(message: str) -> dict: + return {"result": f"Tool processed: {message}"} + + mock_model = testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call( + name="simple_tool", args={"message": "payload"} + ), + "Answer", + ] + ) + agent = LlmAgent(name="test_agent", model=mock_model, tools=[simple_tool]) + runner = testing_utils.InMemoryRunner(agent) + session = runner.session + + list( + runner.runner.run( + user_id=session.user_id, + session_id=session.id, + new_message=testing_utils.get_user_content("Question"), + run_config=RunConfig( + model_input_context=[ + types.UserContent("Relevant context for this turn") + ] + ), + ) + ) + + assert testing_utils.simplify_contents(mock_model.requests[0].contents) == [ + ("user", "Relevant context for this turn"), + ("user", "Question"), + ] + assert testing_utils.simplify_contents(mock_model.requests[1].contents) == [ + ("user", "Relevant context for this turn"), + ("user", "Question"), + ( + "model", + types.Part.from_function_call( + name="simple_tool", args={"message": "payload"} + ), + ), + ( + "user", + types.Part.from_function_response( + name="simple_tool", + response={"result": "Tool processed: payload"}, + ), + ), + ] + assert testing_utils.simplify_events(runner.session.events) == [ + ("user", "Question"), + ( + "test_agent", + types.Part.from_function_call( + name="simple_tool", args={"message": "payload"} + ), + ), + ( + "test_agent", + types.Part.from_function_response( + name="simple_tool", + response={"result": "Tool processed: payload"}, + ), + ), + ("test_agent", "Answer"), + ] + + +def test_model_input_context_with_include_contents_none_sub_agent(): + agent1_model = testing_utils.MockModel.create( + responses=["Agent1 response: XYZ"] + ) + agent1 = LlmAgent(name="agent1", model=agent1_model) + + agent2_model = testing_utils.MockModel.create( + responses=["Agent2 final response"] + ) + agent2 = LlmAgent( + name="agent2", + model=agent2_model, + include_contents="none", + ) + sequential_agent = SequentialAgent( + name="sequential_test_agent", sub_agents=[agent1, agent2] + ) + runner = testing_utils.InMemoryRunner(sequential_agent) + session = runner.session + + list( + runner.runner.run( + user_id=session.user_id, + session_id=session.id, + new_message=testing_utils.get_user_content("Original user request"), + run_config=RunConfig( + model_input_context=[ + types.UserContent("Relevant context for this turn") + ] + ), + ) + ) + + assert testing_utils.simplify_contents(agent1_model.requests[0].contents) == [ + ("user", "Relevant context for this turn"), + ("user", "Original user request"), + ] + assert testing_utils.simplify_contents(agent2_model.requests[0].contents) == [ + ("user", "Relevant context for this turn"), + ( + "user", + [ + types.Part(text="For context:"), + types.Part(text="[agent1] said: Agent1 response: XYZ"), + ], + ), + ] + + +def test_model_input_context_without_user_message_is_prepended_before_history(): + mock_model = testing_utils.MockModel.create( + responses=["First answer", "Second answer"] + ) + agent = LlmAgent(name="test_agent", model=mock_model) + runner = testing_utils.InMemoryRunner(agent) + session = runner.session + + list( + runner.runner.run( + user_id=session.user_id, + session_id=session.id, + new_message=testing_utils.get_user_content("First question"), + ) + ) + # No new_message, so the invocation has no user content to anchor before + # (e.g. live mode or a re-run over existing history). The transient context + # falls back to the front of the request, before all prior history. + list( + runner.runner.run( + user_id=session.user_id, + session_id=session.id, + new_message=None, + run_config=RunConfig( + model_input_context=[ + types.UserContent("Relevant context for this turn") + ] + ), + ) + ) + + assert testing_utils.simplify_contents(mock_model.requests[1].contents) == [ + ("user", "Relevant context for this turn"), + ("user", "First question"), + ("model", "First answer"), + ] + assert testing_utils.simplify_events(runner.session.events) == [ + ("user", "First question"), + ("test_agent", "First answer"), + ("test_agent", "Second answer"), + ] + + @pytest.mark.asyncio async def test_include_contents_none_sequential_agents(): """Test include_contents='none' with sequential agents.""" diff --git a/tests/unittests/agents/test_run_config.py b/tests/unittests/agents/test_run_config.py index a8f9eed0bf1..16eba04835d 100644 --- a/tests/unittests/agents/test_run_config.py +++ b/tests/unittests/agents/test_run_config.py @@ -129,3 +129,11 @@ def test_avatar_config_with_name(): assert run_config.avatar_config == avatar_config assert run_config.avatar_config.avatar_name == "test_avatar" assert run_config.avatar_config.customized_avatar is None + + +def test_model_input_context_accepts_transient_contents(): + context_content = types.UserContent("Relevant context for this turn") + + run_config = RunConfig(model_input_context=[context_content]) + + assert run_config.model_input_context == [context_content] From c14258dffc77804f638f5abbeb434d979ec3149b Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Mon, 6 Jul 2026 14:20:44 -0700 Subject: [PATCH 11/15] feat(bigquery): expose thinking and tool-use token columns in analytics views Add usage_thinking_tokens and usage_tool_use_tokens to the LLM_RESPONSE analytics view, sourced from the usage_metadata proto the plugin already logs to attributes.usage_metadata (thoughts_token_count and tool_use_prompt_token_count). Per the genai GenerateContentResponseUsageMetadata contract, total_token_count = prompt_token_count + candidates_token_count + tool_use_prompt_token_count + thoughts_token_count, so thinking and tool-use tokens are separate addends rather than subsets of prompt/candidates. Surfacing them lets analytics account for them without double counting. Both fields are optional and resolve to NULL for models/responses that do not report them, so the change stays model-agnostic. No plugin logging change is needed since the full usage_metadata is already persisted to attributes. Co-authored-by: Haiyuan Cao PiperOrigin-RevId: 943493796 --- .../bigquery_agent_analytics_plugin.py | 10 ++++++++++ .../test_bigquery_agent_analytics_plugin.py | 20 +++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index dd0d99d5dab..7afd3b3e2d4 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -2060,6 +2060,16 @@ def _parse_custom_metadata_allowlist( " '$.usage_metadata.cached_content_token_count') AS INT64) AS" " usage_cached_tokens" ), + ( + "CAST(JSON_VALUE(attributes," + " '$.usage_metadata.thoughts_token_count') AS INT64) AS" + " usage_thinking_tokens" + ), + ( + "CAST(JSON_VALUE(attributes," + " '$.usage_metadata.tool_use_prompt_token_count') AS INT64) AS" + " usage_tool_use_tokens" + ), ( "SAFE_DIVIDE(CAST(JSON_VALUE(attributes," " '$.usage_metadata.cached_content_token_count') AS" diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index 3dcbce2e9a6..063dcee392c 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -5978,6 +5978,26 @@ def test_view_sql_contains_correct_event_filter(self): view_name = "v_" + event_type.lower() assert view_name in all_sql, f"View {view_name} not found in SQL" + def test_llm_response_view_exposes_token_usage_columns(self): + """LLM_RESPONSE view surfaces cached/thinking/tool-use token columns. + + These are read from the full ``usage_metadata`` proto that is already + logged to ``attributes.usage_metadata``, so they are sourced from + ``attributes`` rather than the ``content.usage`` summary. + """ + plugin = self._make_plugin(create_views=True) + plugin.client.get_table.side_effect = cloud_exceptions.NotFound("not found") + plugin.client.query.return_value = mock.MagicMock() + + plugin._ensure_schema_exists() + + all_sql = " ".join(c[0][0] for c in plugin.client.query.call_args_list) + assert "usage_cached_tokens" in all_sql + assert "usage_thinking_tokens" in all_sql + assert "usage_tool_use_tokens" in all_sql + assert "$.usage_metadata.thoughts_token_count" in all_sql + assert "$.usage_metadata.tool_use_prompt_token_count" in all_sql + def test_config_create_views_default_true(self): """Config create_views defaults to True.""" config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig() From 1263ed64e30805464fff3391554f65ebbf72746b Mon Sep 17 00:00:00 2001 From: Shangjie Chen Date: Mon, 6 Jul 2026 14:32:19 -0700 Subject: [PATCH 12/15] feat: Implement Workflow as Tool core feature Introduce NodeTool, allowing individual Nodes and entire Workflows to be wrapped and executed as standard ADK Tools. This PR implements the core, single-turn execution capability: - Support auto-wrapping of BaseNode (Workflows) directly in Agent.tools. - Wrapping synchronous and asynchronous function nodes as NodeTools. - Providing a complete sample workflow demonstrating how to run a workflow as a tool. Resumption and multi-turn nested HITL support are skipped in this PR and will be fully enabled in the later PR. Co-authored-by: Shangjie Chen PiperOrigin-RevId: 943499058 --- .../samples/workflows/node_as_tool/README.md | 36 + .../samples/workflows/node_as_tool/agent.py | 73 ++ .../workflows/node_as_tool/tests/go.json | 182 +++ src/google/adk/agents/llm_agent.py | 55 +- src/google/adk/tools/_node_tool.py | 156 +++ src/google/adk/workflow/_function_node.py | 25 +- tests/unittests/workflow/test_node_tool.py | 1003 +++++++++++++++++ 7 files changed, 1526 insertions(+), 4 deletions(-) create mode 100644 contributing/samples/workflows/node_as_tool/README.md create mode 100644 contributing/samples/workflows/node_as_tool/agent.py create mode 100644 contributing/samples/workflows/node_as_tool/tests/go.json create mode 100644 src/google/adk/tools/_node_tool.py create mode 100644 tests/unittests/workflow/test_node_tool.py diff --git a/contributing/samples/workflows/node_as_tool/README.md b/contributing/samples/workflows/node_as_tool/README.md new file mode 100644 index 00000000000..605050c2819 --- /dev/null +++ b/contributing/samples/workflows/node_as_tool/README.md @@ -0,0 +1,36 @@ +# Node as Tool + +## Overview + +Demonstrates wrapping both a regular ADK `Node` (using the `@node` decorator) and a `Workflow` as tools that can be automatically called by a parent `Agent`. + +In this sample: + +1. The parent agent receives an inquiry about a customer's discount. +1. It invokes `customer_lookup_workflow` (a `Workflow` wrapped as a tool) to retrieve customer status. +1. It then invokes `calculate_discount` (a regular `Node` wrapped as a tool) using the retrieved status. + +## Sample Inputs + +- `What discount does customer c123 get?` + + *The parent agent first invokes `customer_lookup_workflow` to verify status, then invokes `calculate_discount` to determine the discount percentage, and summarizes the results.* + +## Agent Topology Graph + +```mermaid +graph TD + customer_service_agent[customer_service_agent] -->|calls| customer_lookup_workflow(customer_lookup_workflow) + customer_service_agent -->|calls| calculate_discount(calculate_discount) +``` + +## How To + +To expose an existing `Node` or `Workflow` as a tool callable by an `Agent`: + +1. Define your `Node` (or `@node`) or `Workflow` and assign both an `input_schema` and a `description`. +1. Pass the node/workflow directly into your parent agent's `tools` list: `Agent(..., tools=[my_node, my_workflow])`. + +## Related Guides + +- [Workflows](../../../../docs/guides/workflows/workflows.md) - Explains building complex multi-step graphs. diff --git a/contributing/samples/workflows/node_as_tool/agent.py b/contributing/samples/workflows/node_as_tool/agent.py new file mode 100644 index 00000000000..9b50feca237 --- /dev/null +++ b/contributing/samples/workflows/node_as_tool/agent.py @@ -0,0 +1,73 @@ +# 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 typing import Generator + +from google.adk import Agent +from google.adk import Event +from google.adk import Workflow +from google.adk.workflow import node +from pydantic import BaseModel +from pydantic import Field + + +# 1. Define schemas +class CustomerLookupArgs(BaseModel): + user_id: str = Field(description="The customer's unique identifier.") + + +# 2. Define a regular Node using the @node decorator. +# This Node is wrapped as a NodeTool automatically by the Agent. +# As a NodeTool, it has the ability to yield intermediate Events during execution. +@node +def calculate_discount(tier: str, ctx) -> Generator[Event | str, None, None]: + """Calculates the discount percentage based on customer tier. + + Args: + tier: The customer's membership tier (e.g., VIP, Standard). + """ + yield Event(message=f"Checking discount rules for tier '{tier}'...") + discount = "20% off" if "VIP" in tier else "5% off" + yield discount + + +# 3. Define a Workflow. +# This Workflow is wrapped as a NodeTool automatically by the Agent. +def lookup_customer_data(node_input: CustomerLookupArgs, ctx) -> dict[str, str]: + return {"user_id": node_input.user_id, "tier": "Verified VIP Member"} + + +customer_lookup_workflow = Workflow( + name="customer_lookup_workflow", + description="Looks up customer status and tier by user_id.", + input_schema=CustomerLookupArgs, + edges=[ + ("START", lookup_customer_data), + ], +) + + +# 4. Define the Agent that uses both Node and Workflow as tools. +root_agent = Agent( + name="customer_service_agent", + instruction=""" + You are a customer service assistant. + 1. First, call `customer_lookup_workflow` using the user_id to get their membership tier. + 2. Then, call `calculate_discount` node with that tier to find out what discount they get. + Summarize these details for the customer. + """, + tools=[customer_lookup_workflow, calculate_discount], +) diff --git a/contributing/samples/workflows/node_as_tool/tests/go.json b/contributing/samples/workflows/node_as_tool/tests/go.json new file mode 100644 index 00000000000..10b36485242 --- /dev/null +++ b/contributing/samples/workflows/node_as_tool/tests/go.json @@ -0,0 +1,182 @@ +{ + "appName": "node_as_tool", + "events": [ + { + "author": "user", + "content": { + "parts": [ + { + "text": "What discount does customer c123 get?" + } + ], + "role": "user" + }, + "id": "e-1", + "invocationId": "i-1", + "nodeInfo": { + "path": "" + } + }, + { + "author": "customer_service_agent", + "content": { + "parts": [ + { + "functionCall": { + "args": { + "user_id": "c123" + }, + "id": "fc-1", + "name": "customer_lookup_workflow" + } + } + ], + "role": "model" + }, + "id": "e-2", + "invocationId": "i-1", + "longRunningToolIds": [ + "fc-1" + ], + "nodeInfo": { + "path": "customer_service_agent@1" + } + }, + { + "author": "customer_lookup_workflow", + "branch": "customer_lookup_workflow@fc-1", + "id": "e-3", + "invocationId": "i-1", + "nodeInfo": { + "outputFor": [ + "customer_lookup_workflow@1/lookup_customer_data@1", + "customer_lookup_workflow@1" + ], + "path": "customer_lookup_workflow@1/lookup_customer_data@1" + }, + "output": { + "tier": "Verified VIP Member", + "user_id": "c123" + } + }, + { + "author": "customer_service_agent", + "content": { + "parts": [ + { + "functionResponse": { + "id": "fc-1", + "name": "customer_lookup_workflow", + "response": { + "tier": "Verified VIP Member", + "user_id": "c123" + } + } + } + ], + "role": "user" + }, + "id": "e-4", + "invocationId": "i-1", + "nodeInfo": { + "path": "customer_service_agent@1" + } + }, + { + "author": "customer_service_agent", + "content": { + "parts": [ + { + "functionCall": { + "args": { + "tier": "Verified VIP Member" + }, + "id": "fc-2", + "name": "calculate_discount" + } + } + ], + "role": "model" + }, + "id": "e-5", + "invocationId": "i-1", + "longRunningToolIds": [ + "fc-2" + ], + "nodeInfo": { + "path": "customer_service_agent@1" + } + }, + + { + "author": "calculate_discount", + "branch": "calculate_discount@fc-2", + "content": { + "parts": [ + { + "text": "Checking discount rules for tier 'Verified VIP Member'..." + } + ], + "role": "user" + }, + "id": "e-6", + "invocationId": "i-1", + "nodeInfo": { + "path": "calculate_discount@1" + } + }, + { + "author": "calculate_discount", + "branch": "calculate_discount@fc-2", + "id": "e-7", + "invocationId": "i-1", + "nodeInfo": { + "outputFor": [ + "calculate_discount@1" + ], + "path": "calculate_discount@1" + }, + "output": "20% off" + }, + { + "author": "customer_service_agent", + "content": { + "parts": [ + { + "functionResponse": { + "id": "fc-2", + "name": "calculate_discount", + "response": { + "result": "20% off" + } + } + } + ], + "role": "user" + }, + "id": "e-8", + "invocationId": "i-1", + "nodeInfo": { + "path": "customer_service_agent@1" + } + }, + { + "author": "customer_service_agent", + "content": { + "parts": [ + { + "text": "Customer c123 is a Verified VIP Member and gets a 20% discount." + } + ], + "role": "model" + }, + "id": "e-9", + "invocationId": "i-1", + "nodeInfo": { + "path": "customer_service_agent@1" + } + } + ], + "id": "12345678-1234-1234-1234-123456789abc", + "userId": "user" +} diff --git a/src/google/adk/agents/llm_agent.py b/src/google/adk/agents/llm_agent.py index 6f0c2af79a4..97c9174854a 100644 --- a/src/google/adk/agents/llm_agent.py +++ b/src/google/adk/agents/llm_agent.py @@ -133,7 +133,6 @@ InstructionProvider: TypeAlias = Callable[ [ReadonlyContext], Union[str, Awaitable[str]] ] - ToolUnion: TypeAlias = Union[Callable, BaseTool, BaseToolset] @@ -175,6 +174,32 @@ async def _convert_tool_union_to_tools( max_results=vais_tool.max_results, ) ] + from ..workflow._base_node import BaseNode + + if isinstance(tool_union, BaseNode): + from ..tools._node_tool import NodeTool + from .base_agent import BaseAgent + + if isinstance(tool_union, BaseAgent): + raise ValueError( + f"Agent '{tool_union.name}' cannot be wrapped as a NodeTool. Agents" + ' should be invoked as sub-agents.' + ) + + description = tool_union.description + if not description: + raise ValueError( + f"Workflow/Node '{tool_union.name}' must have a description to be" + ' wrapped as a tool.' + ) + + return [ + NodeTool( + node=tool_union, + name=tool_union.name, + description=description, + ) + ] if isinstance(tool_union, BaseTool): return [tool_union] @@ -1011,6 +1036,34 @@ def __maybe_accumulate_streaming_output( event.actions.state_delta[self.output_key] = accumulator return accumulator + @model_validator(mode='before') + @classmethod + def _pre_validate_tools(cls, data: Any) -> Any: + if isinstance(data, dict) and 'tools' in data and data['tools']: + from google.adk.agents.base_agent import BaseAgent + from google.adk.tools._node_tool import NodeTool + from google.adk.workflow._base_node import BaseNode + + new_tools = [] + for t in data['tools']: + if isinstance(t, BaseAgent): + raise ValueError( + f"Agent '{t.name}' cannot be wrapped as a NodeTool. Agents should" + ' be invoked as sub-agents.' + ) + elif isinstance(t, BaseNode): + description = t.description + if not description: + raise ValueError( + f"Workflow/Node '{t.name}' must have a description to be" + ' wrapped as a tool.' + ) + new_tools.append(NodeTool(node=t, description=description)) + else: + new_tools.append(t) + data['tools'] = new_tools + return data + @model_validator(mode='after') def __model_validator_after(self) -> LlmAgent: return self diff --git a/src/google/adk/tools/_node_tool.py b/src/google/adk/tools/_node_tool.py new file mode 100644 index 00000000000..f69c0ef3feb --- /dev/null +++ b/src/google/adk/tools/_node_tool.py @@ -0,0 +1,156 @@ +# 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 typing import Any + +from google.genai import types +from typing_extensions import override + +from ..utils._schema_utils import schema_to_json_schema +from ..workflow._base_node import BaseNode +from ..workflow._errors import NodeInterruptedError +from .base_tool import BaseTool +from .tool_context import ToolContext + + +class NodeTool(BaseTool): + """A tool wrapper that executes a BaseNode (e.g. a Workflow or loop node).""" + + def __init__( + self, + node: BaseNode, + name: str | None = None, + description: str | None = None, + ): + from ..agents.base_agent import BaseAgent + from ..workflow._function_node import FunctionNode + + if isinstance(node, BaseAgent): + raise ValueError( + f"Agent '{node.name}' cannot be wrapped as a NodeTool. Agents should" + ' be invoked as Sub-Agents instead.' + ) + + # Automatically align FunctionNode binding + if ( + isinstance(node, FunctionNode) + and node.parameter_binding != 'node_input' + ): + orig_input_schema = getattr(node, 'input_schema', None) + orig_output_schema = getattr(node, 'output_schema', None) + node = FunctionNode( + func=node._func, + name=node.name, + rerun_on_resume=node.rerun_on_resume, + retry_config=node.retry_config, + timeout=node.timeout, + auth_config=node.auth_config, + parameter_binding='node_input', # Force binding to node_input + state_schema=node.state_schema, + ) + if orig_input_schema is not None: + node.input_schema = orig_input_schema + if orig_output_schema is not None: + node.output_schema = orig_output_schema + + if not getattr(node, 'input_schema', None): + raise ValueError( + f"Node '{node.name}' does not have an input_schema defined." + ' NodeTool requires an explicit Pydantic input_schema on the wrapped' + ' node.' + ) + + self.node = node + super().__init__( + name=name or node.name, + description=description + or node.description + or f'Executes the node: {node.name}', + ) + self.is_long_running = True + + @override + def _get_declaration(self) -> types.FunctionDeclaration: + schema = schema_to_json_schema(self.node.input_schema) + + # The GenAI API strictly requires parameters_json_schema to be an 'object' + # type schema. If the node has a primitive input schema (e.g., str, int), + # we wrap it into an object schema with a 'request' property. + if isinstance(schema, dict) and schema.get('type') != 'object': + schema = { + 'type': 'object', + 'properties': { + 'request': schema, + }, + 'required': ['request'], + } + + decl = types.FunctionDeclaration( + name=self.name, + description=self.description, + parameters_json_schema=schema, + ) + + output_schema = getattr(self.node, 'output_schema', None) + if output_schema: + decl.response_json_schema = schema_to_json_schema(output_schema) + + return decl + + @override + async def run_async( + self, + *, + args: dict[str, Any], + tool_context: ToolContext, + ) -> Any: + import inspect + + from pydantic import BaseModel + + input_schema = self.node.input_schema + node_input: Any + if inspect.isclass(input_schema) and issubclass(input_schema, BaseModel): + try: + # Convert input based on Pydantic schema + node_input = input_schema.model_validate(args) + except Exception as e: + return f'Error validating input for node: {e}' + else: + schema = schema_to_json_schema(input_schema) + if isinstance(schema, dict) and schema.get('type') != 'object': + node_input = args.get('request') + else: + node_input = args + + fc_id = tool_context.function_call_id + base_branch = tool_context.branch + segment = f'{self.name}@{fc_id}' + tool_branch = f'{base_branch}.{segment}' if base_branch else segment + + try: + return await tool_context.run_node( + self.node, + node_input=node_input, + override_branch=tool_branch, + use_sub_branch=False, + raise_on_wait=True, + ) + except NodeInterruptedError as nie: + # Propagates the interrupt up so the runner pauses the invocation + raise nie + except Exception as e: + return f'Error running node {self.name}: {e}' diff --git a/src/google/adk/workflow/_function_node.py b/src/google/adk/workflow/_function_node.py index c6e03ae1c4b..15670634af2 100644 --- a/src/google/adk/workflow/_function_node.py +++ b/src/google/adk/workflow/_function_node.py @@ -328,9 +328,15 @@ def _bind_parameters(self, ctx: Context, node_input: Any) -> dict[str, Any]: is passed through directly and all other non-context parameters are looked up in ``ctx.state``. """ + from pydantic import BaseModel + input_bound = self.parameter_binding == 'node_input' + source: Any if input_bound: - source = node_input if isinstance(node_input, dict) else {} + if isinstance(node_input, (dict, BaseModel)): + source = node_input + else: + source = {} else: source = ctx.state source_name = 'node_input' if input_bound else 'state' @@ -353,8 +359,21 @@ def _bind_parameters(self, ctx: Context, node_input: Any) -> dict[str, Any]: kwargs[param_name] = value continue - if param_name in source: - value = source[param_name] + has_param = False + value = None + if isinstance(source, BaseModel): + if hasattr(source, param_name): + has_param = True + value = getattr(source, param_name) + else: + try: + if param_name in source: + has_param = True + value = source[param_name] + except (TypeError, KeyError): + pass + + if has_param: if param_name in self._type_hints: value = self._coerce_param( param_name, diff --git a/tests/unittests/workflow/test_node_tool.py b/tests/unittests/workflow/test_node_tool.py new file mode 100644 index 00000000000..826562323d5 --- /dev/null +++ b/tests/unittests/workflow/test_node_tool.py @@ -0,0 +1,1003 @@ +# 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 copy +from typing import Any + +from google.adk.agents.llm_agent import LlmAgent +from google.adk.apps.app import App +from google.adk.apps.app import ResumabilityConfig +from google.adk.events.event import Event +from google.adk.tools._node_tool import NodeTool +from google.adk.tools.long_running_tool import LongRunningFunctionTool +from google.adk.workflow import JoinNode +from google.adk.workflow import node +from google.adk.workflow import START +from google.adk.workflow._node_status import NodeStatus +from google.adk.workflow._workflow import Workflow +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 REQUEST_INPUT_FUNCTION_CALL_NAME +from google.genai import types +from pydantic import BaseModel +from pydantic import Field +import pytest + +from . import workflow_testing_utils +from .. import testing_utils +from .workflow_testing_utils import RequestInputNode + + +class UserInfo(BaseModel): + name: str + age: int + + +class DummyRequest(BaseModel): + request: str = '' + + +def test_node_tool_requires_input_schema(): + """NodeTool raises ValueError if wrapped node has no input_schema.""" + wf = Workflow(name='no_schema_wf', edges=[]) + with pytest.raises(ValueError, match='does not have an input_schema defined'): + NodeTool(node=wf) + + +@pytest.mark.skip(reason='Requires CL 2 subagent branch refactor') +@pytest.mark.asyncio +async def test_workflow_as_tool_hitl_resume(request: pytest.FixtureRequest): + """Workflow-as-a-tool suspends on RequestInput and resumes successfully. + + Setup: + - LlmAgent 'parent_agent' uses WorkflowTool 'collect_user_info_tool'. + - The tool wraps 'sub_workflow' which has a RequestInputNode and a + format_response node. + Act: + - Turn 1: Run with 'Start task'. The model calls the tool, which suspends. + - Turn 2: Resume with the user input response to the interrupt. + Assert: + - Turn 1: Event history contains the RequestInput function call. + - Turn 2: The workflow tool resumes and finishes, and parent agent produces + final text response. + """ + # 1. Define the sub-workflow that has an input interrupt + input_node = RequestInputNode( + name='input_node', + message='What is your name and age?', + response_schema=UserInfo.model_json_schema(), + ) + + def format_response(node_input: dict[str, Any]): + yield Event( + output=f"User {node_input['name']} is {node_input['age']} years old." + ) + + sub_workflow = Workflow( + name='sub_workflow', + edges=[ + (START, input_node), + (input_node, format_response), + ], + ) + sub_workflow.input_schema = DummyRequest + + # 2. Wrap the sub-workflow as a WorkflowTool + wf_tool = NodeTool( + node=sub_workflow, + name='collect_user_info_tool', + description='Call this tool to collect customer name and age.', + ) + + # 3. Define the parent agent that calls this tool + # In the first turn, the model decides to call the tool. + # In the second turn, after the tool resumes and returns output, the model replies to the user. + parent_agent = LlmAgent( + name='parent_agent', + model=testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call( + name='collect_user_info_tool', + args={}, + ), + types.Part.from_text( + text='Thank you! I received the user details.' + ), + ] + ), + tools=[wf_tool], + ) + + # 4. Wrap the parent agent in an App with resumability enabled + app = App( + name=request.function.__name__, + root_agent=parent_agent, + resumability_config=ResumabilityConfig(is_resumable=True), + ) + runner = testing_utils.InMemoryRunner(app=app) + + # Turn 1: Run the agent, triggering the tool call. + # The sub-workflow starts, hits the RequestInputNode, and suspends. + user_event = testing_utils.get_user_content('Start task') + events1 = await runner.run_async(user_event) + + simplified_events1 = ( + workflow_testing_utils.simplify_events_with_node_and_agent_state( + copy.deepcopy(events1), + ) + ) + + # Verify that we got a RequestInput event + request_input_event = workflow_testing_utils.find_function_call_event( + events1, REQUEST_INPUT_FUNCTION_CALL_NAME + ) + assert request_input_event is not None + args = request_input_event.content.parts[0].function_call.args + assert args['message'] == 'What is your name and age?' + + interrupt_id = get_request_input_interrupt_ids(request_input_event)[0] + invocation_id = request_input_event.invocation_id + + # Turn 2: Resume with the user input resolving the interrupt. + user_input = create_request_input_response( + interrupt_id, {'name': 'Alice', 'age': 25} + ) + events2 = await runner.run_async( + new_message=testing_utils.UserContent(user_input), + invocation_id=invocation_id, + ) + + simplified_events2 = ( + workflow_testing_utils.simplify_events_with_node_and_agent_state( + copy.deepcopy(events2), + ) + ) + + # Verify the tool workflow finished executing, returned the output, + # and the parent agent LLM produced its final response. + text_responses = [ + event.content.parts[0].text + for event in events2 + if event.content and event.content.parts and event.content.parts[0].text + ] + assert 'Thank you! I received the user details.' in text_responses + + +@pytest.mark.skip(reason='Requires CL 2 subagent branch refactor') +@pytest.mark.asyncio +async def test_workflow_as_tool_hitl_resume_non_resumable_app( + request: pytest.FixtureRequest, +): + """Workflow-as-a-tool suspends and resumes successfully even when the App has resumability disabled.""" + # 1. Define the sub-workflow that has an input interrupt + input_node = RequestInputNode( + name='input_node', + message='What is your name and age?', + response_schema=UserInfo.model_json_schema(), + ) + + def format_response(node_input: dict[str, Any]): + yield Event( + output=f"User {node_input['name']} is {node_input['age']} years old." + ) + + sub_workflow = Workflow( + name='sub_workflow', + edges=[ + (START, input_node), + (input_node, format_response), + ], + ) + sub_workflow.input_schema = DummyRequest + + # 2. Wrap the sub-workflow as a WorkflowTool + wf_tool = NodeTool( + node=sub_workflow, + name='collect_user_info_tool', + description='Call this tool to collect customer name and age.', + ) + + # 3. Define the parent agent that calls this tool + parent_agent = LlmAgent( + name='parent_agent', + model=testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call( + name='collect_user_info_tool', + args={}, + ), + types.Part.from_text( + text='Thank you! I received the user details.' + ), + ] + ), + tools=[wf_tool], + ) + + # 4. Wrap the parent agent in an App with resumability disabled + app = App( + name=request.function.__name__, + root_agent=parent_agent, + resumability_config=None, + ) + runner = testing_utils.InMemoryRunner(app=app) + + # Turn 1: Run the agent, triggering the tool call. + user_event = testing_utils.get_user_content('Start task') + events1 = await runner.run_async(user_event) + + # Verify that we got a RequestInput event + request_input_event = workflow_testing_utils.find_function_call_event( + events1, REQUEST_INPUT_FUNCTION_CALL_NAME + ) + assert request_input_event is not None + args = request_input_event.content.parts[0].function_call.args + assert args['message'] == 'What is your name and age?' + + interrupt_id = get_request_input_interrupt_ids(request_input_event)[0] + invocation_id = request_input_event.invocation_id + + # Turn 2: Resume with the user input resolving the interrupt. + user_input = create_request_input_response( + interrupt_id, {'name': 'Alice', 'age': 25} + ) + events2 = await runner.run_async( + new_message=testing_utils.UserContent(user_input), + invocation_id=invocation_id, + ) + + # Verify the tool workflow finished executing, returned the output, + # and the parent agent LLM produced its final response. + text_responses = [ + event.content.parts[0].text + for event in events2 + if event.content and event.content.parts and event.content.parts[0].text + ] + assert 'Thank you! I received the user details.' in text_responses + + +def test_node_tool_rejects_agent(): + """NodeTool raises ValueError if initialized with any BaseAgent.""" + agent = LlmAgent( + name='my_agent', + instruction='Answer questions', + ) + with pytest.raises(ValueError, match='cannot be wrapped as a NodeTool'): + NodeTool(node=agent) + + +class GreetRequest(BaseModel): + request: str + + +@pytest.mark.asyncio +async def test_function_node_wrapped_as_tool_returns_output( + request: pytest.FixtureRequest, +): + """NodeTool wraps a function node and returns expected output.""" + + @node + def greet_node(request: str) -> str: + return f'Hello, {request}!' + + greet_node.input_schema = GreetRequest + greet_tool = NodeTool(node=greet_node, name='greet_tool') + + parent_agent = LlmAgent( + name='parent_agent', + model=testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call( + name='greet_tool', + args={'request': 'world'}, + ), + types.Part.from_text(text='Processed greet.'), + ] + ), + tools=[greet_tool], + ) + + app = App( + name=request.function.__name__, + root_agent=parent_agent, + ) + runner = testing_utils.InMemoryRunner(app=app) + events = await runner.run_async(testing_utils.get_user_content('Greet world')) + + func_response_events = [ + e + for e in events + if e.content and e.content.parts and e.content.parts[0].function_response + ] + assert len(func_response_events) == 1 + assert func_response_events[0].content.parts[ + 0 + ].function_response.response == {'result': 'Hello, world!'} + + +@pytest.mark.asyncio +async def test_workflow_tool_with_join_node(request: pytest.FixtureRequest): + """WorkflowTool containing a JoinNode works correctly when wrapped as a tool.""" + node_a = workflow_testing_utils.TestingNode(name='NodeA', output={'a': 1}) + node_b = workflow_testing_utils.TestingNode(name='NodeB', output={'b': 2}) + node_join = JoinNode(name='NodeJoin') + + def format_response(node_input: dict[str, Any]): + yield Event( + output=( + f"A is {node_input['NodeA']['a']} and B is" + f" {node_input['NodeB']['b']}." + ) + ) + + sub_workflow = Workflow( + name='sub_workflow', + edges=[ + (START, node_a), + (START, node_b), + (node_a, node_join), + (node_b, node_join), + (node_join, format_response), + ], + ) + sub_workflow.input_schema = DummyRequest + + wf_tool = NodeTool( + node=sub_workflow, + name='my_join_tool', + description='Collect parallel items.', + ) + + parent_agent = LlmAgent( + name='parent_agent', + model=testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call( + name='my_join_tool', + args={}, + ), + types.Part.from_text(text='Done.'), + ] + ), + tools=[wf_tool], + ) + + app = App( + name=request.function.__name__, + root_agent=parent_agent, + ) + runner = testing_utils.InMemoryRunner(app=app) + events = await runner.run_async(testing_utils.get_user_content('Run join')) + + func_response_events = [ + e + for e in events + if e.content and e.content.parts and e.content.parts[0].function_response + ] + assert len(func_response_events) == 1 + assert func_response_events[0].content.parts[ + 0 + ].function_response.response == {'result': 'A is 1 and B is 2.'} + + +@pytest.mark.asyncio +async def test_workflow_tool_with_dynamic_node(request: pytest.FixtureRequest): + """WorkflowTool containing a dynamic node schedules and executes it correctly.""" + + @node + async def child(*, ctx, node_input): + yield f'child got: {node_input}' + + @node(rerun_on_resume=True) + async def parent_node(*, ctx, node_input): + result = await ctx.run_node(child, node_input='hello') + yield f'parent got: {result}' + + sub_workflow = Workflow( + name='sub_workflow', + edges=[ + (START, parent_node), + ], + ) + sub_workflow.input_schema = DummyRequest + + wf_tool = NodeTool( + node=sub_workflow, + name='my_dynamic_tool', + description='Call dynamic node.', + ) + + parent_agent = LlmAgent( + name='parent_agent', + model=testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call( + name='my_dynamic_tool', + args={}, + ), + types.Part.from_text(text='Done.'), + ] + ), + tools=[wf_tool], + ) + + app = App( + name=request.function.__name__, + root_agent=parent_agent, + ) + runner = testing_utils.InMemoryRunner(app=app) + events = await runner.run_async(testing_utils.get_user_content('Run dynamic')) + + func_response_events = [ + e + for e in events + if e.content and e.content.parts and e.content.parts[0].function_response + ] + assert len(func_response_events) == 1 + assert func_response_events[0].content.parts[ + 0 + ].function_response.response == {'result': 'parent got: child got: hello'} + + +@pytest.mark.asyncio +async def test_workflow_tool_with_nested_workflows( + request: pytest.FixtureRequest, +): + """WorkflowTool wrapping a nested workflow executes successfully.""" + inner_node = workflow_testing_utils.TestingNode( + name='inner_node', output='inner_output' + ) + inner_wf = Workflow( + name='inner_wf', + edges=[ + (START, inner_node), + ], + ) + inner_wf.input_schema = None + + outer_node = workflow_testing_utils.TestingNode( + name='outer_node', output='outer_output' + ) + outer_wf = Workflow( + name='outer_wf', + edges=[ + (START, outer_node, inner_wf), + ], + ) + outer_wf.input_schema = DummyRequest + + wf_tool = NodeTool( + node=outer_wf, + name='nested_wf_tool', + description='Call nested workflow.', + ) + + parent_agent = LlmAgent( + name='parent_agent', + model=testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call( + name='nested_wf_tool', + args={}, + ), + types.Part.from_text(text='Done.'), + ] + ), + tools=[wf_tool], + ) + + app = App( + name=request.function.__name__, + root_agent=parent_agent, + ) + runner = testing_utils.InMemoryRunner(app=app) + events = await runner.run_async(testing_utils.get_user_content('Run nested')) + + func_response_events = [ + e + for e in events + if e.content and e.content.parts and e.content.parts[0].function_response + ] + assert len(func_response_events) == 1 + assert func_response_events[0].content.parts[ + 0 + ].function_response.response == {'result': 'inner_output'} + + +@pytest.mark.skip(reason='Requires CL 2 subagent branch refactor') +@pytest.mark.asyncio +async def test_workflow_tool_with_dynamic_node_hitl_resume( + request: pytest.FixtureRequest, +): + """WorkflowTool with a dynamic node containing HITL suspends and resumes successfully.""" + # 1. Define dynamic node calling a child RequestInputNode + input_node = RequestInputNode( + name='input_node', + message='Enter value:', + response_schema=UserInfo.model_json_schema(), + ) + + @node(rerun_on_resume=True) + async def parent_node(*, ctx, node_input): + result = await ctx.run_node(input_node) + yield f'parent got: {result["name"]}' + + sub_workflow = Workflow( + name='sub_workflow', + edges=[ + (START, parent_node), + ], + ) + sub_workflow.input_schema = DummyRequest + + # 2. Wrap as WorkflowTool + wf_tool = NodeTool( + node=sub_workflow, + name='my_dynamic_hitl_tool', + description='Call dynamic HITL node.', + ) + + # 3. Define parent agent + parent_agent = LlmAgent( + name='parent_agent', + model=testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call( + name='my_dynamic_hitl_tool', + args={}, + ), + types.Part.from_text(text='Task completed.'), + ] + ), + tools=[wf_tool], + ) + + # 4. App with resumability enabled + app = App( + name=request.function.__name__, + root_agent=parent_agent, + resumability_config=ResumabilityConfig(is_resumable=True), + ) + runner = testing_utils.InMemoryRunner(app=app) + + # Turn 1: Run the agent, triggering the tool call and dynamic node suspend. + user_event = testing_utils.get_user_content('Start') + events1 = await runner.run_async(user_event) + + request_input_event = workflow_testing_utils.find_function_call_event( + events1, REQUEST_INPUT_FUNCTION_CALL_NAME + ) + assert request_input_event is not None + interrupt_id = get_request_input_interrupt_ids(request_input_event)[0] + invocation_id = request_input_event.invocation_id + + # Turn 2: Resume with the user input response. + user_input = create_request_input_response( + interrupt_id, {'name': 'Bob', 'age': 30} + ) + events2 = await runner.run_async( + new_message=testing_utils.UserContent(user_input), + invocation_id=invocation_id, + ) + + # Verify the tool workflow finished executing, returned output, + # and parent agent replied. + text_responses = [ + event.content.parts[0].text + for event in events2 + if event.content and event.content.parts and event.content.parts[0].text + ] + assert 'Task completed.' in text_responses + + +@pytest.mark.skip( + reason='Known framework issue with MockModel nested HITL in sub-workflow' +) +@pytest.mark.asyncio +async def test_workflow_as_tool_nested_hitl(request: pytest.FixtureRequest): + """Parent LLM agent -> workflow -> LLM agent -> NodeTool(HITL) propagation.""" + # 1. Define the deepest node that raises RequestInput + input_node = RequestInputNode( + name='deep_input_node', + message='Give me some input:', + response_schema={ + 'type': 'object', + 'properties': {'val': {'type': 'string'}}, + }, + ) + + # 2. Wrap it as a NodeTool + input_node.input_schema = DummyRequest + node_tool = NodeTool(node=input_node, name='my_node_tool') + + # 3. Define the child agent that uses this NodeTool + child_agent = LlmAgent( + name='child_agent', + model=testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call( + name='my_node_tool', + args={}, + ), + types.Part.from_text(text='Child agent processed.'), + ] + ), + tools=[node_tool], + ) + + # 4. Define the sub-workflow containing the child agent + sub_workflow = Workflow( + name='sub_workflow', + edges=[ + (START, child_agent), + ], + ) + sub_workflow.input_schema = DummyRequest + + # 5. Wrap the sub-workflow as a WorkflowTool + wf_tool = NodeTool( + node=sub_workflow, + name='my_wf_tool', + description='Call sub workflow.', + ) + + # 6. Define the parent agent that calls the WorkflowTool + parent_agent = LlmAgent( + name='parent_agent', + model=testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call( + name='my_wf_tool', + args={}, + ), + types.Part.from_text(text='Parent agent finished successfully.'), + ] + ), + tools=[wf_tool], + ) + + # 7. Wrap in App and Runner + app = App( + name=request.function.__name__, + root_agent=parent_agent, + resumability_config=ResumabilityConfig(is_resumable=True), + ) + runner = testing_utils.InMemoryRunner(app=app) + + # Turn 1: Run + events1 = await runner.run_async(testing_utils.get_user_content('Start task')) + + # Assert Turn 1: Expect RequestInput event + request_input_event = workflow_testing_utils.find_function_call_event( + events1, REQUEST_INPUT_FUNCTION_CALL_NAME + ) + assert request_input_event is not None + args = request_input_event.content.parts[0].function_call.args + assert args['message'] == 'Give me some input:' + + interrupt_id = get_request_input_interrupt_ids(request_input_event)[0] + invocation_id = request_input_event.invocation_id + + # Turn 2: Resume + user_input = create_request_input_response(interrupt_id, {'val': 'hello'}) + events2 = await runner.run_async( + new_message=testing_utils.UserContent(user_input), + invocation_id=invocation_id, + ) + + # Assert Turn 2: Expect completion + text_responses = [ + event.content.parts[0].text + for event in events2 + if event.content and event.content.parts and event.content.parts[0].text + ] + assert 'Parent agent finished successfully.' in text_responses + + +@pytest.mark.skip( + reason='Known framework issue with MockModel multi-HITL in nested workflow' +) +@pytest.mark.asyncio +async def test_workflow_as_tool_nested_multi_hitl( + request: pytest.FixtureRequest, +): + """Parent LLM agent -> workflow -> LLM agent -> NodeTool(HITL) twice.""" + # 1. Define the deepest node that raises RequestInput + input_node = RequestInputNode( + name='deep_input_node', + message='Give me some input:', + response_schema={ + 'type': 'object', + 'properties': {'val': {'type': 'string'}}, + }, + ) + + # 2. Wrap it as a NodeTool + input_node.input_schema = DummyRequest + node_tool = NodeTool(node=input_node, name='my_node_tool') + + # 3. Define the child agent that uses this NodeTool twice + child_agent = LlmAgent( + name='child_agent', + model=testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call( + name='my_node_tool', + args={}, + ), + types.Part.from_function_call( + name='my_node_tool', + args={}, + ), + types.Part.from_text(text='Child agent finished.'), + ] + ), + tools=[node_tool], + ) + + # 4. Define the sub-workflow containing the child agent + sub_workflow = Workflow( + name='sub_workflow', + edges=[ + (START, child_agent), + ], + ) + sub_workflow.input_schema = DummyRequest + + # 5. Wrap the sub-workflow as a WorkflowTool + wf_tool = NodeTool( + node=sub_workflow, + name='my_wf_tool', + description='Call sub workflow.', + ) + + # 6. Define the parent agent that calls the WorkflowTool + parent_agent = LlmAgent( + name='parent_agent', + model=testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call( + name='my_wf_tool', + args={}, + ), + types.Part.from_text(text='Parent agent finished successfully.'), + ] + ), + tools=[wf_tool], + ) + + # 7. Wrap in App and Runner + app = App( + name=request.function.__name__, + root_agent=parent_agent, + resumability_config=ResumabilityConfig(is_resumable=True), + ) + runner = testing_utils.InMemoryRunner(app=app) + + # Turn 1: Run -> triggers first HITL + events1 = await runner.run_async(testing_utils.get_user_content('Start task')) + request_input_event1 = workflow_testing_utils.find_function_call_event( + events1, REQUEST_INPUT_FUNCTION_CALL_NAME + ) + assert request_input_event1 is not None + interrupt_id1 = get_request_input_interrupt_ids(request_input_event1)[0] + invocation_id = request_input_event1.invocation_id + + # Turn 2: Resume first HITL -> triggers second HITL + user_input1 = create_request_input_response(interrupt_id1, {'val': 'hello'}) + events2 = await runner.run_async( + new_message=testing_utils.UserContent(user_input1), + invocation_id=invocation_id, + ) + request_input_event2 = workflow_testing_utils.find_function_call_event( + events2, REQUEST_INPUT_FUNCTION_CALL_NAME + ) + assert request_input_event2 is not None + interrupt_id2 = get_request_input_interrupt_ids(request_input_event2)[0] + assert interrupt_id1 != interrupt_id2 + + # Turn 3: Resume second HITL -> finishes + user_input2 = create_request_input_response(interrupt_id2, {'val': 'world'}) + events3 = await runner.run_async( + new_message=testing_utils.UserContent(user_input2), + invocation_id=invocation_id, + ) + + # Assert Turn 3: Expect completion + text_responses = [ + event.content.parts[0].text + for event in events3 + if event.content and event.content.parts and event.content.parts[0].text + ] + assert 'Parent agent finished successfully.' in text_responses + + +@pytest.mark.skip(reason='Requires CL 2 subagent branch refactor') +@pytest.mark.asyncio +async def test_workflow_as_tool_nested_lro(request: pytest.FixtureRequest): + """Parent LLM agent -> workflow -> LLM agent -> LRO tool.""" + + # 1. Define LRO tool function + def my_lro_func(): + return None + + # 2. Define child agent with LRO tool + child_agent = LlmAgent( + name='child_agent', + model=testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call( + name='my_lro_func', + args={}, + ), + types.Part.from_text(text='Child agent finished after LRO.'), + ] + ), + tools=[LongRunningFunctionTool(func=my_lro_func)], + ) + + # 3. Define sub-workflow + sub_workflow = Workflow( + name='sub_workflow', + edges=[ + (START, child_agent), + ], + ) + sub_workflow.input_schema = DummyRequest + + # 4. Wrap as WorkflowTool + wf_tool = NodeTool( + node=sub_workflow, + name='my_wf_tool', + description='Call sub workflow.', + ) + + # 5. Define parent agent + parent_agent = LlmAgent( + name='parent_agent', + model=testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call( + name='my_wf_tool', + args={}, + ), + types.Part.from_text(text='Parent agent finished successfully.'), + ] + ), + tools=[wf_tool], + ) + + # 6. Wrap in App and Runner + app = App( + name=request.function.__name__, + root_agent=parent_agent, + resumability_config=ResumabilityConfig(is_resumable=True), + ) + runner = testing_utils.InMemoryRunner(app=app) + + # Turn 1: Run -> should pause on LRO + events1 = await runner.run_async(testing_utils.get_user_content('Start task')) + 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, 'my_lro_func' + ) + assert fc_event is not None + function_call_id = fc_event.content.parts[0].function_call.id + + # Turn 2: Resume with LRO response + tool_response = testing_utils.UserContent( + types.Part( + function_response=types.FunctionResponse( + id=function_call_id, + name='my_lro_func', + response={'result': 'LRO finished'}, + ) + ) + ) + events2 = await runner.run_async( + new_message=tool_response, + invocation_id=invocation_id, + ) + + # Assert Turn 2: Expect completion + text_responses = [ + event.content.parts[0].text + for event in events2 + if event.content and event.content.parts and event.content.parts[0].text + ] + assert 'Parent agent finished successfully.' in text_responses + + +def test_node_tool_auto_converts_function_node_binding(): + """NodeTool automatically converts FunctionNode parameter_binding to 'node_input'.""" + + @node + def my_func_node(request: str) -> str: + """A dummy node.""" + return f'Result: {request}' + + # Originally it is 'state' mode by default + assert my_func_node.parameter_binding == 'state' + # input_schema is originally None + assert getattr(my_func_node, 'input_schema', None) is None + + # Wrap it + tool = NodeTool(node=my_func_node) + + # Check that the wrapped node copy is converted to 'node_input' mode + assert tool.node.parameter_binding == 'node_input' + # And input_schema is automatically inferred + schema = tool.node.input_schema + assert 'request' in schema['properties'] + + +@pytest.mark.asyncio +async def test_node_tool_primitive_input_schema(request: pytest.FixtureRequest): + """NodeTool automatically wraps primitive input_schema to object in declaration and unwraps in run.""" + + def echo_func(node_input: str): + yield Event(output=f'Echo: {node_input}') + + sub_workflow = Workflow( + name='sub_workflow', + edges=[ + (START, echo_func), + ], + ) + sub_workflow.input_schema = str + tool = NodeTool(node=sub_workflow, name='primitive_tool') + + # 1. Check declaration is wrapped to object schema + decl = tool._get_declaration() + assert decl.parameters_json_schema is not None + assert decl.parameters_json_schema['type'] == 'object' + assert 'request' in decl.parameters_json_schema['properties'] + assert ( + decl.parameters_json_schema['properties']['request']['type'] == 'string' + ) + + # 2. Run the tool (passing wrapped argument) and check execution + parent_agent = LlmAgent( + name='parent_agent', + model=testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call( + name='primitive_tool', + args={'request': 'hello_world'}, + ), + types.Part.from_text(text='Finished.'), + ] + ), + tools=[tool], + ) + app = App( + name=request.function.__name__, + root_agent=parent_agent, + ) + runner = testing_utils.InMemoryRunner(app=app) + events = await runner.run_async(testing_utils.get_user_content('Run')) + + func_response_events = [ + e + for e in events + if e.content and e.content.parts and e.content.parts[0].function_response + ] + assert len(func_response_events) == 1 + assert func_response_events[0].content.parts[ + 0 + ].function_response.response == {'result': 'Echo: hello_world'} From 4aa0fd8df4d93f6b2ab8c79373cde08f2cf3a9f7 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Mon, 6 Jul 2026 14:35:33 -0700 Subject: [PATCH 13/15] fix: prevent warning logs when cancelling MCP session runner task PiperOrigin-RevId: 943500722 --- src/google/adk/tools/mcp_tool/session_context.py | 4 ++-- .../tools/mcp_tool/test_session_context.py | 14 +++++++++++--- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/src/google/adk/tools/mcp_tool/session_context.py b/src/google/adk/tools/mcp_tool/session_context.py index 0c6bc9369bc..82da207a197 100644 --- a/src/google/adk/tools/mcp_tool/session_context.py +++ b/src/google/adk/tools/mcp_tool/session_context.py @@ -347,8 +347,8 @@ async def _run(self): # Wait for close signal - the session remains valid while we wait await self._close_event.wait() - except BaseException as e: - logger.warning(f'Error on session runner task: {e}') + except Exception as e: + logger.warning('Error on session runner task: %s', e) raise finally: self._ready_event.set() diff --git a/tests/unittests/tools/mcp_tool/test_session_context.py b/tests/unittests/tools/mcp_tool/test_session_context.py index 71b2879d513..9634a4013a2 100644 --- a/tests/unittests/tools/mcp_tool/test_session_context.py +++ b/tests/unittests/tools/mcp_tool/test_session_context.py @@ -609,9 +609,14 @@ async def test_close_handles_cancelled_error(self): mock_session = MockClientSession() - with patch( - 'google.adk.tools.mcp_tool.session_context.ClientSession' - ) as mock_session_class: + with ( + patch( + 'google.adk.tools.mcp_tool.session_context.ClientSession' + ) as mock_session_class, + patch( + 'google.adk.tools.mcp_tool.session_context.logger' + ) as mock_logger, + ): mock_session_class.return_value = mock_session await session_context.start() @@ -626,6 +631,9 @@ async def test_close_handles_cancelled_error(self): # Should not raise exception assert session_context._close_event.is_set() + # Verify no warning logs were generated + mock_logger.warning.assert_not_called() + @pytest.mark.asyncio async def test_close_handles_exception_during_cleanup(self): """Test that close() handles exceptions during cleanup gracefully.""" From 3ebef82a52f2b650d1ee1f94dbdeeb31fbd0ae10 Mon Sep 17 00:00:00 2001 From: Kathy Wu Date: Mon, 6 Jul 2026 14:39:40 -0700 Subject: [PATCH 14/15] fix: Truncate MCP http debug logs if greater than 1000 chars To prevent OOM issues if HTTP request/response bodies are too long Co-authored-by: Kathy Wu PiperOrigin-RevId: 943502772 --- .../adk/tools/mcp_tool/mcp_session_manager.py | 8 ++++ .../mcp_tool/test_mcp_session_manager.py | 43 +++++++++++++++++++ 2 files changed, 51 insertions(+) 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 8407ee7ab80..3a61929e76d 100644 --- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py +++ b/src/google/adk/tools/mcp_tool/mcp_session_manager.py @@ -81,6 +81,8 @@ class AsyncAuthorizedSession: # pylint: disable=g-bad-classes logger = logging.getLogger('google_adk.' + __name__) +_MAX_LOG_BODY_LENGTH = 1000 + def create_mcp_http_client( headers: dict[str, str] | None = None, @@ -271,6 +273,8 @@ async def _response_hook(self, response: httpx.Response): request_body = response.request.content.decode( 'utf-8', errors='replace' ) + if len(request_body) > _MAX_LOG_BODY_LENGTH: + request_body = request_body[:_MAX_LOG_BODY_LENGTH] + '... [truncated]' except Exception: # pylint: disable=broad-exception-caught request_body = '' @@ -278,6 +282,10 @@ async def _response_hook(self, response: httpx.Response): try: await response.aread() response_body = response.text + if len(response_body) > _MAX_LOG_BODY_LENGTH: + response_body = ( + response_body[:_MAX_LOG_BODY_LENGTH] + '... [truncated]' + ) except Exception as e: # pylint: disable=broad-exception-caught response_body = f'' else: 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 b769eca55e1..2f6a11305d5 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py @@ -1471,3 +1471,46 @@ def keyword_only_factory(**kwargs) -> httpx.AsyncClient: client = debug_factory({"X-Test": "Val"}, None, None) assert client is base_client await base_client.aclose() + + @pytest.mark.asyncio + async def test_response_hook_truncates_large_bodies(self): + """Test that response hook truncates request and response bodies exceeding limit.""" + base_client = httpx.AsyncClient() + base_factory = Mock(return_value=base_client) + debug_factory = _DebugHttpxClientFactory(base_factory) + + # Mock request and response with large content + large_req_body = b"a" * 1500 + large_resp_body = "b" * 1500 + + mock_request = Mock(spec=httpx.Request) + mock_request.method = "POST" + mock_request.content = large_req_body + mock_request.headers = httpx.Headers() + + mock_response = Mock(spec=httpx.Response) + mock_response.url = httpx.URL("https://example.com/large") + mock_response.status_code = 200 + mock_response.request = mock_request + mock_response.headers = httpx.Headers({"content-type": "application/json"}) + mock_response.text = large_resp_body + mock_response.aread = AsyncMock() + + debug_list = [] + token = _http_debug_var.set(debug_list) + try: + await debug_factory._response_hook(mock_response) + finally: + _http_debug_var.reset(token) + + assert len(debug_list) == 1 + record = debug_list[0] + assert len(record["request_body"]) == 1015 # 1000 + len("... [truncated]") + assert record["request_body"].endswith("... [truncated]") + assert record["request_body"].startswith("a" * 1000) + + assert len(record["response_body"]) == 1015 # 1000 + len("... [truncated]") + assert record["response_body"].endswith("... [truncated]") + assert record["response_body"].startswith("b" * 1000) + + await base_client.aclose() From 63561ce7192faa77cc0abb45b29f07f512cfbb9c Mon Sep 17 00:00:00 2001 From: Haran Rajkumar Date: Mon, 6 Jul 2026 15:03:55 -0700 Subject: [PATCH 15/15] feat(scripts): add check for private-by-default new Python files Add scripts/check_new_py_files.py, which enforces ADK's private-by-default policy: a newly-added Python file under src/google/adk/ must have a '_'-prefixed basename (expose public API via __init__.py / __all__ instead). Newly-added files are detected by comparing the checked-out tree against a baseline source tree, so the check needs no git history. Co-authored-by: Haran Rajkumar PiperOrigin-RevId: 943513793 --- scripts/check_new_py_files.py | 137 ++++++++++++++++++++++++++++++++++ 1 file changed, 137 insertions(+) create mode 100644 scripts/check_new_py_files.py diff --git a/scripts/check_new_py_files.py b/scripts/check_new_py_files.py new file mode 100644 index 00000000000..ffcc9731f54 --- /dev/null +++ b/scripts/check_new_py_files.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python3 +# 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. + +"""Checks that newly-added Python files under src/google/adk/ have a '_' prefix. + +ADK is private-by-default: a newly-added Python file under src/google/adk/ must +have a '_'-prefixed basename. To make it public, add the symbol to the package +__init__.py / __all__ instead. See +.agents/skills/adk-style/references/visibility.md. + +Newly-added files are detected by diffing the working tree against a baseline +source tree (e.g. an origin/main checkout), so it works in a checkout that has +no local git history: + + python scripts/check_new_py_files.py --baseline-dir /path/to/origin-main + +Exit codes: 0 = ok, 1 = violation(s) found, 2 = usage/setup error. +""" + +from __future__ import annotations + +import argparse +import os +import sys + +_PACKAGE_RELPATH = os.path.join('src', 'google', 'adk') + +_VIOLATION_LINE = "Error: New Python file '{path}' must have a '_' prefix." + +_GUIDANCE = ( + 'All new Python files in src/google/adk/ must be private by default.\n' + 'To expose a public interface, use __init__.py and list public symbols' + ' in __all__.\n' + 'See .agents/skills/adk-style/references/visibility.md for details.' +) + +# Subtrees that may exist in the working tree but are intentionally absent from +# the baseline tree; ignore them so the diff does not report them as newly +# added. +_IGNORED_PREFIXES = ( + 'src/google/adk/internal/', + 'src/google/adk/v1/', + 'src/google/adk/platform/internal/', +) + + +def find_py_files(root: str) -> set[str]: + """Returns root-relative paths of every *.py under /src/google/adk. + + Each path includes the src/google/adk/ prefix (e.g. + 'src/google/adk/agents/foo.py'). Symlinks are followed so that a src/google/adk + tree assembled from symlinked subdirectories is walked correctly. + """ + package_root = os.path.join(root, _PACKAGE_RELPATH) + found: set[str] = set() + for dirpath, _, filenames in os.walk(package_root, followlinks=True): + for name in filenames: + if name.endswith('.py'): + abs_path = os.path.join(dirpath, name) + found.add(os.path.relpath(abs_path, root)) + return found + + +def _should_check(relpath: str) -> bool: + """Returns False for paths under an ignored prefix.""" + return not any(relpath.startswith(prefix) for prefix in _IGNORED_PREFIXES) + + +def added_py_files(new_root: str, baseline_root: str) -> set[str]: + """Returns .py files present in new_root but not in baseline_root. + + Paths under _IGNORED_PREFIXES are skipped: they may exist in the working tree + but are intentionally absent from the baseline, so a plain diff would + otherwise report them as newly added. + """ + added = find_py_files(new_root) - find_py_files(baseline_root) + return {path for path in added if _should_check(path)} + + +def find_violations(added: set[str]) -> list[str]: + """Returns the sorted added files whose basename does not start with '_'.""" + return sorted( + path for path in added if not os.path.basename(path).startswith('_') + ) + + +def _has_package_dir(root: str) -> bool: + return os.path.isdir(os.path.join(root, _PACKAGE_RELPATH)) + + +def _parse_args(argv: list[str]) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + '--baseline-dir', + required=True, + help='Baseline source tree to diff against (an origin/main checkout).', + ) + parser.add_argument( + '--new-dir', + default='.', + help='New source tree to check (default: current directory).', + ) + return parser.parse_args(argv) + + +def main(argv: list[str]) -> int: + args = _parse_args(argv) + for label, root in (('baseline', args.baseline_dir), ('new', args.new_dir)): + if not _has_package_dir(root): + print( + f'Error: {label} tree has no {_PACKAGE_RELPATH} directory: {root}', + file=sys.stderr, + ) + return 2 + + violations = find_violations(added_py_files(args.new_dir, args.baseline_dir)) + for path in violations: + print(_VIOLATION_LINE.format(path=path), file=sys.stderr) + if violations: + print(_GUIDANCE, file=sys.stderr) + return 1 if violations else 0 + + +if __name__ == '__main__': + sys.exit(main(sys.argv[1:]))