diff --git a/contributing/samples/environment_and_skills/daytona_environment/README.md b/contributing/samples/environment_and_skills/daytona_environment/README.md new file mode 100644 index 00000000000..dc344f1645b --- /dev/null +++ b/contributing/samples/environment_and_skills/daytona_environment/README.md @@ -0,0 +1,87 @@ +# Daytona Environment Sample + +## Overview + +A small data analysis agent that uses the `DaytonaEnvironment` with the +`EnvironmentToolset` to download public datasets and analyze them inside a +[Daytona](https://daytona.io) remote sandbox. + +Instead of running on the local machine, all commands and file operations +execute in an isolated remote sandbox with internet access. Asked a question, +the agent downloads a public dataset (a GCS-hosted world population / +demographics dataset by default), installs `pandas` on demand, writes a short +analysis script, runs it, and reports the result — all without touching the +user's machine. This makes the sandbox a natural fit for running +model-generated code safely and keeping the host clean. + +## Prerequisites + +1. Install the `daytona` extra: + + ```bash + pip install google-adk[daytona] + ``` + +1. Set your Daytona configuration. Get a server and API key by following the + Daytona installation guide (e.g. self-hosted or via Daytona Cloud). + + If you are using Daytona Cloud, you only need to set: + + ```bash + export DAYTONA_API_KEY="your-api-key" + ``` + + If you are using a self-hosted Daytona server, also set: + + ```bash + export DAYTONA_API_URL="your-api-url" + ``` + +## Sample Inputs + +- `Download the world demographics dataset and tell me which country has the largest population.` + + The agent downloads the dataset, installs `pandas`, filters to country-level + rows, and finds the maximum. Expected: China (`CN`), ≈ 1.44 billion, just + ahead of India (`IN`) at ≈ 1.38 billion. + +- `For the United States, what is the urban vs rural population split?` + + A follow-up to the previous turn. Because the sandbox persists across the + session, the agent reuses the already-downloaded CSV and the installed + `pandas` — it only writes and runs a new script. Expected for `US`: urban + ≈ 270.7 million vs rural ≈ 57.6 million (out of ≈ 331 million total). + +- `Using https://storage.googleapis.com/cloud-samples-data/bigquery/us-states/us-states.csv, how many US states are listed?` + + Demonstrates pointing the agent at your own dataset URL instead of the + default. + +## Graph + +```mermaid +graph TD + User -->|question| Agent[data_analysis_agent] + Agent -->|EnvironmentToolset| Sandbox[DaytonaEnvironment sandbox] + Sandbox -->|download / install / run| Agent + Agent -->|answer| User +``` + +## How To + +The agent is a standalone `Agent` (no workflow graph) wired to a single +`EnvironmentToolset` whose `environment` is a `DaytonaEnvironment`: + +```python +from google.adk.integrations.daytona import DaytonaEnvironment +from google.adk.tools.environment import EnvironmentToolset + +EnvironmentToolset( + environment=DaytonaEnvironment(timeout=300), +) +``` + +- `timeout` bounds the sandbox lifetime in seconds. +- By default, it will spin up a sandbox from the built-in default Python snapshot. + If you want to use a custom Docker image instead, you can pass it to the + `image` parameter (e.g. `image="python:3.12"`). diff --git a/contributing/samples/environment_and_skills/daytona_environment/__init__.py b/contributing/samples/environment_and_skills/daytona_environment/__init__.py new file mode 100644 index 00000000000..4015e47d6e4 --- /dev/null +++ b/contributing/samples/environment_and_skills/daytona_environment/__init__.py @@ -0,0 +1,15 @@ +# 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 . import agent diff --git a/contributing/samples/environment_and_skills/daytona_environment/agent.py b/contributing/samples/environment_and_skills/daytona_environment/agent.py new file mode 100644 index 00000000000..71c01a1adbe --- /dev/null +++ b/contributing/samples/environment_and_skills/daytona_environment/agent.py @@ -0,0 +1,43 @@ +# 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. + +"""A data analysis agent that runs Python in a Daytona remote sandbox.""" + +from google.adk import Agent +from google.adk.integrations.daytona import DaytonaEnvironment +from google.adk.tools.environment import EnvironmentToolset + +root_agent = Agent( + name="data_analysis_agent", + description=( + "A data analysis agent that downloads public datasets and analyzes" + " them inside a Daytona remote sandbox." + ), + instruction="""\ +You are a data analysis assistant. You work inside an isolated Daytona remote +sandbox that has internet access, where you can safely download data and run +Python, so you never touch the user's machine. + +To analyze a dataset: +1. Download it from the internet into the working directory, e.g. with + `curl -O ` or `wget `. +2. Install whatever you need on demand, e.g. `pip install pandas`. +3. Write a short Python script that loads the data and computes the answer. +4. Run the script and report the result, showing the numbers you found. + +Prefer writing a script and executing it over guessing. If a command fails, +read the error, fix the script, and try again. +""", + tools=[EnvironmentToolset(environment=DaytonaEnvironment())], +) diff --git a/pyproject.toml b/pyproject.toml index 0fd1049383a..0972f353d6e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,6 +65,7 @@ optional-dependencies.agent-identity = [ ] optional-dependencies.all = [ "anyio>=4.9,<5", + "daytona>=0.191", "e2b>=2,<3", "google-api-python-client>=2.157,<3", "google-cloud-aiplatform[agent-engines]>=1.148.1,<2", @@ -101,6 +102,9 @@ optional-dependencies.benchmark = [ optional-dependencies.community = [ "google-adk-community", ] +optional-dependencies.daytona = [ + "daytona>=0.191", # For DaytonaEnvironment remote sandbox. +] optional-dependencies.db = [ "sqlalchemy>=2,<3", "sqlalchemy-spanner>=1.14", @@ -199,6 +203,7 @@ optional-dependencies.test = [ "anyio>=4.9,<5", "beautifulsoup4>=3.2.2", "crewai[tools]; python_version>='3.11' and python_version<'3.12'", # For CrewaiTool tests; chromadb/pypika fail on 3.12+ + "daytona>=0.191", "docker>=7", # For ContainerCodeExecutor tests "e2b>=2,<3", "gepa>=0.1", diff --git a/src/google/adk/auth/auth_preprocessor.py b/src/google/adk/auth/auth_preprocessor.py index b0fa1e0ba82..20ed53ac4c9 100644 --- a/src/google/adk/auth/auth_preprocessor.py +++ b/src/google/adk/auth/auth_preprocessor.py @@ -134,7 +134,7 @@ async def run_async( agent = invocation_context.agent if not hasattr(agent, 'canonical_tools'): return - events = invocation_context.session.events + events = invocation_context._get_events(current_branch=True) if not events: return diff --git a/src/google/adk/cli/cli_tools_click.py b/src/google/adk/cli/cli_tools_click.py index f884bce61a9..29484a3bcad 100644 --- a/src/google/adk/cli/cli_tools_click.py +++ b/src/google/adk/cli/cli_tools_click.py @@ -1806,8 +1806,8 @@ def cli_web( """Starts a FastAPI server with Web UI for agents. AGENTS_DIR: The directory of agents (where each subdirectory is a single - agent containing `agent.py` or `root_agent.yaml` files) or a path pointing - directly to a single agent folder. + agent containing `agent.py`, `__init__.py`, or `root_agent.yaml`) or a path + pointing directly to a single agent folder. Example: @@ -1947,8 +1947,8 @@ def cli_api_server( """Starts a FastAPI server for agents. AGENTS_DIR: The directory of agents (where each subdirectory is a single - agent containing `agent.py` or `root_agent.yaml` files) or a path pointing - directly to a single agent folder. + agent containing `agent.py`, `__init__.py`, or `root_agent.yaml`) or a path + pointing directly to a single agent folder. Example: diff --git a/src/google/adk/cli/utils/agent_loader.py b/src/google/adk/cli/utils/agent_loader.py index dd5e0ed81a6..1a865fd104a 100644 --- a/src/google/adk/cli/utils/agent_loader.py +++ b/src/google/adk/cli/utils/agent_loader.py @@ -471,6 +471,8 @@ def _determine_agent_language( return "python" elif (base_path / "__init__.py").exists(): return "python" + elif (base_path.parent / f"{agent_name}.py").exists(): + return "python" raise ValueError(f"Could not determine agent type for '{agent_name}'.") diff --git a/src/google/adk/features/_feature_registry.py b/src/google/adk/features/_feature_registry.py index b2c1623052e..8e473e4a6fe 100644 --- a/src/google/adk/features/_feature_registry.py +++ b/src/google/adk/features/_feature_registry.py @@ -37,6 +37,8 @@ class FeatureName(str, Enum): COMPUTER_USE = "COMPUTER_USE" DATA_AGENT_TOOL_CONFIG = "DATA_AGENT_TOOL_CONFIG" DATA_AGENT_TOOLSET = "DATA_AGENT_TOOLSET" + DAYTONA_ENVIRONMENT = "DAYTONA_ENVIRONMENT" + E2B_ENVIRONMENT = "E2B_ENVIRONMENT" ENVIRONMENT_SIMULATION = "ENVIRONMENT_SIMULATION" GCS_ADMIN_TOOLSET = "GCS_ADMIN_TOOLSET" GCS_TOOL_SETTINGS = "GCS_TOOL_SETTINGS" @@ -128,6 +130,12 @@ class FeatureConfig: FeatureName.DATA_AGENT_TOOLSET: FeatureConfig( FeatureStage.EXPERIMENTAL, default_on=True ), + FeatureName.DAYTONA_ENVIRONMENT: FeatureConfig( + FeatureStage.EXPERIMENTAL, default_on=True + ), + FeatureName.E2B_ENVIRONMENT: FeatureConfig( + FeatureStage.EXPERIMENTAL, default_on=True + ), FeatureName.ENVIRONMENT_SIMULATION: FeatureConfig( FeatureStage.EXPERIMENTAL, default_on=True ), diff --git a/src/google/adk/integrations/daytona/__init__.py b/src/google/adk/integrations/daytona/__init__.py new file mode 100644 index 00000000000..38db3b66a68 --- /dev/null +++ b/src/google/adk/integrations/daytona/__init__.py @@ -0,0 +1,38 @@ +# 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. + +"""Daytona sandbox integration. + +This module provides a BaseEnvironment implementation backed by a Daytona +remote sandbox, offering a persistent remote workspace for file CRUD and +shell execution. + +Requires the ``daytona`` extra: ``pip install google-adk[daytona]``. + +Example: + ```python + from google.adk.integrations.daytona import DaytonaEnvironment + + env = DaytonaEnvironment() + await env.initialize() + result = await env.execute("pip install requests") + await env.close() + ``` +""" + +from ._daytona_environment import DaytonaEnvironment + +__all__ = [ + "DaytonaEnvironment", +] diff --git a/src/google/adk/integrations/daytona/_daytona_environment.py b/src/google/adk/integrations/daytona/_daytona_environment.py new file mode 100644 index 00000000000..6c8990c86f3 --- /dev/null +++ b/src/google/adk/integrations/daytona/_daytona_environment.py @@ -0,0 +1,243 @@ +# 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. + +"""Daytona sandbox code execution environment.""" + +from __future__ import annotations + +import logging +import os +from pathlib import Path +from pathlib import PurePosixPath +from typing import TYPE_CHECKING + +from typing_extensions import override + +from ...environment._base_environment import BaseEnvironment +from ...environment._base_environment import ExecutionResult +from ...features import experimental +from ...features import FeatureName + +if TYPE_CHECKING: + from daytona import AsyncDaytona + from daytona import AsyncSandbox + from daytona import ExecuteResponse + from daytona import Image + +logger = logging.getLogger("google_adk." + __name__) + +_DEFAULT_TIMEOUT = 300 +_SANDBOX_HOME = "/workspaces" + + +@experimental(FeatureName.DAYTONA_ENVIRONMENT) +class DaytonaEnvironment(BaseEnvironment): + """A persistent remote workspace backed by a Daytona sandbox. + + Provides file CRUD and shell execution inside an isolated remote sandbox. + One sandbox is created on ``initialize()`` and killed on ``close()``. + + Requires the ``daytona`` extra: ``pip install google-adk[daytona]``. + """ + + def __init__( + self, + *, + image: str | Image | None = None, + timeout: int = _DEFAULT_TIMEOUT, + api_key: str | None = None, + api_url: str | None = None, + env_vars: dict[str, str] | None = None, + ): + """Create a Daytona environment. + + Args: + image: Daytona template/image name used to create the sandbox. + timeout: Sandbox time-to-live / timeout in seconds. + api_key: Daytona API key. If ``None``, the environment variable is used. + api_url: Daytona API URL. If ``None``, defaults to Daytona Cloud API. + env_vars: Environment variables set inside the sandbox. + """ + self._image = image + self._timeout = timeout + self._api_key = api_key + self._api_url = api_url + self._env_vars = env_vars + self._sandbox: AsyncSandbox | None = None + self._client: AsyncDaytona | None = None # To hold AsyncDaytona instance + + @property + @override + def working_dir(self) -> Path: + if self._sandbox is None: + raise RuntimeError("Sandbox is not started. Call initialize() first.") + return Path(_SANDBOX_HOME) + + @override + async def initialize(self) -> None: + if self._sandbox is not None: + return + self._sandbox = await self._create_sandbox() + self._is_initialized = True + + @override + async def close(self) -> None: + if self._sandbox is not None: + await self._sandbox.delete() + self._sandbox = None + self._client = None + self._is_initialized = False + + @override + async def execute( + self, + command: str, + *, + timeout: float | None = None, + ) -> ExecutionResult: + sandbox = await self._ensure_sandbox() + + # timeout needs to be int for daytona SDK + timeout_int = int(timeout) if timeout is not None else self._timeout + + try: + response: ExecuteResponse = await sandbox.process.exec( + command=command, + timeout=timeout_int, + ) + except Exception as e: + # If Daytona has specific Timeout exceptions, they should be handled, + # but a generic fallback is catching and logging/translating. + from daytona import DaytonaError + + if isinstance(e, DaytonaError) and "timeout" in str(e).lower(): + return ExecutionResult(exit_code=-1, timed_out=True) + # Otherwise raise + raise e + + return ExecutionResult( + exit_code=response.exit_code or 0, + stdout=response.artifacts.stdout if response.artifacts else "", + # Daytona process.exec combines stdout and stderr into stdout. + stderr="", + ) + + @override + async def read_file(self, path: str | os.PathLike[str]) -> bytes: + sandbox = await self._ensure_sandbox() + resolved = self._resolve_path(path) + try: + content = await sandbox.fs.download_file(resolved) + if content is None: + raise FileNotFoundError(resolved) + return bytes(content) + except Exception as e: + from daytona import DaytonaNotFoundError + + if isinstance(e, DaytonaNotFoundError): + raise FileNotFoundError(resolved) from e + raise e + + @override + async def write_file( + self, path: str | os.PathLike[str], content: str | bytes + ) -> None: + sandbox = await self._ensure_sandbox() + resolved = self._resolve_path(path) + + # Create parent directory recursively to prevent upload failures + resolved_path = PurePosixPath(resolved) + parent = resolved_path.parent + if parent and parent != PurePosixPath("/"): + parts = parent.parts + for i in range(2, len(parts) + 1): + current = PurePosixPath(*parts[:i]) + try: + await sandbox.fs.create_folder(str(current), mode="755") + except Exception as e: + # If folder already exists, Daytona may raise DaytonaConflictError + # or similar. We check the message or type if we can, but safely + # ignoring is fine since the ultimate upload will fail if it's a + # real issue. + from daytona import DaytonaConflictError + + if ( + isinstance(e, DaytonaConflictError) + or "already exists" in str(e).lower() + ): + continue + + # Daytona's upload_file accepts bytes directly + if isinstance(content, str): + content_bytes = content.encode("utf-8") + else: + content_bytes = content + + await sandbox.fs.upload_file(content_bytes, resolved) + + async def _create_sandbox(self) -> AsyncSandbox: + try: + from daytona import AsyncDaytona + from daytona import CreateSandboxFromImageParams + from daytona import CreateSandboxFromSnapshotParams + from daytona import DaytonaConfig + except ImportError as e: + raise ImportError( + "The daytona package is required to use DaytonaEnvironment. Install" + " it with `pip install google-adk[daytona]`." + ) from e + + config_args = {} + if self._api_key: + config_args["api_key"] = self._api_key + if self._api_url: + config_args["api_url"] = self._api_url + + config = DaytonaConfig(**config_args) if config_args else None + self._client = AsyncDaytona(config=config) + + auto_stop_interval_mins = self._timeout // 60 + if self._timeout > 0 and auto_stop_interval_mins == 0: + auto_stop_interval_mins = 1 + + if self._image: + params = CreateSandboxFromImageParams( + image=self._image, + env_vars=self._env_vars or {}, + auto_stop_interval=auto_stop_interval_mins, + auto_delete_interval=0, + ) + else: + params = CreateSandboxFromSnapshotParams( + language="python", + env_vars=self._env_vars or {}, + auto_stop_interval=auto_stop_interval_mins, + auto_delete_interval=0, + ) + + return await self._client.create(params) + + async def _ensure_sandbox(self) -> AsyncSandbox: + sandbox = self._sandbox + if sandbox is None: + raise RuntimeError("Sandbox is not started. Call initialize() first.") + await sandbox.refresh_activity() + return sandbox + + def _resolve_path(self, path: str | os.PathLike[str]) -> str: + """Resolve a relative path against the sandbox working directory.""" + pure = PurePosixPath(os.fspath(path)) + if pure.is_absolute(): + return str(pure) + return str(PurePosixPath(_SANDBOX_HOME) / pure) diff --git a/src/google/adk/integrations/e2b/_e2b_environment.py b/src/google/adk/integrations/e2b/_e2b_environment.py index 55140d7d354..246a27e6ed8 100644 --- a/src/google/adk/integrations/e2b/_e2b_environment.py +++ b/src/google/adk/integrations/e2b/_e2b_environment.py @@ -27,19 +27,20 @@ from ...environment._base_environment import BaseEnvironment from ...environment._base_environment import ExecutionResult -from ...utils.feature_decorator import experimental +from ...features import experimental +from ...features import FeatureName if TYPE_CHECKING: from e2b import AsyncSandbox -logger = logging.getLogger('google_adk.' + __name__) +logger = logging.getLogger("google_adk." + __name__) -_DEFAULT_IMAGE = 'base' +_DEFAULT_IMAGE = "base" _DEFAULT_TIMEOUT = 300 -_SANDBOX_HOME = '/home/user' +_SANDBOX_HOME = "/home/user" -@experimental +@experimental(FeatureName.E2B_ENVIRONMENT) class E2BEnvironment(BaseEnvironment): """A persistent remote workspace backed by an E2B sandbox. @@ -86,7 +87,7 @@ def __init__( @override def working_dir(self) -> Path: if self._sandbox is None: - raise RuntimeError('Sandbox is not started. Call initialize() first.') + raise RuntimeError("Sandbox is not started. Call initialize() first.") return Path(_SANDBOX_HOME) @override @@ -94,12 +95,14 @@ async def initialize(self) -> None: if self._sandbox is not None: return self._sandbox = await self._create_sandbox() + self._is_initialized = True @override async def close(self) -> None: if self._sandbox is not None: await self._sandbox.kill() self._sandbox = None + self._is_initialized = False @override async def execute( @@ -137,7 +140,7 @@ async def read_file(self, path: str | os.PathLike[str]) -> bytes: sandbox = await self._ensure_sandbox() resolved = self._resolve_path(path) try: - content = await sandbox.files.read(resolved, format='bytes') + content = await sandbox.files.read(resolved, format="bytes") except FileNotFoundException as e: raise FileNotFoundError(resolved) from e return bytes(content) @@ -155,8 +158,8 @@ async def _create_sandbox(self) -> AsyncSandbox: from e2b import AsyncSandbox except ImportError as e: raise ImportError( - 'The e2b package is required to use E2BEnvironment. Install it with' - ' `pip install google-adk[e2b]`.' + "The e2b package is required to use E2BEnvironment. Install it with" + " `pip install google-adk[e2b]`." ) from e return await AsyncSandbox.create( @@ -168,15 +171,15 @@ async def _create_sandbox(self) -> AsyncSandbox: async def _ensure_sandbox(self) -> AsyncSandbox: if self._sandbox is None: - raise RuntimeError('Sandbox is not started. Call initialize() first.') + raise RuntimeError("Sandbox is not started. Call initialize() first.") if await self._sandbox.is_running(): # Keepalive: extend the TTL while the workspace is actively used. await self._sandbox.set_timeout(self._timeout) else: logger.warning( - 'E2B sandbox expired; recreating a fresh sandbox. Workspace state' - ' (installed packages and files) has been lost.' + "E2B sandbox expired; recreating a fresh sandbox. Workspace state" + " (installed packages and files) has been lost." ) self._sandbox = await self._create_sandbox() return self._sandbox diff --git a/src/google/adk/integrations/parameter_manager/parameter_client.py b/src/google/adk/integrations/parameter_manager/parameter_client.py index 812202d0ef1..7ff77255071 100644 --- a/src/google/adk/integrations/parameter_manager/parameter_client.py +++ b/src/google/adk/integrations/parameter_manager/parameter_client.py @@ -25,7 +25,7 @@ from google.oauth2 import service_account from ... import version -from ...utils import _mtls_utils +from ...utils._mtls_utils import get_api_endpoint USER_AGENT = f"google-adk/{version.__version__}" @@ -112,7 +112,7 @@ def __init__( client_options = None if location: client_options = { - "api_endpoint": _mtls_utils.get_api_endpoint( + "api_endpoint": get_api_endpoint( location, _DEFAULT_REGIONAL_ENDPOINT_TEMPLATE, _DEFAULT_MTLS_REGIONAL_ENDPOINT_TEMPLATE, diff --git a/src/google/adk/telemetry/_instrumentation.py b/src/google/adk/telemetry/_instrumentation.py index 1f8cc31feba..8bda08edd82 100644 --- a/src/google/adk/telemetry/_instrumentation.py +++ b/src/google/adk/telemetry/_instrumentation.py @@ -29,6 +29,8 @@ from . import _metrics from . import tracing from ..events import event as event_lib +from ._schema_version import resolve_schema_version +from ._schema_version import SCHEMA_VERSION_SEMCONV_ALIGNED if TYPE_CHECKING: from ..agents.base_agent import BaseAgent @@ -41,52 +43,42 @@ logger = logging.getLogger("google_adk." + __name__) -def _get_elapsed_s( - span: trace.Span | tracing.GenerateContentSpan | None, - fallback_start: float, -) -> float: - """Guarantees consistent time source for duration calculation. - - Note: This must be called with an ended span. - - Args: - span (trace.Span | tracing.GenerateContentSpan | None): The ended span to - extract duration from. - fallback_start (float): Fallback start time in seconds (monotonic). - - Returns: - float: Elapsed duration in seconds. - """ - if span is None: - return time.monotonic() - fallback_start - - span = span.span if hasattr(span, "span") else span - start_ns = getattr(span, "start_time", None) - end_ns = getattr(span, "end_time", None) - - if isinstance(start_ns, int) and isinstance(end_ns, int): - return (end_ns - start_ns) / 1e9 # Convert ns to s - - # Fallback if span times are missing - return time.monotonic() - fallback_start - - @contextlib.contextmanager def record_invocation( entrypoint_node: BaseNode | None, conversation_id: str, ) -> Iterator[None]: - """Top-level ``invocation`` span for a runner invocation. + """Top-level invocation span for a runner invocation. + + Schema v1 emits the legacy ``invocation`` span. Schema v2 replaces it with an + entrypoint ``invoke_workflow {entrypoint}`` span (entrypoint = root agent or + root node name), which omits the ``gen_ai.workflow.nested`` attribute, and a + ``gen_ai.invoke_workflow.duration`` metric -- unless the entrypoint is itself + a workflow, in which case its own node span is the entrypoint + ``invoke_workflow`` span and we avoid double-emitting it here. Args: entrypoint_node: The runner's root agent/node. - conversation_id: Session/conversation id. + conversation_id: Session/conversation id (stamped on the v2 span). Yields: - Nothing; the span is active for the duration of the block. + Nothing; the span (if any) is active for the duration of the block. """ - del entrypoint_node, conversation_id # Unused until schema v2 lands. - with tracing.tracer.start_as_current_span("invocation"): + if resolve_schema_version() < SCHEMA_VERSION_SEMCONV_ALIGNED: + with tracing.tracer.start_as_current_span("invocation"): + yield + return + + from . import node_tracing + from ..workflow._workflow import Workflow + + if isinstance(entrypoint_node, Workflow): + # The workflow's own node span is the entrypoint `invoke_workflow` span. + yield + return + + entrypoint_name = entrypoint_node.name if entrypoint_node else "" + with node_tracing._use_invoke_workflow_span(entrypoint_name, conversation_id): yield @@ -152,7 +144,7 @@ async def record_agent_invocation( finally: _record_agent_metrics( agent.name, - _get_elapsed_s(span, start_time), + _metrics.get_elapsed_s(span, start_time), getattr(ctx, "user_content", None), getattr(getattr(ctx, "session", None), "events", []), caught_error, @@ -198,7 +190,7 @@ async def record_tool_execution( tool_name=tool.name, tool_type=tool.__class__.__name__, agent_name=agent.name, - elapsed_s=_get_elapsed_s(span, start_time), + elapsed_s=_metrics.get_elapsed_s(span, start_time), error=caught_error, ) except Exception: # pylint: disable=broad-exception-caught @@ -227,7 +219,7 @@ async def record_inference_telemetry( finally: inference_error = sys.exc_info()[1] agent = invocation_context.agent - elapsed_s = _get_elapsed_s(tel_ctx.span, start_time) + elapsed_s = _metrics.get_elapsed_s(tel_ctx.span, start_time) try: if agent is not None and tracing._should_emit_native_telemetry(agent): _metrics.record_client_operation_duration( diff --git a/src/google/adk/telemetry/_metrics.py b/src/google/adk/telemetry/_metrics.py index eac65b48ca0..04ec0e9c777 100644 --- a/src/google/adk/telemetry/_metrics.py +++ b/src/google/adk/telemetry/_metrics.py @@ -15,6 +15,7 @@ from __future__ import annotations import logging +import time from typing import TYPE_CHECKING from google.adk import version @@ -30,6 +31,10 @@ from google.adk.events.event import Event from google.adk.models.llm_request import LlmRequest from google.adk.models.llm_response import LlmResponse + from opentelemetry.trace import Span + from opentelemetry.util.types import AttributeValue + + from .tracing import GenerateContentSpan logger = logging.getLogger("google_adk." + __name__) @@ -42,7 +47,7 @@ ) _agent_invocation_duration = meter.create_histogram( - "gen_ai.agent.invocation.duration", + "gen_ai.invoke_agent.duration", unit="s", description="Duration of agent invocations.", explicit_bucket_boundaries_advisory=[ @@ -61,8 +66,13 @@ 409.6, ], ) +_workflow_invocation_duration = meter.create_histogram( + "gen_ai.invoke_workflow.duration", + unit="s", + description="Duration of workflow invocations.", +) _tool_execution_duration = meter.create_histogram( - "gen_ai.tool.execution.duration", + "gen_ai.execute_tool.duration", unit="s", description="Duration of tool executions.", explicit_bucket_boundaries_advisory=[ @@ -160,6 +170,27 @@ def record_agent_invocation_duration( _agent_invocation_duration.record(elapsed_s, attributes=attrs) +def record_workflow_invocation_duration( + *, + workflow_name: str, + elapsed_s: float, + nested: bool, + error: BaseException | None = None, +) -> None: + """Records the duration of a workflow invocation.""" + attrs: dict[str, AttributeValue] = { + gen_ai_attributes.GEN_AI_OPERATION_NAME: "invoke_workflow", + } + # Root workflow omits the attribute entirely; only nested ones emit it. + if nested: + attrs["gen_ai.workflow.nested"] = True + if error is not None: + attrs[error_attributes.ERROR_TYPE] = type(error).__name__ + if workflow_name: + attrs["gen_ai.workflow.name"] = workflow_name + _workflow_invocation_duration.record(elapsed_s, attributes=attrs) + + def record_agent_request_size( agent_name: str, user_content: types.Content | None ): @@ -303,3 +334,33 @@ def _get_content_size( def _get_provider_name() -> str: return tracing._guess_gemini_system_name() + + +def get_elapsed_s( + span: Span | GenerateContentSpan | None, + fallback_start: float, +) -> float: + """Guarantees consistent time source for duration calculation. + + Note: This must be called with an ended span. + + Args: + span (trace.Span | tracing.GenerateContentSpan | None): The ended span to + extract duration from. + fallback_start (float): Fallback start time in seconds (monotonic). + + Returns: + float: Elapsed duration in seconds. + """ + if span is None: + return time.monotonic() - fallback_start + + span = span.span if hasattr(span, "span") else span + start_ns = getattr(span, "start_time", None) + end_ns = getattr(span, "end_time", None) + + if isinstance(start_ns, int) and isinstance(end_ns, int): + return (end_ns - start_ns) / 1e9 # Convert ns to s + + # Fallback if span times are missing + return time.monotonic() - fallback_start diff --git a/src/google/adk/telemetry/node_tracing.py b/src/google/adk/telemetry/node_tracing.py index 500cb894594..9fc4a75f411 100644 --- a/src/google/adk/telemetry/node_tracing.py +++ b/src/google/adk/telemetry/node_tracing.py @@ -15,24 +15,46 @@ from __future__ import annotations from collections.abc import AsyncIterator +from collections.abc import Iterator from contextlib import asynccontextmanager +from contextlib import contextmanager from dataclasses import dataclass from dataclasses import field +import sys +import time from typing import TYPE_CHECKING from opentelemetry import context as context_api from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_CONVERSATION_ID from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_OPERATION_NAME -from opentelemetry.util.types import Attributes +from opentelemetry.trace import Span +from . import _metrics from ..agents.context import Context from ..workflow._base_node import BaseNode from .tracing import tracer if TYPE_CHECKING: + from opentelemetry.util.types import AttributeValue + + from ..agents.base_agent import BaseAgent from ..events.event import Event from ..workflow._workflow import Workflow +# Span/metric attribute flagging that an `invoke_workflow` span is nested +# within another workflow. Only emitted for nested workflows; the root +# (entrypoint) workflow omits it entirely. +GEN_AI_WORKFLOW_NESTED = "gen_ai.workflow.nested" + +# OTel-context key recording that an entrypoint workflow is already active. It +# rides along the otel_context propagated to child nodes, so only the first +# workflow invoked within an invocation is treated as the root -- nested +# workflows (incl. agents-as-tool that spin up their own runner) see the key +# already set and report nested=true. +_ENTRYPOINT_WORKFLOW_KEY = context_api.create_key( + "adk-entrypoint-workflow-active" +) + @dataclass(frozen=True) class TelemetryContext: @@ -49,12 +71,6 @@ def add_event(self, event: Event) -> None: self._associated_event_ids.append(event.id) -@dataclass -class _SpanMetadata: - name: str - attributes: Attributes - - @asynccontextmanager async def start_as_current_node_span( context: Context, node: BaseNode @@ -83,64 +99,123 @@ async def start_as_current_node_span( Context with the started span. """ - span_metadata = _span_metadata(context, node) - if span_metadata is None: - token = context_api.attach(context.telemetry_context.otel_context) - try: - yield TelemetryContext( - otel_context=context.telemetry_context.otel_context - ) - finally: - context_api.detach(token) - return - - with tracer.start_as_current_span( - span_metadata.name, - attributes=span_metadata.attributes, - context=context.telemetry_context.otel_context, - ) as span: - telemetry_context = TelemetryContext(otel_context=context_api.get_current()) - yield telemetry_context - - if span.is_recording() and len(telemetry_context._associated_event_ids) > 0: - span.set_attribute( - "gcp.vertex.agent.associated_event_ids", - telemetry_context._associated_event_ids, - ) - - -def _span_metadata(context: Context, node: BaseNode) -> _SpanMetadata | None: from ..agents.base_agent import BaseAgent from ..workflow._workflow import Workflow if isinstance(node, BaseAgent): - return None + with _invoke_agent_span(context, node) as tel_ctx: + yield tel_ctx elif isinstance(node, Workflow): - return _workflow_span_metadata(context, node) + with _invoke_workflow_span(context, node) as tel_ctx: + yield tel_ctx else: - return _default_node_span_metadata(context, node) + with _invoke_node_span(context, node) as tel_ctx: + yield tel_ctx + + +@contextmanager +def _invoke_agent_span( + context: Context, agent: BaseAgent +) -> Iterator[TelemetryContext]: + """Passes through an agent node; agents emit their own `invoke_agent` span.""" + del agent + token = context_api.attach(context.telemetry_context.otel_context) + try: + yield TelemetryContext(otel_context=context.telemetry_context.otel_context) + finally: + context_api.detach(token) -def _workflow_span_metadata( +@contextmanager +def _invoke_workflow_span( context: Context, workflow: Workflow -) -> _SpanMetadata: - return _SpanMetadata( - name=f"invoke_workflow {workflow.name}", - attributes={ - GEN_AI_OPERATION_NAME: "invoke_workflow", - "gen_ai.workflow.name": workflow.name, - GEN_AI_CONVERSATION_ID: context.session.id, - }, - ) +) -> Iterator[TelemetryContext]: + """Opens an `invoke_workflow` span plus its duration metric for ``node``.""" + with _use_invoke_workflow_span( + workflow.name, + context.session.id, + otel_context=context.telemetry_context.otel_context, + ) as span: + tel_ctx = TelemetryContext(otel_context=context_api.get_current()) + yield tel_ctx + _maybe_set_associated_events(span, tel_ctx) -def _default_node_span_metadata( +@contextmanager +def _invoke_node_span( context: Context, node: BaseNode -) -> _SpanMetadata: - return _SpanMetadata( - name=f"invoke_node {node.name}", +) -> Iterator[TelemetryContext]: + """Opens an `invoke_node` span for a plain node.""" + with tracer.start_as_current_span( + f"invoke_node {node.name}", attributes={ GEN_AI_OPERATION_NAME: "invoke_node", GEN_AI_CONVERSATION_ID: context.session.id, }, + context=context.telemetry_context.otel_context, + ) as span: + tel_ctx = TelemetryContext(otel_context=context_api.get_current()) + yield tel_ctx + _maybe_set_associated_events(span, tel_ctx) + + +def _maybe_set_associated_events( + span: Span, telemetry_context: TelemetryContext +) -> None: + """Stamps the node's associated event IDs onto its span, if any.""" + if span.is_recording() and len(telemetry_context._associated_event_ids) > 0: + span.set_attribute( + "gcp.vertex.agent.associated_event_ids", + telemetry_context._associated_event_ids, + ) + + +@contextmanager +def _use_invoke_workflow_span( + workflow_name: str, + conversation_id: str, + *, + otel_context: context_api.Context | None = None, +) -> Iterator[Span]: + """Opens an `invoke_workflow {workflow_name}` span.""" + if otel_context is None: + otel_context = context_api.get_current() + # First workflow in the invocation is the root; subsequent ones are nested. + # The flag rides along the otel_context propagated to child nodes, so nested + # workflows see it set. + nested = bool(context_api.get_value(_ENTRYPOINT_WORKFLOW_KEY, otel_context)) + if not nested: + otel_context = context_api.set_value( + _ENTRYPOINT_WORKFLOW_KEY, True, otel_context + ) + attributes: dict[str, AttributeValue] = { + GEN_AI_OPERATION_NAME: "invoke_workflow", + GEN_AI_CONVERSATION_ID: conversation_id, + } + # Root workflow omits the attribute entirely; only nested ones emit it. + if nested: + attributes[GEN_AI_WORKFLOW_NESTED] = True + if workflow_name: + attributes["gen_ai.workflow.name"] = workflow_name + + span_name = ( + f"invoke_workflow {workflow_name}" if workflow_name else "invoke_workflow" ) + + start_s = time.monotonic() + workflow_span: Span | None = None + try: + with tracer.start_as_current_span( + name=span_name, + attributes=attributes, + context=otel_context, + ) as span: + workflow_span = span + yield span + finally: + _metrics.record_workflow_invocation_duration( + workflow_name=workflow_name, + elapsed_s=_metrics.get_elapsed_s(workflow_span, start_s), + nested=nested, + error=sys.exc_info()[1], + ) diff --git a/src/google/adk/tools/application_integration_tool/application_integration_toolset.py b/src/google/adk/tools/application_integration_tool/application_integration_toolset.py index 030a73aa8c5..93fd055d97d 100644 --- a/src/google/adk/tools/application_integration_tool/application_integration_toolset.py +++ b/src/google/adk/tools/application_integration_tool/application_integration_toolset.py @@ -103,6 +103,7 @@ def __init__( auth_scheme: Optional[AuthScheme] = None, auth_credential: Optional[AuthCredential] = None, tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, + credential_key: Optional[str] = None, ): """Args: @@ -146,12 +147,14 @@ def __init__( self._service_account_json = service_account_json self._auth_scheme = auth_scheme self._auth_credential = auth_credential + self._credential_key = credential_key # Store auth config as instance variable so ADK can populate # exchanged_auth_credential in-place before calling get_tools() self._auth_config: Optional[AuthConfig] = ( AuthConfig( auth_scheme=auth_scheme, raw_auth_credential=auth_credential, + credential_key=credential_key, ) if auth_scheme else None @@ -221,6 +224,7 @@ def _parse_spec_to_toolset( spec_dict=spec_dict, auth_credential=auth_credential, auth_scheme=auth_scheme, + credential_key=self._credential_key, tool_filter=self.tool_filter, ) return @@ -274,6 +278,7 @@ def _parse_spec_to_toolset( rest_api_tool=rest_api_tool, auth_scheme=connector_auth_scheme, auth_credential=connector_auth_credential, + credential_key=self._credential_key, ) ) diff --git a/src/google/adk/tools/application_integration_tool/integration_connector_tool.py b/src/google/adk/tools/application_integration_tool/integration_connector_tool.py index 822965eb21e..fdafa746ad4 100644 --- a/src/google/adk/tools/application_integration_tool/integration_connector_tool.py +++ b/src/google/adk/tools/application_integration_tool/integration_connector_tool.py @@ -80,6 +80,7 @@ def __init__( rest_api_tool: RestApiTool, auth_scheme: Optional[Union[AuthScheme, str]] = None, auth_credential: Optional[Union[AuthCredential, str]] = None, + credential_key: Optional[str] = None, ): """Initializes the ApplicationIntegrationTool. @@ -115,6 +116,7 @@ def __init__( self._rest_api_tool = rest_api_tool self._auth_scheme = auth_scheme self._auth_credential = auth_credential + self._credential_key = credential_key @override def _get_declaration(self) -> FunctionDeclaration: @@ -156,7 +158,10 @@ async def run_async( ) -> Dict[str, Any]: tool_auth_handler = ToolAuthHandler.from_tool_context( - tool_context, self._auth_scheme, self._auth_credential + tool_context, + self._auth_scheme, + self._auth_credential, + credential_key=self._credential_key, ) auth_result = await tool_auth_handler.prepare_auth_credentials() 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 39d5e719423..66e6dad88e4 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 @@ -14,11 +14,13 @@ from __future__ import annotations +import logging from typing import Dict from typing import List from typing import Optional from typing import Union +import httpx from typing_extensions import override from ...agents.readonly_context import ReadonlyContext @@ -26,10 +28,14 @@ from ...auth.auth_schemes import OpenIdConnectWithConfig from ...tools.base_toolset import BaseToolset from ...tools.base_toolset import ToolPredicate +from ...utils._mtls_utils import MtlsClientCerts +from ...utils._mtls_utils import use_client_cert_effective from ..openapi_tool import OpenAPIToolset from .google_api_tool import GoogleApiTool from .googleapi_to_openapi_converter import GoogleApiToOpenApiConverter +logger = logging.getLogger('google_adk.' + __name__) + class GoogleApiToolset(BaseToolset): """Google API Toolset contains tools for interacting with Google APIs. @@ -75,6 +81,20 @@ def __init__( self._additional_headers = additional_headers self._additional_scopes = additional_scopes self._discovery_url = discovery_url + + self._httpx_client_factory = None + use_client_cert = use_client_cert_effective() + + if use_client_cert: + self._mtls_certs = MtlsClientCerts() + cert_path, key_path, passphrase = self._mtls_certs.get_certs() + if cert_path and key_path and passphrase: + + def client_factory(): + return httpx.AsyncClient(cert=(cert_path, key_path, passphrase)) + + self._httpx_client_factory = client_factory + self._openapi_toolset = self._load_toolset_with_oidc_auth() @override @@ -134,6 +154,7 @@ def _load_toolset_with_oidc_auth(self) -> OpenAPIToolset: grant_types_supported=['authorization_code'], scopes=scopes, ), + httpx_client_factory=self._httpx_client_factory, ) def configure_auth(self, client_id: str, client_secret: str): @@ -147,3 +168,5 @@ def configure_sa_auth(self, service_account: ServiceAccount): async def close(self): if self._openapi_toolset: await self._openapi_toolset.close() + if hasattr(self, '_mtls_certs') and self._mtls_certs: + self._mtls_certs.close() diff --git a/src/google/adk/tools/google_api_tool/googleapi_to_openapi_converter.py b/src/google/adk/tools/google_api_tool/googleapi_to_openapi_converter.py index 9c2aaa5f6dc..b6eab32be88 100644 --- a/src/google/adk/tools/google_api_tool/googleapi_to_openapi_converter.py +++ b/src/google/adk/tools/google_api_tool/googleapi_to_openapi_converter.py @@ -17,6 +17,7 @@ import argparse import json import logging +import socket from typing import Any from typing import Dict from typing import List @@ -24,6 +25,10 @@ # Google API client from googleapiclient.discovery import build from googleapiclient.errors import HttpError +import httplib2 + +from ...utils._mtls_utils import MtlsClientCerts +from ...utils._mtls_utils import use_client_cert_effective # Configure logging logger = logging.getLogger("google_adk." + __name__) @@ -54,6 +59,8 @@ def __init__( "paths": {}, "components": {"schemas": {}, "securitySchemes": {}}, } + self._use_client_cert = use_client_cert_effective() + self._mtls_certs = MtlsClientCerts() if self._use_client_cert else None def fetch_google_api_spec(self) -> None: """Fetches the Google API specification using discovery service.""" @@ -63,15 +70,42 @@ def fetch_google_api_spec(self) -> None: self._api_name, self._api_version, ) + + # Determine if we should use mTLS + # self._use_client_cert is already initialized in __init__ + + http_client = None + discovery_url = self._discovery_url + + if self._use_client_cert and self._mtls_certs: + cert_path, key_path, passphrase = self._mtls_certs.get_certs() + if cert_path and key_path and passphrase: + # Set default HTTP timeout similar to googleapiclient.http.build_http() + http_timeout = socket.getdefaulttimeout() or 60 + http_client = httplib2.Http(timeout=http_timeout) + try: + http_client.redirect_codes = http_client.redirect_codes - {308} + except AttributeError: + pass + http_client.add_certificate(key_path, cert_path, "", passphrase) + + if not discovery_url: + discovery_url = "https://www.mtls.googleapis.com/discovery/v1/apis/{api}/{apiVersion}/rest" + # Build a resource object for the specified API - if self._discovery_url: + if discovery_url: self._google_api_resource = build( self._api_name, self._api_version, - discoveryServiceUrl=self._discovery_url, + discoveryServiceUrl=discovery_url, + http=http_client, ) else: - self._google_api_resource = build(self._api_name, self._api_version) + self._google_api_resource = build( + self._api_name, + self._api_version, + http=http_client, + ) # Access the underlying API discovery document self._google_api_spec = self._google_api_resource._rootDesc @@ -136,9 +170,13 @@ def _convert_info(self) -> None: def _convert_servers(self) -> None: """Convert server information.""" - base_url = self._google_api_spec.get( - "rootUrl", "" - ) + self._google_api_spec.get("servicePath", "") + use_client_cert = getattr(self, "_use_client_cert", False) + if use_client_cert and "mtlsRootUrl" in self._google_api_spec: + root_url = self._google_api_spec["mtlsRootUrl"] + else: + root_url = self._google_api_spec.get("rootUrl", "") + + base_url = root_url + self._google_api_spec.get("servicePath", "") # Remove trailing slash if present if base_url.endswith("/"): diff --git a/src/google/adk/utils/_mtls_utils.py b/src/google/adk/utils/_mtls_utils.py index 58848627cd6..c36ed82740e 100644 --- a/src/google/adk/utils/_mtls_utils.py +++ b/src/google/adk/utils/_mtls_utils.py @@ -19,6 +19,8 @@ import enum import logging import os +import tempfile +import threading from typing import TYPE_CHECKING from urllib.parse import urlsplit from urllib.parse import urlunsplit @@ -149,3 +151,69 @@ def configure_session_for_mtls(session: requests.Session) -> bool: if is_mtls: session.mount("https://", _MutualTlsAdapter(cert, key)) return bool(is_mtls) + + +class MtlsClientCerts: + """Manages the creation and lifecycle of client certificates for mTLS. + + Extracts certificates to a temporary directory that is automatically cleaned up + when the instance is garbage collected. + """ + + def __init__(self) -> None: + self._tempdir: tempfile.TemporaryDirectory[str] | None = None + self.cert_path: str | None = None + self.key_path: str | None = None + self.passphrase: bytes | None = None + self._lock = threading.Lock() + self._initialized = False + + def get_certs(self) -> tuple[str | None, str | None, bytes | None]: + """Extracts and returns the certificate paths and passphrase. + + Returns: + A tuple of (cert_path, key_path, passphrase) if client certificates + are available, otherwise (None, None, None). + """ + with self._lock: + if self._initialized: + return self.cert_path, self.key_path, self.passphrase + + if not mtls.has_default_client_cert_source(): + self._initialized = True + return None, None, None + + self._tempdir = tempfile.TemporaryDirectory() + cert_path_tmp = os.path.join(self._tempdir.name, "cert.pem") + key_path_tmp = os.path.join(self._tempdir.name, "key.pem") + + try: + cert_source = mtls.default_client_encrypted_cert_source( + cert_path_tmp, key_path_tmp + ) + _, _, passphrase = cert_source() + except Exception as e: + # If extraction fails, we should fail loud. + self._tempdir.cleanup() + self._tempdir = None + raise RuntimeError( + f"Failed to extract default client certificates for mTLS: {e}" + ) from e + + self.cert_path = cert_path_tmp + self.key_path = key_path_tmp + self.passphrase = passphrase + self._initialized = True + + return self.cert_path, self.key_path, self.passphrase + + def close(self) -> None: + """Manually cleans up the temporary directory.""" + with self._lock: + if self._tempdir: + self._tempdir.cleanup() + self._tempdir = None + self.cert_path = None + self.key_path = None + self.passphrase = None + self._initialized = False diff --git a/tests/unittests/auth/test_auth_preprocessor.py b/tests/unittests/auth/test_auth_preprocessor.py index f1b026cde4c..21c9b30fb92 100644 --- a/tests/unittests/auth/test_auth_preprocessor.py +++ b/tests/unittests/auth/test_auth_preprocessor.py @@ -69,6 +69,7 @@ def mock_invocation_context(self, mock_llm_agent, mock_session): context = Mock(spec=InvocationContext) context.agent = mock_llm_agent context.session = mock_session + context._get_events.side_effect = lambda **_: context.session.events return context @pytest.fixture @@ -165,7 +166,8 @@ async def test_non_llm_agent_returns_early( ): """Test that non-LLM agents return early.""" mock_context = Mock(spec=InvocationContext) - mock_context.agent = Mock() + # Using spec=[] ensures hasattr(agent, 'canonical_tools') returns False. + mock_context.agent = Mock(spec=[]) mock_context.agent.__class__.__name__ = 'BaseAgent' mock_context.session = mock_session @@ -273,6 +275,38 @@ async def test_last_event_no_auth_responses_returns_early( assert result == [] + @pytest.mark.asyncio + @patch('google.adk.auth.auth_preprocessor.AuthHandler') + @patch('google.adk.auth.auth_tool.AuthConfig.model_validate') + async def test_ignores_auth_responses_outside_current_branch( + self, + mock_auth_config_validate, + mock_auth_handler_class, + processor, + mock_invocation_context, + mock_llm_request, + mock_user_event_with_auth_response, + ): + """Test auth responses hidden by branch filtering are ignored.""" + mock_invocation_context.session.events = [ + mock_user_event_with_auth_response + ] + mock_invocation_context._get_events.side_effect = None + mock_invocation_context._get_events.return_value = [] + + result = [] + async for event in processor.run_async( + mock_invocation_context, mock_llm_request + ): + result.append(event) + + mock_invocation_context._get_events.assert_called_once_with( + current_branch=True + ) + mock_auth_config_validate.assert_not_called() + mock_auth_handler_class.assert_not_called() + assert result == [] + @pytest.mark.asyncio @patch('google.adk.auth.auth_preprocessor.AuthHandler') @patch('google.adk.auth.auth_tool.AuthConfig.model_validate') @@ -536,7 +570,8 @@ async def test_isinstance_check_for_llm_agent( # Create a mock that fails isinstance check mock_context = Mock(spec=InvocationContext) - mock_context.agent = Mock() # This will fail isinstance(agent, LlmAgent) + # This will fail isinstance(agent, LlmAgent) + mock_context.agent = Mock(spec=[]) mock_context.session = mock_session result = [] diff --git a/tests/unittests/cli/utils/test_agent_loader.py b/tests/unittests/cli/utils/test_agent_loader.py index 6e682ba7d8b..52935a9a8d0 100644 --- a/tests/unittests/cli/utils/test_agent_loader.py +++ b/tests/unittests/cli/utils/test_agent_loader.py @@ -1039,3 +1039,56 @@ def test_validate_agent_name_rejects_nonexistent_agent(self): # 'subprocess' is a valid identifier but shouldn't be importable as an agent with pytest.raises(ValueError, match="Agent not found"): loader.load_agent("subprocess") + + +class TestDetermineAgentLanguage: + """Tests for AgentLoader._determine_agent_language covering all 4 load patterns.""" + + def test_flat_module_returns_python(self): + """Flat-module agent (agents_dir/agent_name.py) is detected as python.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + (temp_path / "my_agent.py").write_text("root_agent = None\n") + loader = AgentLoader(temp_dir) + assert loader._determine_agent_language("my_agent") == "python" + + def test_agent_py_subdirectory_returns_python(self): + """Subdirectory with agent.py is detected as python.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + agent_dir = temp_path / "my_agent" + agent_dir.mkdir() + (agent_dir / "agent.py").write_text("root_agent = None\n") + loader = AgentLoader(temp_dir) + assert loader._determine_agent_language("my_agent") == "python" + + def test_init_py_subdirectory_returns_python(self): + """Subdirectory with __init__.py is detected as python.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + agent_dir = temp_path / "my_agent" + agent_dir.mkdir() + (agent_dir / "__init__.py").write_text("root_agent = None\n") + loader = AgentLoader(temp_dir) + assert loader._determine_agent_language("my_agent") == "python" + + def test_root_agent_yaml_returns_yaml(self): + """Subdirectory with root_agent.yaml is detected as yaml.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + agent_dir = temp_path / "my_agent" + agent_dir.mkdir() + (agent_dir / "root_agent.yaml").write_text("root_agent: {}\n") + loader = AgentLoader(temp_dir) + assert loader._determine_agent_language("my_agent") == "yaml" + + def test_unrecognized_structure_raises_value_error(self): + """A directory with no recognized structure raises ValueError.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + agent_dir = temp_path / "my_agent" + agent_dir.mkdir() + (agent_dir / "main.py").write_text("root_agent = None\n") + loader = AgentLoader(temp_dir) + with pytest.raises(ValueError, match="Could not determine agent type"): + loader._determine_agent_language("my_agent") diff --git a/tests/unittests/integrations/daytona/test_daytona_environment.py b/tests/unittests/integrations/daytona/test_daytona_environment.py new file mode 100644 index 00000000000..c4ee0b1b806 --- /dev/null +++ b/tests/unittests/integrations/daytona/test_daytona_environment.py @@ -0,0 +1,242 @@ +# 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 DaytonaEnvironment.""" + +from unittest import mock + +import daytona +from daytona import CreateSandboxFromImageParams +from daytona import CreateSandboxFromSnapshotParams +from daytona import DaytonaError +from daytona import DaytonaNotFoundError +from google.adk.integrations.daytona._daytona_environment import DaytonaEnvironment +import pytest + + +def _make_sandbox() -> mock.MagicMock: + """Build a mock AsyncSandbox with async method stubs.""" + sandbox = mock.MagicMock(name="AsyncSandbox") + sandbox.delete = mock.AsyncMock() + sandbox.process.exec = mock.AsyncMock() + sandbox.fs.download_file = mock.AsyncMock() + sandbox.fs.upload_file = mock.AsyncMock() + sandbox.fs.create_folder = mock.AsyncMock() + sandbox.refresh_activity = mock.AsyncMock() + return sandbox + + +@pytest.fixture(name="sandbox") +def _sandbox() -> mock.MagicMock: + return _make_sandbox() + + +@pytest.fixture(name="daytona_patch") +def _daytona_patch(sandbox: mock.MagicMock): + """Patch AsyncDaytona to return a mock client.""" + mock_client = mock.MagicMock(name="AsyncDaytona") + mock_client.create = mock.AsyncMock(return_value=sandbox) + + with mock.patch.object(daytona, "AsyncDaytona", autospec=True) as mock_class: + mock_class.return_value = mock_client + yield mock_class + + +async def test_initialize_creates_sandbox(daytona_patch, sandbox): + env = DaytonaEnvironment(image="custom-image", env_vars={"A": "1"}) + assert env.is_initialized is False + await env.initialize() + assert env.is_initialized is True + + daytona_patch.assert_called_once() + client = daytona_patch.return_value + client.create.assert_awaited_once() + + args, _ = client.create.call_args + params = args[0] + assert isinstance(params, CreateSandboxFromImageParams) + assert params.image == "custom-image" + assert params.env_vars == {"A": "1"} + assert params.auto_stop_interval == 5 + assert params.auto_delete_interval == 0 + assert env._sandbox is sandbox + + +async def test_initialize_creates_sandbox_default(daytona_patch, sandbox): + env = DaytonaEnvironment(env_vars={"B": "2"}) + await env.initialize() + + daytona_patch.assert_called_once() + client = daytona_patch.return_value + client.create.assert_awaited_once() + + args, _ = client.create.call_args + params = args[0] + assert isinstance(params, CreateSandboxFromSnapshotParams) + assert params.language == "python" + assert params.env_vars == {"B": "2"} + assert params.auto_stop_interval == 5 + assert params.auto_delete_interval == 0 + assert env._sandbox is sandbox + + +async def test_initialize_is_idempotent(daytona_patch, sandbox): + env = DaytonaEnvironment() + await env.initialize() + await env.initialize() + client = daytona_patch.return_value + client.create.assert_awaited_once() + + +async def test_close_deletes_sandbox_and_is_idempotent(daytona_patch, sandbox): + env = DaytonaEnvironment() + await env.initialize() + assert env.is_initialized is True + await env.close() + sandbox.delete.assert_awaited_once() + assert env._sandbox is None + assert env.is_initialized is False + + # Second close is a no-op. + await env.close() + sandbox.delete.assert_awaited_once() + + +async def test_working_dir_requires_initialize(): + env = DaytonaEnvironment() + with pytest.raises(RuntimeError): + _ = env.working_dir + + +async def test_execute_before_initialize_raises(): + env = DaytonaEnvironment() + with pytest.raises(RuntimeError): + await env.execute("echo hi") + + +async def test_execute_success(daytona_patch, sandbox): + class MockArtifacts: + stdout = "out" + + class MockResponse: + exit_code = 0 + artifacts = MockArtifacts() + + sandbox.process.exec.return_value = MockResponse() + + env = DaytonaEnvironment() + await env.initialize() + + result = await env.execute("echo out") + + assert result.exit_code == 0 + assert result.stdout == "out" + assert result.stderr == "" + assert result.timed_out is False + sandbox.refresh_activity.assert_awaited_once() + + +async def test_execute_timeout(daytona_patch, sandbox): + # Simulate a Daytona timeout error + sandbox.process.exec.side_effect = DaytonaError("timeout occurred") + env = DaytonaEnvironment() + await env.initialize() + + result = await env.execute("sleep 999") + + assert result.timed_out is True + + +async def test_read_file_returns_bytes(daytona_patch, sandbox): + sandbox.fs.download_file.return_value = b"data" + env = DaytonaEnvironment() + await env.initialize() + + data = await env.read_file("notes.txt") + + assert data == b"data" + sandbox.fs.download_file.assert_awaited_once_with("/workspaces/notes.txt") + sandbox.refresh_activity.assert_awaited_once() + + +async def test_read_file_absolute_path_passthrough(daytona_patch, sandbox): + sandbox.fs.download_file.return_value = b"x" + env = DaytonaEnvironment() + await env.initialize() + + await env.read_file("/etc/hostname") + + sandbox.fs.download_file.assert_awaited_once_with("/etc/hostname") + + +async def test_read_file_missing_raises(daytona_patch, sandbox): + sandbox.fs.download_file.return_value = None + env = DaytonaEnvironment() + await env.initialize() + + with pytest.raises(FileNotFoundError): + await env.read_file("missing.txt") + + +async def test_write_file_resolves_relative_path(daytona_patch, sandbox): + env = DaytonaEnvironment() + await env.initialize() + + await env.write_file("sub/out.txt", "hello") + sandbox.refresh_activity.assert_awaited_once() + + sandbox.fs.upload_file.assert_awaited_once_with( + b"hello", "/workspaces/sub/out.txt" + ) + + +async def test_initialize_propagates_api_key_and_url(daytona_patch, sandbox): + from daytona import DaytonaConfig + + env = DaytonaEnvironment(api_key="my-key", api_url="my-url") + await env.initialize() + + daytona_patch.assert_called_once() + _, kwargs = daytona_patch.call_args + config = kwargs.get("config") + assert isinstance(config, DaytonaConfig) + assert config.api_key == "my-key" + assert config.api_url == "my-url" + + +async def test_write_file_creates_parent_directory(daytona_patch, sandbox): + env = DaytonaEnvironment() + await env.initialize() + + await env.write_file("sub/nested/file.txt", "content") + + sandbox.fs.create_folder.assert_has_calls([ + mock.call("/workspaces", mode="755"), + mock.call("/workspaces/sub", mode="755"), + mock.call("/workspaces/sub/nested", mode="755"), + ]) + sandbox.fs.upload_file.assert_awaited_once_with( + b"content", "/workspaces/sub/nested/file.txt" + ) + + +async def test_read_file_raises_file_not_found_on_daytona_not_found( + daytona_patch, sandbox +): + sandbox.fs.download_file.side_effect = DaytonaNotFoundError("not found") + env = DaytonaEnvironment() + await env.initialize() + + with pytest.raises(FileNotFoundError): + await env.read_file("missing.txt") diff --git a/tests/unittests/integrations/e2b/test_e2b_environment.py b/tests/unittests/integrations/e2b/test_e2b_environment.py index 16943b9d79a..c30673eb55f 100644 --- a/tests/unittests/integrations/e2b/test_e2b_environment.py +++ b/tests/unittests/integrations/e2b/test_e2b_environment.py @@ -53,7 +53,9 @@ def _create_patch(sandbox: mock.MagicMock): @pytest.mark.asyncio async def test_initialize_creates_sandbox(create_patch, sandbox): env = E2BEnvironment(image='custom', timeout=120, env_vars={'A': '1'}) + assert env.is_initialized is False await env.initialize() + assert env.is_initialized is True create_patch.assert_awaited_once() _, kwargs = create_patch.call_args @@ -75,9 +77,11 @@ async def test_initialize_is_idempotent(create_patch, sandbox): async def test_close_kills_sandbox_and_is_idempotent(create_patch, sandbox): env = E2BEnvironment() await env.initialize() + assert env.is_initialized is True await env.close() sandbox.kill.assert_awaited_once() assert env._sandbox is None + assert env.is_initialized is False # Second close is a no-op. await env.close() sandbox.kill.assert_awaited_once() diff --git a/tests/unittests/integrations/parameter_manager/test_parameter_client.py b/tests/unittests/integrations/parameter_manager/test_parameter_client.py index 8ca4060851e..b0f75b3bb1a 100644 --- a/tests/unittests/integrations/parameter_manager/test_parameter_client.py +++ b/tests/unittests/integrations/parameter_manager/test_parameter_client.py @@ -115,7 +115,7 @@ def test_init_with_auth_token(self, mock_pm_client_class): "google.adk.integrations.parameter_manager.parameter_client.default_service_credential" ) @patch( - "google.adk.integrations.parameter_manager.parameter_client._mtls_utils.get_api_endpoint" + "google.adk.integrations.parameter_manager.parameter_client.get_api_endpoint" ) def test_init_with_location( self, diff --git a/tests/unittests/telemetry/functional_node_test_cases.py b/tests/unittests/telemetry/functional_node_test_cases.py index 56e58a7fa7a..f562565f2be 100644 --- a/tests/unittests/telemetry/functional_node_test_cases.py +++ b/tests/unittests/telemetry/functional_node_test_cases.py @@ -1513,525 +1513,433 @@ EXPECTED_STABLE_NO_CAPTURE_V2 = SpanDigest( - name="invocation", - attributes={}, + name=f"invoke_workflow {WORKFLOW_NAME}", + attributes={ + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": WORKFLOW_NAME, + "gen_ai.conversation.id": PRESENT, + }, children=[ SpanDigest( - name=f"invoke_workflow {WORKFLOW_NAME}", + name=f"invoke_agent {AGENT_NAME}", attributes={ - "gen_ai.operation.name": "invoke_workflow", - "gen_ai.workflow.name": WORKFLOW_NAME, + "gen_ai.operation.name": "invoke_agent", + "gen_ai.agent.description": AGENT_DESCRIPTION, + "gen_ai.agent.name": AGENT_NAME, "gen_ai.conversation.id": PRESENT, }, children=[ SpanDigest( - name=f"invoke_agent {AGENT_NAME}", + name="call_llm", attributes={ - "gen_ai.operation.name": "invoke_agent", - "gen_ai.agent.description": AGENT_DESCRIPTION, - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": PRESENT, + "gen_ai.system": "gcp.vertex.agent", + "gen_ai.request.model": "mock", + "gcp.vertex.agent.invocation_id": PRESENT, + "gcp.vertex.agent.session_id": PRESENT, + "gcp.vertex.agent.event_id": PRESENT, + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gen_ai.response.finish_reasons": ["stop"], }, children=[ SpanDigest( - name="call_llm", + name="generate_content mock", attributes={ - "gen_ai.system": "gcp.vertex.agent", + "gen_ai.system": "gemini", + "gen_ai.operation.name": "generate_content", "gen_ai.request.model": "mock", - "gcp.vertex.agent.invocation_id": PRESENT, - "gcp.vertex.agent.session_id": PRESENT, + "gen_ai.agent.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, "gcp.vertex.agent.event_id": PRESENT, - "gcp.vertex.agent.llm_request": "{}", - "gcp.vertex.agent.llm_response": "{}", + "gcp.vertex.agent.invocation_id": PRESENT, "gen_ai.response.finish_reasons": ["stop"], }, + logs=[ + LogDigest( + event_name=GEN_AI_CHOICE_EVENT, + body={ + "content": "", + "index": 0, + "finish_reason": "STOP", + }, + attributes={"gen_ai.system": "gemini"}, + ), + LogDigest( + event_name=GEN_AI_SYSTEM_MESSAGE_EVENT, + body={"content": ""}, + attributes={"gen_ai.system": "gemini"}, + ), + LogDigest( + event_name=GEN_AI_USER_MESSAGE_EVENT, + body={"content": ""}, + attributes={"gen_ai.system": "gemini"}, + ), + ], children=[ SpanDigest( - name="generate_content mock", + name=f"execute_tool {TOOL_NAME}", attributes={ - "gen_ai.system": "gemini", - "gen_ai.operation.name": ( - "generate_content" - ), - "gen_ai.request.model": "mock", - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": PRESENT, + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.description": ( + TOOL_DESCRIPTION + ), + "gen_ai.tool.name": TOOL_NAME, + "gen_ai.tool.type": "FunctionTool", + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gcp.vertex.agent.tool_call_args": "{}", + "gen_ai.tool.call.id": PRESENT, "gcp.vertex.agent.event_id": PRESENT, - "gcp.vertex.agent.invocation_id": ( - PRESENT - ), - "gen_ai.response.finish_reasons": [ - "stop" - ], + "gcp.vertex.agent.tool_response": "{}", }, - logs=[ - LogDigest( - event_name=GEN_AI_CHOICE_EVENT, - body={ - "content": "", - "index": 0, - "finish_reason": "STOP", - }, - attributes={ - "gen_ai.system": "gemini" - }, - ), - LogDigest( - event_name=GEN_AI_SYSTEM_MESSAGE_EVENT, - body={"content": ""}, - attributes={ - "gen_ai.system": "gemini" - }, - ), - LogDigest( - event_name=GEN_AI_USER_MESSAGE_EVENT, - body={"content": ""}, - attributes={ - "gen_ai.system": "gemini" - }, - ), - ], - children=[ - SpanDigest( - name=f"execute_tool {TOOL_NAME}", - attributes={ - "gen_ai.operation.name": ( - "execute_tool" - ), - "gen_ai.tool.description": ( - TOOL_DESCRIPTION - ), - "gen_ai.tool.name": TOOL_NAME, - "gen_ai.tool.type": ( - "FunctionTool" - ), - "gcp.vertex.agent.llm_request": ( - "{}" - ), - "gcp.vertex.agent.llm_response": ( - "{}" - ), - "gcp.vertex.agent.tool_call_args": ( - "{}" - ), - "gen_ai.tool.call.id": PRESENT, - "gcp.vertex.agent.event_id": ( - PRESENT - ), - "gcp.vertex.agent.tool_response": ( - "{}" - ), - }, - ), - ], ), ], ), + ], + ), + SpanDigest( + name="call_llm", + attributes={ + "gen_ai.system": "gcp.vertex.agent", + "gen_ai.request.model": "mock", + "gcp.vertex.agent.invocation_id": PRESENT, + "gcp.vertex.agent.session_id": PRESENT, + "gcp.vertex.agent.event_id": PRESENT, + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gen_ai.response.finish_reasons": ["stop"], + }, + children=[ SpanDigest( - name="call_llm", + name="generate_content mock", attributes={ - "gen_ai.system": "gcp.vertex.agent", + "gen_ai.system": "gemini", + "gen_ai.operation.name": "generate_content", "gen_ai.request.model": "mock", - "gcp.vertex.agent.invocation_id": PRESENT, - "gcp.vertex.agent.session_id": PRESENT, + "gen_ai.agent.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, "gcp.vertex.agent.event_id": PRESENT, - "gcp.vertex.agent.llm_request": "{}", - "gcp.vertex.agent.llm_response": "{}", + "gcp.vertex.agent.invocation_id": PRESENT, "gen_ai.response.finish_reasons": ["stop"], }, - children=[ - SpanDigest( - name="generate_content mock", - attributes={ - "gen_ai.system": "gemini", - "gen_ai.operation.name": ( - "generate_content" - ), - "gen_ai.request.model": "mock", - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": PRESENT, - "gcp.vertex.agent.event_id": PRESENT, - "gcp.vertex.agent.invocation_id": ( - PRESENT - ), - "gen_ai.response.finish_reasons": [ - "stop" - ], + logs=[ + LogDigest( + event_name=GEN_AI_CHOICE_EVENT, + body={ + "content": "", + "index": 0, + "finish_reason": "STOP", }, - logs=[ - LogDigest( - event_name=GEN_AI_CHOICE_EVENT, - body={ - "content": "", - "index": 0, - "finish_reason": "STOP", - }, - attributes={ - "gen_ai.system": "gemini" - }, - ), - LogDigest( - event_name=GEN_AI_SYSTEM_MESSAGE_EVENT, - body={"content": ""}, - attributes={ - "gen_ai.system": "gemini" - }, - ), - LogDigest( - event_name=GEN_AI_USER_MESSAGE_EVENT, - body={"content": ""}, - attributes={ - "gen_ai.system": "gemini" - }, - ), - LogDigest( - event_name=GEN_AI_USER_MESSAGE_EVENT, - body={"content": ""}, - attributes={ - "gen_ai.system": "gemini" - }, - ), - LogDigest( - event_name=GEN_AI_USER_MESSAGE_EVENT, - body={"content": ""}, - attributes={ - "gen_ai.system": "gemini" - }, - ), - ], + attributes={"gen_ai.system": "gemini"}, + ), + LogDigest( + event_name=GEN_AI_SYSTEM_MESSAGE_EVENT, + body={"content": ""}, + attributes={"gen_ai.system": "gemini"}, + ), + LogDigest( + event_name=GEN_AI_USER_MESSAGE_EVENT, + body={"content": ""}, + attributes={"gen_ai.system": "gemini"}, + ), + LogDigest( + event_name=GEN_AI_USER_MESSAGE_EVENT, + body={"content": ""}, + attributes={"gen_ai.system": "gemini"}, + ), + LogDigest( + event_name=GEN_AI_USER_MESSAGE_EVENT, + body={"content": ""}, + attributes={"gen_ai.system": "gemini"}, ), ], ), ], ), - SpanDigest( - name=f"invoke_node {NODE_NAME}", - attributes={ - "gen_ai.operation.name": "invoke_node", - "gen_ai.conversation.id": PRESENT, - "gcp.vertex.agent.associated_event_ids": PRESENT, - }, - ), ], ), + SpanDigest( + name=f"invoke_node {NODE_NAME}", + attributes={ + "gen_ai.operation.name": "invoke_node", + "gen_ai.conversation.id": PRESENT, + "gcp.vertex.agent.associated_event_ids": PRESENT, + }, + ), ], ) EXPECTED_STABLE_CAPTURE_V2 = SpanDigest( - name="invocation", - attributes={}, + name=f"invoke_workflow {WORKFLOW_NAME}", + attributes={ + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": WORKFLOW_NAME, + "gen_ai.conversation.id": PRESENT, + }, children=[ SpanDigest( - name=f"invoke_workflow {WORKFLOW_NAME}", + name=f"invoke_agent {AGENT_NAME}", attributes={ - "gen_ai.operation.name": "invoke_workflow", - "gen_ai.workflow.name": WORKFLOW_NAME, + "gen_ai.operation.name": "invoke_agent", + "gen_ai.agent.description": AGENT_DESCRIPTION, + "gen_ai.agent.name": AGENT_NAME, "gen_ai.conversation.id": PRESENT, }, children=[ SpanDigest( - name=f"invoke_agent {AGENT_NAME}", + name="call_llm", attributes={ - "gen_ai.operation.name": "invoke_agent", - "gen_ai.agent.description": AGENT_DESCRIPTION, - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": PRESENT, + "gen_ai.system": "gcp.vertex.agent", + "gen_ai.request.model": "mock", + "gcp.vertex.agent.invocation_id": PRESENT, + "gcp.vertex.agent.session_id": PRESENT, + "gcp.vertex.agent.event_id": PRESENT, + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gen_ai.response.finish_reasons": ["stop"], }, children=[ SpanDigest( - name="call_llm", + name="generate_content mock", attributes={ - "gen_ai.system": "gcp.vertex.agent", + "gen_ai.system": "gemini", + "gen_ai.operation.name": "generate_content", "gen_ai.request.model": "mock", - "gcp.vertex.agent.invocation_id": PRESENT, - "gcp.vertex.agent.session_id": PRESENT, + "gen_ai.agent.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, "gcp.vertex.agent.event_id": PRESENT, - "gcp.vertex.agent.llm_request": "{}", - "gcp.vertex.agent.llm_response": "{}", + "gcp.vertex.agent.invocation_id": PRESENT, "gen_ai.response.finish_reasons": ["stop"], }, + logs=[ + LogDigest( + event_name=GEN_AI_CHOICE_EVENT, + body={ + "content": { + "parts": [{ + "function_call": { + "args": TOOL_ARGS, + "name": TOOL_NAME, + } + }], + "role": "model", + }, + "index": 0, + "finish_reason": "STOP", + }, + attributes={"gen_ai.system": "gemini"}, + ), + LogDigest( + event_name=GEN_AI_SYSTEM_MESSAGE_EVENT, + body={"content": _NODE_SYSTEM_INSTRUCTION}, + attributes={"gen_ai.system": "gemini"}, + ), + LogDigest( + event_name=GEN_AI_USER_MESSAGE_EVENT, + body={ + "content": { + "parts": [ + {"text": _AGENT_USER_INPUT} + ], + "role": "user", + } + }, + attributes={ + "gen_ai.system": "gemini", + "user.id": "some_user", + }, + ), + ], children=[ SpanDigest( - name="generate_content mock", + name=f"execute_tool {TOOL_NAME}", attributes={ - "gen_ai.system": "gemini", - "gen_ai.operation.name": ( - "generate_content" - ), - "gen_ai.request.model": "mock", - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": PRESENT, + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.description": ( + TOOL_DESCRIPTION + ), + "gen_ai.tool.name": TOOL_NAME, + "gen_ai.tool.type": "FunctionTool", + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gcp.vertex.agent.tool_call_args": "{}", + "gen_ai.tool.call.id": PRESENT, "gcp.vertex.agent.event_id": PRESENT, - "gcp.vertex.agent.invocation_id": ( - PRESENT - ), - "gen_ai.response.finish_reasons": [ - "stop" - ], + "gcp.vertex.agent.tool_response": "{}", }, - logs=[ - LogDigest( - event_name=GEN_AI_CHOICE_EVENT, - body={ - "content": { - "parts": [{ - "function_call": { - "args": TOOL_ARGS, - "name": TOOL_NAME, - } - }], - "role": "model", - }, - "index": 0, - "finish_reason": "STOP", - }, - attributes={ - "gen_ai.system": "gemini" - }, - ), - LogDigest( - event_name=GEN_AI_SYSTEM_MESSAGE_EVENT, - body={ - "content": ( - _NODE_SYSTEM_INSTRUCTION - ) - }, - attributes={ - "gen_ai.system": "gemini" - }, - ), - LogDigest( - event_name=GEN_AI_USER_MESSAGE_EVENT, - body={ - "content": { - "parts": [{ - "text": ( - _AGENT_USER_INPUT - ) - }], - "role": "user", - } - }, - attributes={ - "gen_ai.system": "gemini", - "user.id": "some_user", - }, - ), - ], - children=[ - SpanDigest( - name=f"execute_tool {TOOL_NAME}", - attributes={ - "gen_ai.operation.name": ( - "execute_tool" - ), - "gen_ai.tool.description": ( - TOOL_DESCRIPTION - ), - "gen_ai.tool.name": TOOL_NAME, - "gen_ai.tool.type": ( - "FunctionTool" - ), - "gcp.vertex.agent.llm_request": ( - "{}" - ), - "gcp.vertex.agent.llm_response": ( - "{}" - ), - "gcp.vertex.agent.tool_call_args": ( - "{}" - ), - "gen_ai.tool.call.id": PRESENT, - "gcp.vertex.agent.event_id": ( - PRESENT - ), - "gcp.vertex.agent.tool_response": ( - "{}" - ), - }, - ), - ], ), ], ), + ], + ), + SpanDigest( + name="call_llm", + attributes={ + "gen_ai.system": "gcp.vertex.agent", + "gen_ai.request.model": "mock", + "gcp.vertex.agent.invocation_id": PRESENT, + "gcp.vertex.agent.session_id": PRESENT, + "gcp.vertex.agent.event_id": PRESENT, + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gen_ai.response.finish_reasons": ["stop"], + }, + children=[ SpanDigest( - name="call_llm", + name="generate_content mock", attributes={ - "gen_ai.system": "gcp.vertex.agent", + "gen_ai.system": "gemini", + "gen_ai.operation.name": "generate_content", "gen_ai.request.model": "mock", - "gcp.vertex.agent.invocation_id": PRESENT, - "gcp.vertex.agent.session_id": PRESENT, + "gen_ai.agent.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, "gcp.vertex.agent.event_id": PRESENT, - "gcp.vertex.agent.llm_request": "{}", - "gcp.vertex.agent.llm_response": "{}", + "gcp.vertex.agent.invocation_id": PRESENT, "gen_ai.response.finish_reasons": ["stop"], }, - children=[ - SpanDigest( - name="generate_content mock", + logs=[ + LogDigest( + event_name=GEN_AI_CHOICE_EVENT, + body={ + "content": { + "parts": [{"text": FINAL_TEXT}], + "role": "model", + }, + "index": 0, + "finish_reason": "STOP", + }, + attributes={"gen_ai.system": "gemini"}, + ), + LogDigest( + event_name=GEN_AI_SYSTEM_MESSAGE_EVENT, + body={"content": _NODE_SYSTEM_INSTRUCTION}, + attributes={"gen_ai.system": "gemini"}, + ), + LogDigest( + event_name=GEN_AI_USER_MESSAGE_EVENT, + body={ + "content": { + "parts": [{ + "function_call": { + "args": TOOL_ARGS, + "name": TOOL_NAME, + } + }], + "role": "model", + } + }, attributes={ "gen_ai.system": "gemini", - "gen_ai.operation.name": ( - "generate_content" - ), - "gen_ai.request.model": "mock", - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": PRESENT, - "gcp.vertex.agent.event_id": PRESENT, - "gcp.vertex.agent.invocation_id": ( - PRESENT - ), - "gen_ai.response.finish_reasons": [ - "stop" - ], + "user.id": "some_user", }, - logs=[ - LogDigest( - event_name=GEN_AI_CHOICE_EVENT, - body={ - "content": { - "parts": [ - {"text": FINAL_TEXT} - ], - "role": "model", - }, - "index": 0, - "finish_reason": "STOP", - }, - attributes={ - "gen_ai.system": "gemini" - }, - ), - LogDigest( - event_name=GEN_AI_SYSTEM_MESSAGE_EVENT, - body={ - "content": ( - _NODE_SYSTEM_INSTRUCTION - ) - }, - attributes={ - "gen_ai.system": "gemini" - }, - ), - LogDigest( - event_name=GEN_AI_USER_MESSAGE_EVENT, - body={ - "content": { - "parts": [{ - "function_call": { - "args": TOOL_ARGS, - "name": TOOL_NAME, - } - }], - "role": "model", - } - }, - attributes={ - "gen_ai.system": "gemini", - "user.id": "some_user", - }, - ), - LogDigest( - event_name=GEN_AI_USER_MESSAGE_EVENT, - body={ - "content": { - "parts": [{ - "function_response": { - "name": TOOL_NAME, - "response": { - "result": ( - TOOL_RESULT - ) - }, - } - }], - "role": "user", - } - }, - attributes={ - "gen_ai.system": "gemini", - "user.id": "some_user", - }, - ), - LogDigest( - event_name=GEN_AI_USER_MESSAGE_EVENT, - body={ - "content": { - "parts": [{ - "text": ( - _AGENT_USER_INPUT - ) - }], - "role": "user", + ), + LogDigest( + event_name=GEN_AI_USER_MESSAGE_EVENT, + body={ + "content": { + "parts": [{ + "function_response": { + "name": TOOL_NAME, + "response": { + "result": TOOL_RESULT + }, } - }, - attributes={ - "gen_ai.system": "gemini", - "user.id": "some_user", - }, - ), - ], + }], + "role": "user", + } + }, + attributes={ + "gen_ai.system": "gemini", + "user.id": "some_user", + }, + ), + LogDigest( + event_name=GEN_AI_USER_MESSAGE_EVENT, + body={ + "content": { + "parts": [ + {"text": _AGENT_USER_INPUT} + ], + "role": "user", + } + }, + attributes={ + "gen_ai.system": "gemini", + "user.id": "some_user", + }, ), ], ), ], ), - SpanDigest( - name=f"invoke_node {NODE_NAME}", - attributes={ - "gen_ai.operation.name": "invoke_node", - "gen_ai.conversation.id": PRESENT, - "gcp.vertex.agent.associated_event_ids": PRESENT, - }, - ), ], ), + SpanDigest( + name=f"invoke_node {NODE_NAME}", + attributes={ + "gen_ai.operation.name": "invoke_node", + "gen_ai.conversation.id": PRESENT, + "gcp.vertex.agent.associated_event_ids": PRESENT, + }, + ), ], ) EXPECTED_EXPERIMENTAL_NO_CONTENT_V2 = SpanDigest( - name="invocation", - attributes={}, + name=f"invoke_workflow {WORKFLOW_NAME}", + attributes={ + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": WORKFLOW_NAME, + "gen_ai.conversation.id": PRESENT, + }, children=[ SpanDigest( - name=f"invoke_workflow {WORKFLOW_NAME}", + name=f"invoke_agent {AGENT_NAME}", attributes={ - "gen_ai.operation.name": "invoke_workflow", - "gen_ai.workflow.name": WORKFLOW_NAME, + "gen_ai.operation.name": "invoke_agent", + "gen_ai.agent.description": AGENT_DESCRIPTION, + "gen_ai.agent.name": AGENT_NAME, "gen_ai.conversation.id": PRESENT, }, children=[ SpanDigest( - name=f"invoke_agent {AGENT_NAME}", + name="call_llm", attributes={ - "gen_ai.operation.name": "invoke_agent", - "gen_ai.agent.description": AGENT_DESCRIPTION, - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": PRESENT, + "gen_ai.system": "gcp.vertex.agent", + "gen_ai.request.model": "mock", + "gcp.vertex.agent.invocation_id": PRESENT, + "gcp.vertex.agent.session_id": PRESENT, + "gcp.vertex.agent.event_id": PRESENT, + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gen_ai.response.finish_reasons": ["stop"], }, children=[ SpanDigest( - name="call_llm", + name="generate_content mock", attributes={ - "gen_ai.system": "gcp.vertex.agent", + "gen_ai.operation.name": "generate_content", "gen_ai.request.model": "mock", - "gcp.vertex.agent.invocation_id": PRESENT, - "gcp.vertex.agent.session_id": PRESENT, + "gen_ai.agent.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, "gcp.vertex.agent.event_id": PRESENT, - "gcp.vertex.agent.llm_request": "{}", - "gcp.vertex.agent.llm_response": "{}", + "gcp.vertex.agent.invocation_id": PRESENT, "gen_ai.response.finish_reasons": ["stop"], + "gen_ai.tool.definitions": [{ + "name": TOOL_NAME, + "description": TOOL_DESCRIPTION, + "type": "function", + }], }, - children=[ - SpanDigest( - name="generate_content mock", + logs=[ + LogDigest( + event_name=( + GEN_AI_COMPLETION_DETAILS_EVENT + ), + body=None, attributes={ - "gen_ai.operation.name": ( - "generate_content" - ), - "gen_ai.request.model": "mock", "gen_ai.agent.name": AGENT_NAME, "gen_ai.conversation.id": PRESENT, "gcp.vertex.agent.event_id": PRESENT, @@ -2047,197 +1955,66 @@ "type": "function", }], }, - logs=[ - LogDigest( - event_name=( - GEN_AI_COMPLETION_DETAILS_EVENT - ), - body=None, - attributes={ - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": ( - PRESENT - ), - "gcp.vertex.agent.event_id": ( - PRESENT - ), - "gcp.vertex.agent.invocation_id": ( - PRESENT - ), - "gen_ai.response.finish_reasons": [ - "stop" - ], - "gen_ai.tool.definitions": [{ - "name": TOOL_NAME, - "description": ( - TOOL_DESCRIPTION - ), - "type": "function", - }], - }, - ), - ], - children=[ - SpanDigest( - name=f"execute_tool {TOOL_NAME}", - attributes={ - "gen_ai.operation.name": ( - "execute_tool" - ), - "gen_ai.tool.description": ( - TOOL_DESCRIPTION - ), - "gen_ai.tool.name": TOOL_NAME, - "gen_ai.tool.type": ( - "FunctionTool" - ), - "gcp.vertex.agent.llm_request": ( - "{}" - ), - "gcp.vertex.agent.llm_response": ( - "{}" - ), - "gcp.vertex.agent.tool_call_args": ( - "{}" - ), - "gen_ai.tool.call.id": PRESENT, - "gcp.vertex.agent.event_id": ( - PRESENT - ), - "gcp.vertex.agent.tool_response": ( - "{}" - ), - }, - ), - ], ), ], - ), - SpanDigest( - name="call_llm", - attributes={ - "gen_ai.system": "gcp.vertex.agent", - "gen_ai.request.model": "mock", - "gcp.vertex.agent.invocation_id": PRESENT, - "gcp.vertex.agent.session_id": PRESENT, - "gcp.vertex.agent.event_id": PRESENT, - "gcp.vertex.agent.llm_request": "{}", - "gcp.vertex.agent.llm_response": "{}", - "gen_ai.response.finish_reasons": ["stop"], - }, children=[ SpanDigest( - name="generate_content mock", + name=f"execute_tool {TOOL_NAME}", attributes={ - "gen_ai.operation.name": ( - "generate_content" - ), - "gen_ai.request.model": "mock", - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": PRESENT, + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.description": ( + TOOL_DESCRIPTION + ), + "gen_ai.tool.name": TOOL_NAME, + "gen_ai.tool.type": "FunctionTool", + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gcp.vertex.agent.tool_call_args": "{}", + "gen_ai.tool.call.id": PRESENT, "gcp.vertex.agent.event_id": PRESENT, - "gcp.vertex.agent.invocation_id": ( - PRESENT - ), - "gen_ai.response.finish_reasons": [ - "stop" - ], - "gen_ai.tool.definitions": [{ - "name": TOOL_NAME, - "description": TOOL_DESCRIPTION, - "type": "function", - }], + "gcp.vertex.agent.tool_response": "{}", }, - logs=[ - LogDigest( - event_name=( - GEN_AI_COMPLETION_DETAILS_EVENT - ), - body=None, - attributes={ - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": ( - PRESENT - ), - "gcp.vertex.agent.event_id": ( - PRESENT - ), - "gcp.vertex.agent.invocation_id": ( - PRESENT - ), - "gen_ai.response.finish_reasons": [ - "stop" - ], - "gen_ai.tool.definitions": [{ - "name": TOOL_NAME, - "description": ( - TOOL_DESCRIPTION - ), - "type": "function", - }], - }, - ), - ], ), ], ), ], ), SpanDigest( - name=f"invoke_node {NODE_NAME}", + name="call_llm", attributes={ - "gen_ai.operation.name": "invoke_node", - "gen_ai.conversation.id": PRESENT, - "gcp.vertex.agent.associated_event_ids": PRESENT, - }, - ), - ], - ), - ], -) - - -EXPECTED_EXPERIMENTAL_SPAN_ONLY_V2 = SpanDigest( - name="invocation", - attributes={}, - children=[ - SpanDigest( - name=f"invoke_workflow {WORKFLOW_NAME}", - attributes={ - "gen_ai.operation.name": "invoke_workflow", - "gen_ai.workflow.name": WORKFLOW_NAME, - "gen_ai.conversation.id": PRESENT, - }, - children=[ - SpanDigest( - name=f"invoke_agent {AGENT_NAME}", - attributes={ - "gen_ai.operation.name": "invoke_agent", - "gen_ai.agent.description": AGENT_DESCRIPTION, - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": PRESENT, + "gen_ai.system": "gcp.vertex.agent", + "gen_ai.request.model": "mock", + "gcp.vertex.agent.invocation_id": PRESENT, + "gcp.vertex.agent.session_id": PRESENT, + "gcp.vertex.agent.event_id": PRESENT, + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gen_ai.response.finish_reasons": ["stop"], }, children=[ SpanDigest( - name="call_llm", + name="generate_content mock", attributes={ - "gen_ai.system": "gcp.vertex.agent", + "gen_ai.operation.name": "generate_content", "gen_ai.request.model": "mock", - "gcp.vertex.agent.invocation_id": PRESENT, - "gcp.vertex.agent.session_id": PRESENT, + "gen_ai.agent.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, "gcp.vertex.agent.event_id": PRESENT, - "gcp.vertex.agent.llm_request": "{}", - "gcp.vertex.agent.llm_response": "{}", + "gcp.vertex.agent.invocation_id": PRESENT, "gen_ai.response.finish_reasons": ["stop"], + "gen_ai.tool.definitions": [{ + "name": TOOL_NAME, + "description": TOOL_DESCRIPTION, + "type": "function", + }], }, - children=[ - SpanDigest( - name="generate_content mock", + logs=[ + LogDigest( + event_name=( + GEN_AI_COMPLETION_DETAILS_EVENT + ), + body=None, attributes={ - "gen_ai.operation.name": ( - "generate_content" - ), - "gen_ai.request.model": "mock", "gen_ai.agent.name": AGENT_NAME, "gen_ai.conversation.id": PRESENT, "gcp.vertex.agent.event_id": PRESENT, @@ -2247,101 +2024,89 @@ "gen_ai.response.finish_reasons": [ "stop" ], - "gen_ai.input.messages": ( - _TURN_1_INPUT_MESSAGES - ), - "gen_ai.system_instructions": ( - _SYSTEM_INSTRUCTIONS - ), - "gen_ai.tool.definitions": [ - _TOOL_DEFINITION_FULL - ], - "gen_ai.output.messages": ( - _TURN_1_OUTPUT_MESSAGES - ), - }, - logs=[ - LogDigest( - event_name=( - GEN_AI_COMPLETION_DETAILS_EVENT - ), - body=None, - attributes={ - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": ( - PRESENT - ), - "gcp.vertex.agent.event_id": ( - PRESENT - ), - "gcp.vertex.agent.invocation_id": ( - PRESENT - ), - "gen_ai.response.finish_reasons": [ - "stop" - ], - "gen_ai.tool.definitions": [ - _TOOL_DEFINITION_NO_CONTENT - ], - }, - ), - ], - children=[ - SpanDigest( - name=f"execute_tool {TOOL_NAME}", - attributes={ - "gen_ai.operation.name": ( - "execute_tool" - ), - "gen_ai.tool.description": ( - TOOL_DESCRIPTION - ), - "gen_ai.tool.name": TOOL_NAME, - "gen_ai.tool.type": ( - "FunctionTool" - ), - "gcp.vertex.agent.llm_request": ( - "{}" - ), - "gcp.vertex.agent.llm_response": ( - "{}" - ), - "gcp.vertex.agent.tool_call_args": ( - "{}" - ), - "gen_ai.tool.call.id": PRESENT, - "gcp.vertex.agent.event_id": ( - PRESENT - ), - "gcp.vertex.agent.tool_response": ( - "{}" - ), - }, - ), - ], + "gen_ai.tool.definitions": [{ + "name": TOOL_NAME, + "description": TOOL_DESCRIPTION, + "type": "function", + }], + }, ), ], ), + ], + ), + ], + ), + SpanDigest( + name=f"invoke_node {NODE_NAME}", + attributes={ + "gen_ai.operation.name": "invoke_node", + "gen_ai.conversation.id": PRESENT, + "gcp.vertex.agent.associated_event_ids": PRESENT, + }, + ), + ], +) + + +EXPECTED_EXPERIMENTAL_SPAN_ONLY_V2 = SpanDigest( + name=f"invoke_workflow {WORKFLOW_NAME}", + attributes={ + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": WORKFLOW_NAME, + "gen_ai.conversation.id": PRESENT, + }, + children=[ + SpanDigest( + name=f"invoke_agent {AGENT_NAME}", + attributes={ + "gen_ai.operation.name": "invoke_agent", + "gen_ai.agent.description": AGENT_DESCRIPTION, + "gen_ai.agent.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, + }, + children=[ + SpanDigest( + name="call_llm", + attributes={ + "gen_ai.system": "gcp.vertex.agent", + "gen_ai.request.model": "mock", + "gcp.vertex.agent.invocation_id": PRESENT, + "gcp.vertex.agent.session_id": PRESENT, + "gcp.vertex.agent.event_id": PRESENT, + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gen_ai.response.finish_reasons": ["stop"], + }, + children=[ SpanDigest( - name="call_llm", + name="generate_content mock", attributes={ - "gen_ai.system": "gcp.vertex.agent", + "gen_ai.operation.name": "generate_content", "gen_ai.request.model": "mock", - "gcp.vertex.agent.invocation_id": PRESENT, - "gcp.vertex.agent.session_id": PRESENT, + "gen_ai.agent.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, "gcp.vertex.agent.event_id": PRESENT, - "gcp.vertex.agent.llm_request": "{}", - "gcp.vertex.agent.llm_response": "{}", + "gcp.vertex.agent.invocation_id": PRESENT, "gen_ai.response.finish_reasons": ["stop"], + "gen_ai.input.messages": _TURN_1_INPUT_MESSAGES, + "gen_ai.system_instructions": ( + _SYSTEM_INSTRUCTIONS + ), + "gen_ai.tool.definitions": [ + _TOOL_DEFINITION_FULL + ], + "gen_ai.output.messages": ( + _TURN_1_OUTPUT_MESSAGES + ), }, - children=[ - SpanDigest( - name="generate_content mock", + logs=[ + LogDigest( + event_name=( + GEN_AI_COMPLETION_DETAILS_EVENT + ), + body=None, attributes={ - "gen_ai.operation.name": ( - "generate_content" - ), - "gen_ai.request.model": "mock", "gen_ai.agent.name": AGENT_NAME, "gen_ai.conversation.id": PRESENT, "gcp.vertex.agent.event_id": PRESENT, @@ -2351,107 +2116,161 @@ "gen_ai.response.finish_reasons": [ "stop" ], - "gen_ai.input.messages": ( - _TURN_2_INPUT_MESSAGES - ), - "gen_ai.system_instructions": ( - _SYSTEM_INSTRUCTIONS - ), "gen_ai.tool.definitions": [ - _TOOL_DEFINITION_FULL + _TOOL_DEFINITION_NO_CONTENT ], - "gen_ai.output.messages": ( - _TURN_2_OUTPUT_MESSAGES - ), }, - logs=[ - LogDigest( - event_name=( - GEN_AI_COMPLETION_DETAILS_EVENT - ), - body=None, - attributes={ - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": ( - PRESENT - ), - "gcp.vertex.agent.event_id": ( - PRESENT - ), - "gcp.vertex.agent.invocation_id": ( - PRESENT - ), - "gen_ai.response.finish_reasons": [ - "stop" - ], - "gen_ai.tool.definitions": [ - _TOOL_DEFINITION_NO_CONTENT - ], - }, - ), - ], + ), + ], + children=[ + SpanDigest( + name=f"execute_tool {TOOL_NAME}", + attributes={ + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.description": ( + TOOL_DESCRIPTION + ), + "gen_ai.tool.name": TOOL_NAME, + "gen_ai.tool.type": "FunctionTool", + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gcp.vertex.agent.tool_call_args": "{}", + "gen_ai.tool.call.id": PRESENT, + "gcp.vertex.agent.event_id": PRESENT, + "gcp.vertex.agent.tool_response": "{}", + }, ), ], ), ], ), SpanDigest( - name=f"invoke_node {NODE_NAME}", + name="call_llm", attributes={ - "gen_ai.operation.name": "invoke_node", - "gen_ai.conversation.id": PRESENT, - "gcp.vertex.agent.associated_event_ids": PRESENT, + "gen_ai.system": "gcp.vertex.agent", + "gen_ai.request.model": "mock", + "gcp.vertex.agent.invocation_id": PRESENT, + "gcp.vertex.agent.session_id": PRESENT, + "gcp.vertex.agent.event_id": PRESENT, + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gen_ai.response.finish_reasons": ["stop"], }, + children=[ + SpanDigest( + name="generate_content mock", + attributes={ + "gen_ai.operation.name": "generate_content", + "gen_ai.request.model": "mock", + "gen_ai.agent.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, + "gcp.vertex.agent.event_id": PRESENT, + "gcp.vertex.agent.invocation_id": PRESENT, + "gen_ai.response.finish_reasons": ["stop"], + "gen_ai.input.messages": _TURN_2_INPUT_MESSAGES, + "gen_ai.system_instructions": ( + _SYSTEM_INSTRUCTIONS + ), + "gen_ai.tool.definitions": [ + _TOOL_DEFINITION_FULL + ], + "gen_ai.output.messages": ( + _TURN_2_OUTPUT_MESSAGES + ), + }, + logs=[ + LogDigest( + event_name=( + GEN_AI_COMPLETION_DETAILS_EVENT + ), + body=None, + attributes={ + "gen_ai.agent.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, + "gcp.vertex.agent.event_id": PRESENT, + "gcp.vertex.agent.invocation_id": ( + PRESENT + ), + "gen_ai.response.finish_reasons": [ + "stop" + ], + "gen_ai.tool.definitions": [ + _TOOL_DEFINITION_NO_CONTENT + ], + }, + ), + ], + ), + ], ), ], ), + SpanDigest( + name=f"invoke_node {NODE_NAME}", + attributes={ + "gen_ai.operation.name": "invoke_node", + "gen_ai.conversation.id": PRESENT, + "gcp.vertex.agent.associated_event_ids": PRESENT, + }, + ), ], ) EXPECTED_EXPERIMENTAL_EVENT_ONLY_V2 = SpanDigest( - name="invocation", - attributes={}, + name=f"invoke_workflow {WORKFLOW_NAME}", + attributes={ + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": WORKFLOW_NAME, + "gen_ai.conversation.id": PRESENT, + }, children=[ SpanDigest( - name=f"invoke_workflow {WORKFLOW_NAME}", + name=f"invoke_agent {AGENT_NAME}", attributes={ - "gen_ai.operation.name": "invoke_workflow", - "gen_ai.workflow.name": WORKFLOW_NAME, + "gen_ai.operation.name": "invoke_agent", + "gen_ai.agent.description": AGENT_DESCRIPTION, + "gen_ai.agent.name": AGENT_NAME, "gen_ai.conversation.id": PRESENT, }, children=[ SpanDigest( - name=f"invoke_agent {AGENT_NAME}", + name="call_llm", attributes={ - "gen_ai.operation.name": "invoke_agent", - "gen_ai.agent.description": AGENT_DESCRIPTION, - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": PRESENT, + "gen_ai.system": "gcp.vertex.agent", + "gen_ai.request.model": "mock", + "gcp.vertex.agent.invocation_id": PRESENT, + "gcp.vertex.agent.session_id": PRESENT, + "gcp.vertex.agent.event_id": PRESENT, + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gen_ai.response.finish_reasons": ["stop"], }, children=[ SpanDigest( - name="call_llm", + name="generate_content mock", attributes={ - "gen_ai.system": "gcp.vertex.agent", + "gen_ai.operation.name": "generate_content", "gen_ai.request.model": "mock", - "gcp.vertex.agent.invocation_id": PRESENT, - "gcp.vertex.agent.session_id": PRESENT, + "gen_ai.agent.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, "gcp.vertex.agent.event_id": PRESENT, - "gcp.vertex.agent.llm_request": "{}", - "gcp.vertex.agent.llm_response": "{}", + "gcp.vertex.agent.invocation_id": PRESENT, "gen_ai.response.finish_reasons": ["stop"], + "gen_ai.tool.definitions": [ + _TOOL_DEFINITION_NO_CONTENT + ], }, - children=[ - SpanDigest( - name="generate_content mock", + logs=[ + LogDigest( + event_name=( + GEN_AI_COMPLETION_DETAILS_EVENT + ), + body=None, attributes={ - "gen_ai.operation.name": ( - "generate_content" - ), - "gen_ai.request.model": "mock", "gen_ai.agent.name": AGENT_NAME, "gen_ai.conversation.id": PRESENT, + "user.id": "some_user", "gcp.vertex.agent.event_id": PRESENT, "gcp.vertex.agent.invocation_id": ( PRESENT @@ -2459,104 +2278,80 @@ "gen_ai.response.finish_reasons": [ "stop" ], + "gen_ai.input.messages": ( + _TURN_1_INPUT_MESSAGES + ), + "gen_ai.system_instructions": ( + _SYSTEM_INSTRUCTIONS + ), "gen_ai.tool.definitions": [ - _TOOL_DEFINITION_NO_CONTENT + _TOOL_DEFINITION_FULL ], - }, - logs=[ - LogDigest( - event_name=( - GEN_AI_COMPLETION_DETAILS_EVENT - ), - body=None, - attributes={ - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": ( - PRESENT - ), - "user.id": "some_user", - "gcp.vertex.agent.event_id": ( - PRESENT - ), - "gcp.vertex.agent.invocation_id": ( - PRESENT - ), - "gen_ai.response.finish_reasons": [ - "stop" - ], - "gen_ai.input.messages": ( - _TURN_1_INPUT_MESSAGES - ), - "gen_ai.system_instructions": ( - _SYSTEM_INSTRUCTIONS - ), - "gen_ai.tool.definitions": [ - _TOOL_DEFINITION_FULL - ], - "gen_ai.output.messages": ( - _TURN_1_OUTPUT_MESSAGES - ), - }, - ), - ], - children=[ - SpanDigest( - name=f"execute_tool {TOOL_NAME}", - attributes={ - "gen_ai.operation.name": ( - "execute_tool" - ), - "gen_ai.tool.description": ( - TOOL_DESCRIPTION - ), - "gen_ai.tool.name": TOOL_NAME, - "gen_ai.tool.type": ( - "FunctionTool" - ), - "gcp.vertex.agent.llm_request": ( - "{}" - ), - "gcp.vertex.agent.llm_response": ( - "{}" - ), - "gcp.vertex.agent.tool_call_args": ( - "{}" - ), - "gen_ai.tool.call.id": PRESENT, - "gcp.vertex.agent.event_id": ( - PRESENT - ), - "gcp.vertex.agent.tool_response": ( - "{}" - ), - }, + "gen_ai.output.messages": ( + _TURN_1_OUTPUT_MESSAGES ), - ], + }, + ), + ], + children=[ + SpanDigest( + name=f"execute_tool {TOOL_NAME}", + attributes={ + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.description": ( + TOOL_DESCRIPTION + ), + "gen_ai.tool.name": TOOL_NAME, + "gen_ai.tool.type": "FunctionTool", + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gcp.vertex.agent.tool_call_args": "{}", + "gen_ai.tool.call.id": PRESENT, + "gcp.vertex.agent.event_id": PRESENT, + "gcp.vertex.agent.tool_response": "{}", + }, ), ], ), + ], + ), + SpanDigest( + name="call_llm", + attributes={ + "gen_ai.system": "gcp.vertex.agent", + "gen_ai.request.model": "mock", + "gcp.vertex.agent.invocation_id": PRESENT, + "gcp.vertex.agent.session_id": PRESENT, + "gcp.vertex.agent.event_id": PRESENT, + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gen_ai.response.finish_reasons": ["stop"], + }, + children=[ SpanDigest( - name="call_llm", + name="generate_content mock", attributes={ - "gen_ai.system": "gcp.vertex.agent", + "gen_ai.operation.name": "generate_content", "gen_ai.request.model": "mock", - "gcp.vertex.agent.invocation_id": PRESENT, - "gcp.vertex.agent.session_id": PRESENT, + "gen_ai.agent.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, "gcp.vertex.agent.event_id": PRESENT, - "gcp.vertex.agent.llm_request": "{}", - "gcp.vertex.agent.llm_response": "{}", + "gcp.vertex.agent.invocation_id": PRESENT, "gen_ai.response.finish_reasons": ["stop"], - }, - children=[ - SpanDigest( - name="generate_content mock", + "gen_ai.tool.definitions": [ + _TOOL_DEFINITION_NO_CONTENT + ], + }, + logs=[ + LogDigest( + event_name=( + GEN_AI_COMPLETION_DETAILS_EVENT + ), + body=None, attributes={ - "gen_ai.operation.name": ( - "generate_content" - ), - "gen_ai.request.model": "mock", "gen_ai.agent.name": AGENT_NAME, "gen_ai.conversation.id": PRESENT, + "user.id": "some_user", "gcp.vertex.agent.event_id": PRESENT, "gcp.vertex.agent.invocation_id": ( PRESENT @@ -2564,108 +2359,99 @@ "gen_ai.response.finish_reasons": [ "stop" ], + "gen_ai.input.messages": ( + _TURN_2_INPUT_MESSAGES + ), + "gen_ai.system_instructions": ( + _SYSTEM_INSTRUCTIONS + ), "gen_ai.tool.definitions": [ - _TOOL_DEFINITION_NO_CONTENT + _TOOL_DEFINITION_FULL ], - }, - logs=[ - LogDigest( - event_name=( - GEN_AI_COMPLETION_DETAILS_EVENT - ), - body=None, - attributes={ - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": ( - PRESENT - ), - "user.id": "some_user", - "gcp.vertex.agent.event_id": ( - PRESENT - ), - "gcp.vertex.agent.invocation_id": ( - PRESENT - ), - "gen_ai.response.finish_reasons": [ - "stop" - ], - "gen_ai.input.messages": ( - _TURN_2_INPUT_MESSAGES - ), - "gen_ai.system_instructions": ( - _SYSTEM_INSTRUCTIONS - ), - "gen_ai.tool.definitions": [ - _TOOL_DEFINITION_FULL - ], - "gen_ai.output.messages": ( - _TURN_2_OUTPUT_MESSAGES - ), - }, + "gen_ai.output.messages": ( + _TURN_2_OUTPUT_MESSAGES ), - ], + }, ), ], ), ], ), - SpanDigest( - name=f"invoke_node {NODE_NAME}", - attributes={ - "gen_ai.operation.name": "invoke_node", - "gen_ai.conversation.id": PRESENT, - "gcp.vertex.agent.associated_event_ids": PRESENT, - }, - ), ], ), + SpanDigest( + name=f"invoke_node {NODE_NAME}", + attributes={ + "gen_ai.operation.name": "invoke_node", + "gen_ai.conversation.id": PRESENT, + "gcp.vertex.agent.associated_event_ids": PRESENT, + }, + ), ], ) EXPECTED_EXPERIMENTAL_SPAN_AND_EVENT_V2 = SpanDigest( - name="invocation", - attributes={}, + name=f"invoke_workflow {WORKFLOW_NAME}", + attributes={ + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": WORKFLOW_NAME, + "gen_ai.conversation.id": PRESENT, + }, children=[ SpanDigest( - name=f"invoke_workflow {WORKFLOW_NAME}", + name=f"invoke_agent {AGENT_NAME}", attributes={ - "gen_ai.operation.name": "invoke_workflow", - "gen_ai.workflow.name": WORKFLOW_NAME, + "gen_ai.operation.name": "invoke_agent", + "gen_ai.agent.description": AGENT_DESCRIPTION, + "gen_ai.agent.name": AGENT_NAME, "gen_ai.conversation.id": PRESENT, }, children=[ SpanDigest( - name=f"invoke_agent {AGENT_NAME}", + name="call_llm", attributes={ - "gen_ai.operation.name": "invoke_agent", - "gen_ai.agent.description": AGENT_DESCRIPTION, - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": PRESENT, + "gen_ai.system": "gcp.vertex.agent", + "gen_ai.request.model": "mock", + "gcp.vertex.agent.invocation_id": PRESENT, + "gcp.vertex.agent.session_id": PRESENT, + "gcp.vertex.agent.event_id": PRESENT, + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gen_ai.response.finish_reasons": ["stop"], }, children=[ SpanDigest( - name="call_llm", + name="generate_content mock", attributes={ - "gen_ai.system": "gcp.vertex.agent", + "gen_ai.operation.name": "generate_content", "gen_ai.request.model": "mock", - "gcp.vertex.agent.invocation_id": PRESENT, - "gcp.vertex.agent.session_id": PRESENT, + "gen_ai.agent.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, "gcp.vertex.agent.event_id": PRESENT, - "gcp.vertex.agent.llm_request": "{}", - "gcp.vertex.agent.llm_response": "{}", + "gcp.vertex.agent.invocation_id": PRESENT, "gen_ai.response.finish_reasons": ["stop"], + "gen_ai.input.messages": _TURN_1_INPUT_MESSAGES, + "gen_ai.system_instructions": ( + _SYSTEM_INSTRUCTIONS + ), + "gen_ai.tool.definitions": [ + _TOOL_DEFINITION_FULL + ], + "gen_ai.output.messages": ( + _TURN_1_OUTPUT_MESSAGES + ), }, - children=[ - SpanDigest( - name="generate_content mock", + logs=[ + LogDigest( + event_name=( + GEN_AI_COMPLETION_DETAILS_EVENT + ), + body=None, attributes={ - "gen_ai.operation.name": ( - "generate_content" - ), - "gen_ai.request.model": "mock", "gen_ai.agent.name": AGENT_NAME, "gen_ai.conversation.id": PRESENT, + "user.id": "some_user", "gcp.vertex.agent.event_id": PRESENT, "gcp.vertex.agent.invocation_id": ( PRESENT @@ -2686,100 +2472,74 @@ _TURN_1_OUTPUT_MESSAGES ), }, - logs=[ - LogDigest( - event_name=( - GEN_AI_COMPLETION_DETAILS_EVENT - ), - body=None, - attributes={ - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": ( - PRESENT - ), - "user.id": "some_user", - "gcp.vertex.agent.event_id": ( - PRESENT - ), - "gcp.vertex.agent.invocation_id": ( - PRESENT - ), - "gen_ai.response.finish_reasons": [ - "stop" - ], - "gen_ai.input.messages": ( - _TURN_1_INPUT_MESSAGES - ), - "gen_ai.system_instructions": ( - _SYSTEM_INSTRUCTIONS - ), - "gen_ai.tool.definitions": [ - _TOOL_DEFINITION_FULL - ], - "gen_ai.output.messages": ( - _TURN_1_OUTPUT_MESSAGES - ), - }, - ), - ], - children=[ - SpanDigest( - name=f"execute_tool {TOOL_NAME}", - attributes={ - "gen_ai.operation.name": ( - "execute_tool" - ), - "gen_ai.tool.description": ( - TOOL_DESCRIPTION - ), - "gen_ai.tool.name": TOOL_NAME, - "gen_ai.tool.type": ( - "FunctionTool" - ), - "gcp.vertex.agent.llm_request": ( - "{}" - ), - "gcp.vertex.agent.llm_response": ( - "{}" - ), - "gcp.vertex.agent.tool_call_args": ( - "{}" - ), - "gen_ai.tool.call.id": PRESENT, - "gcp.vertex.agent.event_id": ( - PRESENT - ), - "gcp.vertex.agent.tool_response": ( - "{}" - ), - }, - ), - ], + ), + ], + children=[ + SpanDigest( + name=f"execute_tool {TOOL_NAME}", + attributes={ + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.description": ( + TOOL_DESCRIPTION + ), + "gen_ai.tool.name": TOOL_NAME, + "gen_ai.tool.type": "FunctionTool", + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gcp.vertex.agent.tool_call_args": "{}", + "gen_ai.tool.call.id": PRESENT, + "gcp.vertex.agent.event_id": PRESENT, + "gcp.vertex.agent.tool_response": "{}", + }, ), ], ), + ], + ), + SpanDigest( + name="call_llm", + attributes={ + "gen_ai.system": "gcp.vertex.agent", + "gen_ai.request.model": "mock", + "gcp.vertex.agent.invocation_id": PRESENT, + "gcp.vertex.agent.session_id": PRESENT, + "gcp.vertex.agent.event_id": PRESENT, + "gcp.vertex.agent.llm_request": "{}", + "gcp.vertex.agent.llm_response": "{}", + "gen_ai.response.finish_reasons": ["stop"], + }, + children=[ SpanDigest( - name="call_llm", + name="generate_content mock", attributes={ - "gen_ai.system": "gcp.vertex.agent", + "gen_ai.operation.name": "generate_content", "gen_ai.request.model": "mock", - "gcp.vertex.agent.invocation_id": PRESENT, - "gcp.vertex.agent.session_id": PRESENT, + "gen_ai.agent.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, "gcp.vertex.agent.event_id": PRESENT, - "gcp.vertex.agent.llm_request": "{}", - "gcp.vertex.agent.llm_response": "{}", + "gcp.vertex.agent.invocation_id": PRESENT, "gen_ai.response.finish_reasons": ["stop"], + "gen_ai.input.messages": _TURN_2_INPUT_MESSAGES, + "gen_ai.system_instructions": ( + _SYSTEM_INSTRUCTIONS + ), + "gen_ai.tool.definitions": [ + _TOOL_DEFINITION_FULL + ], + "gen_ai.output.messages": ( + _TURN_2_OUTPUT_MESSAGES + ), }, - children=[ - SpanDigest( - name="generate_content mock", + logs=[ + LogDigest( + event_name=( + GEN_AI_COMPLETION_DETAILS_EVENT + ), + body=None, attributes={ - "gen_ai.operation.name": ( - "generate_content" - ), - "gen_ai.request.model": "mock", "gen_ai.agent.name": AGENT_NAME, "gen_ai.conversation.id": PRESENT, + "user.id": "some_user", "gcp.vertex.agent.event_id": PRESENT, "gcp.vertex.agent.invocation_id": ( PRESENT @@ -2800,70 +2560,34 @@ _TURN_2_OUTPUT_MESSAGES ), }, - logs=[ - LogDigest( - event_name=( - GEN_AI_COMPLETION_DETAILS_EVENT - ), - body=None, - attributes={ - "gen_ai.agent.name": AGENT_NAME, - "gen_ai.conversation.id": ( - PRESENT - ), - "user.id": "some_user", - "gcp.vertex.agent.event_id": ( - PRESENT - ), - "gcp.vertex.agent.invocation_id": ( - PRESENT - ), - "gen_ai.response.finish_reasons": [ - "stop" - ], - "gen_ai.input.messages": ( - _TURN_2_INPUT_MESSAGES - ), - "gen_ai.system_instructions": ( - _SYSTEM_INSTRUCTIONS - ), - "gen_ai.tool.definitions": [ - _TOOL_DEFINITION_FULL - ], - "gen_ai.output.messages": ( - _TURN_2_OUTPUT_MESSAGES - ), - }, - ), - ], ), ], ), ], ), - SpanDigest( - name=f"invoke_node {NODE_NAME}", - attributes={ - "gen_ai.operation.name": "invoke_node", - "gen_ai.conversation.id": PRESENT, - "gcp.vertex.agent.associated_event_ids": PRESENT, - }, - ), ], ), + SpanDigest( + name=f"invoke_node {NODE_NAME}", + attributes={ + "gen_ai.operation.name": "invoke_node", + "gen_ai.conversation.id": PRESENT, + "gcp.vertex.agent.associated_event_ids": PRESENT, + }, + ), ], ) # Expected metric points, grouped by metric name. EXPECTED_NODE_METRICS_V1: dict[str, frozenset[MetricPoint]] = { - "gen_ai.agent.invocation.duration": frozenset({ + "gen_ai.invoke_agent.duration": frozenset({ MetricPoint( attributes={"gen_ai.agent.name": AGENT_NAME}, value=NON_DETERMINISTIC, ), }), - "gen_ai.tool.execution.duration": frozenset({ + "gen_ai.execute_tool.duration": frozenset({ MetricPoint( attributes={ "gen_ai.agent.name": AGENT_NAME, @@ -2900,17 +2624,26 @@ value=NON_DETERMINISTIC, ), }), + "gen_ai.invoke_workflow.duration": frozenset({ + MetricPoint( + attributes={ + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": WORKFLOW_NAME, + }, + value=NON_DETERMINISTIC, + ), + }), } EXPECTED_NODE_METRICS_V2: dict[str, frozenset[MetricPoint]] = { - "gen_ai.agent.invocation.duration": frozenset({ + "gen_ai.invoke_agent.duration": frozenset({ MetricPoint( attributes={"gen_ai.agent.name": AGENT_NAME}, value=NON_DETERMINISTIC, ), }), - "gen_ai.tool.execution.duration": frozenset({ + "gen_ai.execute_tool.duration": frozenset({ MetricPoint( attributes={ "gen_ai.agent.name": AGENT_NAME, @@ -2947,6 +2680,15 @@ value=NON_DETERMINISTIC, ), }), + "gen_ai.invoke_workflow.duration": frozenset({ + MetricPoint( + attributes={ + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": WORKFLOW_NAME, + }, + value=NON_DETERMINISTIC, + ), + }), } diff --git a/tests/unittests/telemetry/functional_test_cases.py b/tests/unittests/telemetry/functional_test_cases.py index 44c352a4c88..ea6383acb5b 100644 --- a/tests/unittests/telemetry/functional_test_cases.py +++ b/tests/unittests/telemetry/functional_test_cases.py @@ -1275,8 +1275,12 @@ EXPECTED_STABLE_NO_CAPTURE_V2 = SpanDigest( - name="invocation", - attributes={}, + name="invoke_workflow some_root_agent", + attributes={ + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, + }, children=[ SpanDigest( name="invoke_agent some_root_agent", @@ -1421,8 +1425,12 @@ EXPECTED_STABLE_CAPTURE_V2 = SpanDigest( - name="invocation", - attributes={}, + name="invoke_workflow some_root_agent", + attributes={ + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, + }, children=[ SpanDigest( name="invoke_agent some_root_agent", @@ -1622,8 +1630,12 @@ EXPECTED_EXPERIMENTAL_NO_CONTENT_V2 = SpanDigest( - name="invocation", - attributes={}, + name="invoke_workflow some_root_agent", + attributes={ + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, + }, children=[ SpanDigest( name="invoke_agent some_root_agent", @@ -1768,8 +1780,12 @@ EXPECTED_EXPERIMENTAL_SPAN_ONLY_V2 = SpanDigest( - name="invocation", - attributes={}, + name="invoke_workflow some_root_agent", + attributes={ + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, + }, children=[ SpanDigest( name="invoke_agent some_root_agent", @@ -1920,8 +1936,12 @@ EXPECTED_EXPERIMENTAL_EVENT_ONLY_V2 = SpanDigest( - name="invocation", - attributes={}, + name="invoke_workflow some_root_agent", + attributes={ + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, + }, children=[ SpanDigest( name="invoke_agent some_root_agent", @@ -2078,8 +2098,12 @@ EXPECTED_EXPERIMENTAL_SPAN_AND_EVENT_V2 = SpanDigest( - name="invocation", - attributes={}, + name="invoke_workflow some_root_agent", + attributes={ + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": AGENT_NAME, + "gen_ai.conversation.id": PRESENT, + }, children=[ SpanDigest( name="invoke_agent some_root_agent", @@ -2251,13 +2275,13 @@ # Expected metric points, grouped by metric name. EXPECTED_METRICS_V1: dict[str, frozenset[MetricPoint]] = { - "gen_ai.agent.invocation.duration": frozenset({ + "gen_ai.invoke_agent.duration": frozenset({ MetricPoint( attributes={"gen_ai.agent.name": AGENT_NAME}, value=NON_DETERMINISTIC, ), }), - "gen_ai.tool.execution.duration": frozenset({ + "gen_ai.execute_tool.duration": frozenset({ MetricPoint( attributes={ "gen_ai.agent.name": AGENT_NAME, @@ -2298,13 +2322,13 @@ EXPECTED_METRICS_V2: dict[str, frozenset[MetricPoint]] = { - "gen_ai.agent.invocation.duration": frozenset({ + "gen_ai.invoke_agent.duration": frozenset({ MetricPoint( attributes={"gen_ai.agent.name": AGENT_NAME}, value=NON_DETERMINISTIC, ), }), - "gen_ai.tool.execution.duration": frozenset({ + "gen_ai.execute_tool.duration": frozenset({ MetricPoint( attributes={ "gen_ai.agent.name": AGENT_NAME, @@ -2341,6 +2365,15 @@ value=NON_DETERMINISTIC, ), }), + "gen_ai.invoke_workflow.duration": frozenset({ + MetricPoint( + attributes={ + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": AGENT_NAME, + }, + value=NON_DETERMINISTIC, + ), + }), } diff --git a/tests/unittests/telemetry/functional_test_helpers.py b/tests/unittests/telemetry/functional_test_helpers.py index e4ae30a3607..031e8e0a54a 100644 --- a/tests/unittests/telemetry/functional_test_helpers.py +++ b/tests/unittests/telemetry/functional_test_helpers.py @@ -307,12 +307,12 @@ class HistogramSpec(NamedTuple): HistogramSpec( module=_metrics, attr="_agent_invocation_duration", - metric_name="gen_ai.agent.invocation.duration", + metric_name="gen_ai.invoke_agent.duration", ), HistogramSpec( module=_metrics, attr="_tool_execution_duration", - metric_name="gen_ai.tool.execution.duration", + metric_name="gen_ai.execute_tool.duration", ), HistogramSpec( module=_metrics, @@ -339,6 +339,11 @@ class HistogramSpec(NamedTuple): attr="_client_token_usage", metric_name="gen_ai.client.token.usage", ), + HistogramSpec( + module=_metrics, + attr="_workflow_invocation_duration", + metric_name="gen_ai.invoke_workflow.duration", + ), ) diff --git a/tests/unittests/telemetry/test_instrumentation.py b/tests/unittests/telemetry/test_instrumentation.py index 8711aa979ef..dc41adb51fd 100644 --- a/tests/unittests/telemetry/test_instrumentation.py +++ b/tests/unittests/telemetry/test_instrumentation.py @@ -18,6 +18,7 @@ from unittest import mock from google.adk.telemetry import _instrumentation +from google.adk.telemetry import _metrics from opentelemetry import trace import pytest @@ -26,7 +27,7 @@ def test_get_elapsed_s_span_none(): """Tests fallback when span is None.""" start_time = 10.0 with mock.patch("time.monotonic", return_value=12.0): - elapsed = _instrumentation._get_elapsed_s(None, start_time) + elapsed = _metrics.get_elapsed_s(None, start_time) assert elapsed == 2.0 # 12 - 10 @@ -35,7 +36,7 @@ def test_get_elapsed_s_span_valid(): mock_span = mock.MagicMock(spec=trace.Span) mock_span.start_time = 1000000000 # 1s in ns mock_span.end_time = 2000000000 # 2s in ns - elapsed = _instrumentation._get_elapsed_s(mock_span, time.monotonic()) + elapsed = _metrics.get_elapsed_s(mock_span, time.monotonic()) assert elapsed == 1.0 # (2 - 1) s @@ -46,7 +47,7 @@ def test_get_elapsed_s_span_missing_start(): mock_span.end_time = 2000000000 start_time = 10.0 with mock.patch("time.monotonic", return_value=12.0): - elapsed = _instrumentation._get_elapsed_s(mock_span, start_time) + elapsed = _metrics.get_elapsed_s(mock_span, start_time) assert elapsed == 2.0 @@ -57,7 +58,7 @@ def test_get_elapsed_s_span_missing_end(): del mock_span.end_time start_time = 10.0 with mock.patch("time.monotonic", return_value=12.0): - elapsed = _instrumentation._get_elapsed_s(mock_span, start_time) + elapsed = _metrics.get_elapsed_s(mock_span, start_time) assert elapsed == 2.0 @@ -68,7 +69,7 @@ def test_get_elapsed_s_span_non_int_start(): mock_span.end_time = 2000000000 start_time = 10.0 with mock.patch("time.monotonic", return_value=12.0): - elapsed = _instrumentation._get_elapsed_s(mock_span, start_time) + elapsed = _metrics.get_elapsed_s(mock_span, start_time) assert elapsed == 2.0 @@ -79,7 +80,7 @@ def test_get_elapsed_s_span_non_int_end(): mock_span.end_time = 2000000000.0 start_time = 10.0 with mock.patch("time.monotonic", return_value=12.0): - elapsed = _instrumentation._get_elapsed_s(mock_span, start_time) + elapsed = _metrics.get_elapsed_s(mock_span, start_time) assert elapsed == 2.0 diff --git a/tests/unittests/telemetry/test_metrics.py b/tests/unittests/telemetry/test_metrics.py index 30778d06a5a..d18f01e8977 100644 --- a/tests/unittests/telemetry/test_metrics.py +++ b/tests/unittests/telemetry/test_metrics.py @@ -27,6 +27,7 @@ def _mock_meter_setup(monkeypatch): """Sets up mock meter and histograms for testing.""" mock_meter = mock.MagicMock() agent_duration_hist = mock.MagicMock(spec=metrics.Histogram) + workflow_duration_hist = mock.MagicMock(spec=metrics.Histogram) tool_duration_hist = mock.MagicMock(spec=metrics.Histogram) request_size_hist = mock.MagicMock(spec=metrics.Histogram) response_size_hist = mock.MagicMock(spec=metrics.Histogram) @@ -35,6 +36,7 @@ def _mock_meter_setup(monkeypatch): client_token_usage_hist = mock.MagicMock(spec=metrics.Histogram) agent_duration_hist.name = "agent_invocation_duration" + workflow_duration_hist.name = "workflow_invocation_duration" tool_duration_hist.name = "tool_execution_duration" request_size_hist.name = "agent_request_size" response_size_hist.name = "agent_response_size" @@ -43,9 +45,11 @@ def _mock_meter_setup(monkeypatch): client_token_usage_hist.name = "client_token_usage" def create_histogram_side_effect(name, **_kwargs): - if name == "gen_ai.agent.invocation.duration": + if name == "gen_ai.invoke_agent.duration": return agent_duration_hist - elif name == "gen_ai.tool.execution.duration": + elif name == "gen_ai.invoke_workflow.duration": + return workflow_duration_hist + elif name == "gen_ai.execute_tool.duration": return tool_duration_hist elif name == "gen_ai.agent.request.size": return request_size_hist @@ -66,6 +70,9 @@ def create_histogram_side_effect(name, **_kwargs): monkeypatch.setattr( _metrics, "_agent_invocation_duration", agent_duration_hist ) + monkeypatch.setattr( + _metrics, "_workflow_invocation_duration", workflow_duration_hist + ) monkeypatch.setattr(_metrics, "_tool_execution_duration", tool_duration_hist) monkeypatch.setattr(_metrics, "_agent_request_size", request_size_hist) monkeypatch.setattr(_metrics, "_agent_response_size", response_size_hist) @@ -78,6 +85,7 @@ def create_histogram_side_effect(name, **_kwargs): return { "meter": mock_meter, "agent_duration": agent_duration_hist, + "workflow_duration": workflow_duration_hist, "tool_duration": tool_duration_hist, "request_size": request_size_hist, "response_size": response_size_hist, @@ -131,6 +139,40 @@ def test_record_agent_invocation_duration_with_error(mock_meter_setup): assert kwargs["attributes"]["error.type"] == "ValueError" +def test_record_workflow_invocation_duration_root(mock_meter_setup): + """Tests record_workflow_invocation_duration omits nested for the root.""" + _metrics.record_workflow_invocation_duration( + workflow_name="my_workflow", + elapsed_s=1.0, + nested=False, + ) + hist = mock_meter_setup["workflow_duration"] + hist.record.assert_called_once() + args, kwargs = hist.record.call_args + assert args[0] == 1.0 + assert kwargs["attributes"] == { + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": "my_workflow", + } + + +def test_record_workflow_invocation_duration_nested_with_error( + mock_meter_setup, +): + """Tests record_workflow_invocation_duration records nested + error.""" + _metrics.record_workflow_invocation_duration( + workflow_name="nested_workflow", + elapsed_s=2.0, + nested=True, + error=ValueError("boom"), + ) + hist = mock_meter_setup["workflow_duration"] + hist.record.assert_called_once() + _, kwargs = hist.record.call_args + assert kwargs["attributes"]["gen_ai.workflow.nested"] is True + assert kwargs["attributes"]["error.type"] == "ValueError" + + def test_record_agent_response_size(mock_meter_setup): """Tests record_agent_response_size records correctly.""" response_text = "response" diff --git a/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py b/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py index 5d51b29e25f..062531ed7c9 100644 --- a/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py +++ b/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py @@ -567,6 +567,7 @@ async def test_init_with_connection_and_custom_auth( tool_instructions=tool_instructions, auth_scheme=oauth2_scheme, auth_credential=auth_credential, + credential_key="test-key", ) mock_integration_client.assert_called_once_with( project, @@ -594,6 +595,7 @@ async def test_init_with_connection_and_custom_auth( assert (await toolset.get_tools())[0]._operation == "EXECUTE_ACTION" assert (await toolset.get_tools())[0]._auth_scheme == oauth2_scheme assert (await toolset.get_tools())[0]._auth_credential == auth_credential + assert (await toolset.get_tools())[0]._credential_key == "test-key" @pytest.mark.asyncio diff --git a/tests/unittests/tools/application_integration_tool/test_integration_connector_tool.py b/tests/unittests/tools/application_integration_tool/test_integration_connector_tool.py index dfea6442127..c2e0ea15592 100644 --- a/tests/unittests/tools/application_integration_tool/test_integration_connector_tool.py +++ b/tests/unittests/tools/application_integration_tool/test_integration_connector_tool.py @@ -23,6 +23,7 @@ from google.adk.tools.application_integration_tool.integration_connector_tool import IntegrationConnectorTool from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool from google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler import AuthPreparationResult +from google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler import ToolAuthHandler from google.genai.types import FunctionDeclaration from google.genai.types import Schema from google.genai.types import Type @@ -96,6 +97,7 @@ def integration_tool_with_auth(mock_rest_api_tool): credentials=HttpCredentials(token="mocked_token"), ), ), + credential_key="test-key", ) @@ -190,8 +192,8 @@ async def test_run_with_auth_async_none_token( "sortByColumns": ["a", "b"], } - with mock.patch( - "google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool.ToolAuthHandler.from_tool_context" + with mock.patch.object( + ToolAuthHandler, "from_tool_context", autospec=True ) as mock_from_tool_context: mock_tool_auth_handler_instance = mock.MagicMock() # Simulate an AuthCredential that would cause _prepare_dynamic_euc to return None @@ -214,6 +216,12 @@ async def test_run_with_auth_async_none_token( result = await integration_tool_with_auth.run_async( args=input_args, tool_context={} ) + mock_from_tool_context.assert_called_once_with( + {}, + None, + integration_tool_with_auth._auth_credential, + credential_key="test-key", + ) mock_rest_api_tool.call.assert_called_once_with( args=expected_call_args, tool_context={} @@ -241,8 +249,8 @@ async def test_run_with_auth_async( "action": "TestAction", } - with mock.patch( - "google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool.ToolAuthHandler.from_tool_context" + with mock.patch.object( + ToolAuthHandler, "from_tool_context", autospec=True ) as mock_from_tool_context: mock_tool_auth_handler_instance = mock.MagicMock() @@ -262,6 +270,12 @@ async def test_run_with_auth_async( result = await integration_tool_with_auth.run_async( args=input_args, tool_context={} ) + mock_from_tool_context.assert_called_once_with( + {}, + None, + integration_tool_with_auth._auth_credential, + credential_key="test-key", + ) mock_rest_api_tool.call.assert_called_once_with( args=expected_call_args, tool_context={} ) diff --git a/tests/unittests/tools/google_api_tool/test_google_api_toolset.py b/tests/unittests/tools/google_api_tool/test_google_api_toolset.py index 7c2ef78a647..b9e9b667fd1 100644 --- a/tests/unittests/tools/google_api_tool/test_google_api_toolset.py +++ b/tests/unittests/tools/google_api_tool/test_google_api_toolset.py @@ -523,3 +523,44 @@ def test_init_with_tool_name_prefix( ) assert tool_set.tool_name_prefix == tool_name_prefix + + @mock.patch( + "google.adk.tools.google_api_tool.google_api_toolset.OpenAPIToolset" + ) + @mock.patch( + "google.adk.tools.google_api_tool.google_api_toolset.GoogleApiToOpenApiConverter" + ) + @mock.patch( + "google.adk.tools.google_api_tool.google_api_toolset.MtlsClientCerts" + ) + @mock.patch( + "google.adk.tools.google_api_tool.google_api_toolset.use_client_cert_effective" + ) + async def test_mtls_cleanup_on_close( + self, + mock_use_client_cert, + mock_mtls_certs_class, + mock_converter_class, + mock_openapi_toolset_class, + ): + """Test that mTLS temp files are cleaned up on close.""" + mock_converter_class.return_value = mock.MagicMock() + mock_openapi_toolset_instance = mock.MagicMock() + mock_openapi_toolset_instance.close = mock.AsyncMock() + mock_openapi_toolset_class.return_value = mock_openapi_toolset_instance + + mock_use_client_cert.return_value = True + mock_mtls_certs_instance = mock.MagicMock() + mock_mtls_certs_instance.get_certs.return_value = ("cert", "key", b"pass") + mock_mtls_certs_class.return_value = mock_mtls_certs_instance + + tool_set = GoogleApiToolset( + api_name=TEST_API_NAME, api_version=TEST_API_VERSION + ) + + assert tool_set._httpx_client_factory is not None + + await tool_set.close() + + mock_openapi_toolset_instance.close.assert_called_once() + mock_mtls_certs_instance.close.assert_called_once() diff --git a/tests/unittests/tools/google_api_tool/test_googleapi_to_openapi_converter.py b/tests/unittests/tools/google_api_tool/test_googleapi_to_openapi_converter.py index 9242b0e133f..e6320e45f91 100644 --- a/tests/unittests/tools/google_api_tool/test_googleapi_to_openapi_converter.py +++ b/tests/unittests/tools/google_api_tool/test_googleapi_to_openapi_converter.py @@ -197,6 +197,14 @@ def calendar_api_spec(): } +@pytest.fixture(autouse=True) +def disable_mtls_by_default(monkeypatch): + monkeypatch.setattr( + "google.auth.transport.mtls.should_use_client_cert", + lambda: False, + ) + + @pytest.fixture def converter(): """Fixture that provides a basic converter instance.""" @@ -279,7 +287,50 @@ def test_fetch_google_api_spec_with_discovery_url( assert converter._google_api_spec == calendar_api_spec mock_build.assert_called_once_with( - "calendar", "v3", discoveryServiceUrl=discovery_url + "calendar", "v3", discoveryServiceUrl=discovery_url, http=None + ) + + def test_fetch_google_api_spec_with_mtls( + self, monkeypatch, mock_api_resource, calendar_api_spec + ): + """Test fetching Google API specification with mTLS enabled.""" + mock_build = MagicMock(return_value=mock_api_resource) + monkeypatch.setattr( + "google.adk.tools.google_api_tool.googleapi_to_openapi_converter.build", + mock_build, + ) + + # Enable mTLS + monkeypatch.setattr( + "google.auth.transport.mtls.should_use_client_cert", + lambda: True, + ) + monkeypatch.setattr( + "google.auth.transport.mtls.has_default_client_cert_source", + lambda: True, + ) + + mock_cert_source = MagicMock( + return_value=("/path/to/cert", "/path/to/key", b"passphrase") + ) + monkeypatch.setattr( + "google.auth.transport.mtls.default_client_encrypted_cert_source", + lambda c, k: mock_cert_source, + ) + + converter = GoogleApiToOpenApiConverter("calendar", "v3") + converter.fetch_google_api_spec() + + assert converter._google_api_spec == calendar_api_spec + + # Verify build was called with the http parameter set and mtls url + mock_build.assert_called_once() + _, kwargs = mock_build.call_args + assert "http" in kwargs + assert kwargs["http"] is not None + assert ( + kwargs["discoveryServiceUrl"] + == "https://www.mtls.googleapis.com/discovery/v1/apis/{api}/{apiVersion}/rest" ) def test_fetch_google_api_spec_error(self, monkeypatch, converter): diff --git a/tests/unittests/utils/test_mtls_utils.py b/tests/unittests/utils/test_mtls_utils.py index 9a2572a437c..252721d0544 100644 --- a/tests/unittests/utils/test_mtls_utils.py +++ b/tests/unittests/utils/test_mtls_utils.py @@ -20,6 +20,7 @@ from google.adk.utils import _mtls_utils from google.auth import exceptions as ga_exceptions +from google.auth.transport import mtls import pytest _DEFAULT_TEMPLATE = "service.{location}.rep.googleapis.com" @@ -30,7 +31,7 @@ class TestMtlsUtils: """Tests for _mtls_utils functions.""" - @patch("google.auth.transport.mtls.should_use_client_cert") + @patch.object(mtls, "should_use_client_cert", autospec=True) def test_use_client_cert_effective_with_mtls_cert_true( self, mock_should_use_client_cert ): @@ -38,7 +39,7 @@ def test_use_client_cert_effective_with_mtls_cert_true( assert _mtls_utils.use_client_cert_effective() is True mock_should_use_client_cert.assert_called_once() - @patch("google.auth.transport.mtls.should_use_client_cert") + @patch.object(mtls, "should_use_client_cert", autospec=True) def test_use_client_cert_effective_with_mtls_cert_false( self, mock_should_use_client_cert ): @@ -46,7 +47,7 @@ def test_use_client_cert_effective_with_mtls_cert_false( assert _mtls_utils.use_client_cert_effective() is False mock_should_use_client_cert.assert_called_once() - @patch("google.auth.transport.mtls.should_use_client_cert") + @patch.object(mtls, "should_use_client_cert", autospec=True) @patch.dict("os.environ", {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}) def test_use_client_cert_effective_fallback_true( self, mock_should_use_client_cert @@ -54,7 +55,7 @@ def test_use_client_cert_effective_fallback_true( mock_should_use_client_cert.side_effect = AttributeError assert _mtls_utils.use_client_cert_effective() is True - @patch("google.auth.transport.mtls.should_use_client_cert") + @patch.object(mtls, "should_use_client_cert", autospec=True) @patch.dict("os.environ", {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "false"}) def test_use_client_cert_effective_fallback_false( self, mock_should_use_client_cert @@ -62,7 +63,7 @@ def test_use_client_cert_effective_fallback_false( mock_should_use_client_cert.side_effect = AttributeError assert _mtls_utils.use_client_cert_effective() is False - @patch("google.auth.transport.mtls.should_use_client_cert") + @patch.object(mtls, "should_use_client_cert", autospec=True) @patch.dict("os.environ", {}, clear=True) def test_use_client_cert_effective_fallback_default_false( self, mock_should_use_client_cert @@ -70,7 +71,7 @@ def test_use_client_cert_effective_fallback_default_false( mock_should_use_client_cert.side_effect = AttributeError assert _mtls_utils.use_client_cert_effective() is False - @patch("google.adk.utils._mtls_utils.use_client_cert_effective") + @patch.object(_mtls_utils, "use_client_cert_effective", autospec=True) @patch.dict("os.environ", {"GOOGLE_API_USE_MTLS_ENDPOINT": "always"}) def test_get_api_endpoint_always(self, mock_use_client_cert): endpoint = _mtls_utils.get_api_endpoint( @@ -79,7 +80,7 @@ def test_get_api_endpoint_always(self, mock_use_client_cert): assert endpoint == _MTLS_TEMPLATE.format(location=_LOCATION) mock_use_client_cert.assert_not_called() - @patch("google.adk.utils._mtls_utils.use_client_cert_effective") + @patch.object(_mtls_utils, "use_client_cert_effective", autospec=True) @patch.dict("os.environ", {"GOOGLE_API_USE_MTLS_ENDPOINT": "never"}) def test_get_api_endpoint_never(self, mock_use_client_cert): endpoint = _mtls_utils.get_api_endpoint( @@ -88,7 +89,7 @@ def test_get_api_endpoint_never(self, mock_use_client_cert): assert endpoint == _DEFAULT_TEMPLATE.format(location=_LOCATION) mock_use_client_cert.assert_not_called() - @patch("google.adk.utils._mtls_utils.use_client_cert_effective") + @patch.object(_mtls_utils, "use_client_cert_effective", autospec=True) @patch.dict("os.environ", {"GOOGLE_API_USE_MTLS_ENDPOINT": "auto"}) def test_get_api_endpoint_auto_with_cert(self, mock_use_client_cert): mock_use_client_cert.return_value = True @@ -98,7 +99,7 @@ def test_get_api_endpoint_auto_with_cert(self, mock_use_client_cert): assert endpoint == _MTLS_TEMPLATE.format(location=_LOCATION) mock_use_client_cert.assert_called_once() - @patch("google.adk.utils._mtls_utils.use_client_cert_effective") + @patch.object(_mtls_utils, "use_client_cert_effective", autospec=True) @patch.dict("os.environ", {"GOOGLE_API_USE_MTLS_ENDPOINT": "auto"}) def test_get_api_endpoint_auto_without_cert(self, mock_use_client_cert): mock_use_client_cert.return_value = False @@ -108,7 +109,7 @@ def test_get_api_endpoint_auto_without_cert(self, mock_use_client_cert): assert endpoint == _DEFAULT_TEMPLATE.format(location=_LOCATION) mock_use_client_cert.assert_called_once() - @patch("google.adk.utils._mtls_utils.use_client_cert_effective") + @patch.object(_mtls_utils, "use_client_cert_effective", autospec=True) @patch.dict("os.environ", {"GOOGLE_API_USE_MTLS_ENDPOINT": "invalid_value"}) def test_get_api_endpoint_invalid_fallback_to_auto( self, mock_use_client_cert @@ -217,3 +218,86 @@ def test_configure_session_for_mtls_cert_error_falls_back( assert result is False session.mount.assert_not_called() + + +class TestMtlsClientCerts: + """Tests for MtlsClientCerts.""" + + @patch.object(mtls, "has_default_client_cert_source", autospec=True) + def test_get_certs_no_default_source(self, mock_has_cert): + mock_has_cert.return_value = False + certs = _mtls_utils.MtlsClientCerts() + cert_path, key_path, passphrase = certs.get_certs() + assert cert_path is None + assert key_path is None + assert passphrase is None + mock_has_cert.assert_called_once() + + @patch.object(mtls, "has_default_client_cert_source", autospec=True) + @patch.object(mtls, "default_client_encrypted_cert_source", autospec=True) + def test_get_certs_with_default_source( + self, mock_encrypted_source, mock_has_cert + ): + mock_has_cert.return_value = True + + mock_cert_source = MagicMock() + mock_cert_source.return_value = (None, None, b"test_passphrase") + mock_encrypted_source.return_value = mock_cert_source + + certs = _mtls_utils.MtlsClientCerts() + cert_path, key_path, passphrase = certs.get_certs() + + assert cert_path is not None + assert key_path is not None + assert passphrase == b"test_passphrase" + + assert os.path.exists(certs._tempdir.name) + assert cert_path.startswith(certs._tempdir.name) + assert key_path.startswith(certs._tempdir.name) + + # Getting certs again should return cached values without calling mtls again + cert_path2, key_path2, passphrase2 = certs.get_certs() + assert cert_path2 == cert_path + assert key_path2 == key_path + assert passphrase2 == passphrase + mock_has_cert.assert_called_once() + mock_encrypted_source.assert_called_once() + + @patch.object(mtls, "has_default_client_cert_source", autospec=True) + @patch.object(mtls, "default_client_encrypted_cert_source", autospec=True) + def test_get_certs_extraction_failure( + self, mock_encrypted_source, mock_has_cert + ): + mock_has_cert.return_value = True + mock_encrypted_source.side_effect = Exception("extraction failed") + + certs = _mtls_utils.MtlsClientCerts() + with pytest.raises( + RuntimeError, match="Failed to extract default client certificates" + ): + certs.get_certs() + + assert certs._tempdir is None + + @patch.object(mtls, "has_default_client_cert_source", autospec=True) + @patch.object(mtls, "default_client_encrypted_cert_source", autospec=True) + def test_close_cleans_up_tempdir(self, mock_encrypted_source, mock_has_cert): + mock_has_cert.return_value = True + mock_cert_source = MagicMock() + mock_cert_source.return_value = (None, None, b"test_passphrase") + mock_encrypted_source.return_value = mock_cert_source + + certs = _mtls_utils.MtlsClientCerts() + certs.get_certs() + + tempdir_name = certs._tempdir.name + assert os.path.exists(tempdir_name) + + certs.close() + + assert not os.path.exists(tempdir_name) + assert certs._tempdir is None + assert certs.cert_path is None + assert certs.key_path is None + assert certs.passphrase is None + assert not certs._initialized