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, 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/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:])) 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/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/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/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/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/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/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/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/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/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/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/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/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/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/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', 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/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/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/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/_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/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/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/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_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/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_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 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] 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( 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@") 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.""" 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 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() 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() 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( 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.""" 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'} 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'