From 8930d9b19338873d091215a1e56f99ce6fd2060b Mon Sep 17 00:00:00 2001 From: Max Ind Date: Wed, 29 Jul 2026 09:17:09 -0700 Subject: [PATCH] feat(telemetry): add request-driven metric export for Agent Engine On the Vertex AI Agent Runtime CPU is throttled the instant a request finishes, so a background periodic metric exporter is starved between requests and drops data. This adds a request-driven metric reader (and the span processor + middleware that drive it) that collects and exports metrics on the request path, where CPU is guaranteed, without adding latency to any request. The default GCP metric exporter is also switched to a raw OTLP push exporter over telemetry.googleapis.com. Co-authored-by: Max Ind PiperOrigin-RevId: 955921308 --- src/google/adk/cli/api_server.py | 3 +- src/google/adk/cli/fast_api.py | 3 + src/google/adk/telemetry/_agent_engine.py | 182 ++++++- .../_agent_engine_metric_exporter.py | 489 ++++++++++++++++++ src/google/adk/telemetry/google_cloud.py | 161 ++++-- src/google/adk/telemetry/setup.py | 2 + .../unittests/telemetry/test_agent_engine.py | 345 +++++++++++- .../test_agent_engine_metric_exporter.py | 287 ++++++++++ .../unittests/telemetry/test_google_cloud.py | 235 ++++++++- 9 files changed, 1633 insertions(+), 74 deletions(-) create mode 100644 src/google/adk/telemetry/_agent_engine_metric_exporter.py create mode 100644 tests/unittests/telemetry/test_agent_engine_metric_exporter.py diff --git a/src/google/adk/cli/api_server.py b/src/google/adk/cli/api_server.py index 94aac65ae4e..54708b28111 100644 --- a/src/google/adk/cli/api_server.py +++ b/src/google/adk/cli/api_server.py @@ -587,8 +587,7 @@ def _setup_gcp_telemetry( # TODO - use trace_to_cloud here as well once otel_to_cloud is no # longer experimental. enable_cloud_tracing=True, - # TODO - re-enable metrics once errors during shutdown are fixed. - enable_cloud_metrics=False, + enable_cloud_metrics=True, enable_cloud_logging=True, google_auth=(credentials, project_id), ) diff --git a/src/google/adk/cli/fast_api.py b/src/google/adk/cli/fast_api.py index f096717dc29..ad033df7951 100644 --- a/src/google/adk/cli/fast_api.py +++ b/src/google/adk/cli/fast_api.py @@ -52,6 +52,7 @@ from ..auth.credential_service.in_memory_credential_service import InMemoryCredentialService from ..runners import Runner from ..telemetry._agent_engine import get_propagated_context +from ..telemetry._agent_engine import maybe_install_request_metrics_middleware from ..telemetry._agent_engine import TopSpanProcessor from .api_server import ApiServer from .cli_deploy import _AGENT_ENGINE_CLASS_METHODS @@ -678,6 +679,8 @@ async def _a2a_lifespan(app_instance: FastAPI): **extra_fast_api_args, ) + maybe_install_request_metrics_middleware(app, otel_to_cloud=otel_to_cloud) + # --- Builder endpoints (agent editor UI) --- _register_builder_endpoints(app, web, agents_dir) diff --git a/src/google/adk/telemetry/_agent_engine.py b/src/google/adk/telemetry/_agent_engine.py index 99e6b55d501..97afc18819e 100644 --- a/src/google/adk/telemetry/_agent_engine.py +++ b/src/google/adk/telemetry/_agent_engine.py @@ -14,14 +14,40 @@ from __future__ import annotations -from typing import Mapping -from typing import Optional +import asyncio +import contextlib +import functools +import logging +import os +from types import ModuleType +from typing import cast +from typing import TYPE_CHECKING -import fastapi from opentelemetry import baggage from opentelemetry import context +from opentelemetry import metrics from opentelemetry.sdk import trace +from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.trace.propagation import tracecontext +from starlette.middleware.base import BaseHTTPMiddleware + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator + from collections.abc import AsyncIterator + from typing import Mapping + from typing import Optional + + import fastapi + from starlette.middleware.base import DispatchFunction + from starlette.middleware.base import RequestResponseEndpoint + from starlette.requests import Request + from starlette.responses import Response + from starlette.responses import StreamingResponse + + from ._agent_engine_metric_exporter import _RequestDrivenMetricReader + from ._agent_engine_metric_exporter import MetricsState + +logger = logging.getLogger("google_adk." + __name__) _GOOGLE_AE_TRACEPARENT_HEADER = "Google-Agent-Engine-Traceparent" _TRACEPARENT_BAGGAGE_KEY = "traceparent" @@ -110,3 +136,153 @@ def _is_top_span( except ValueError: return False return span.parent.span_id == parent_span_id + + +def _metrics_flushing_dispatch( + reader: _RequestDrivenMetricReader, +) -> DispatchFunction: + """Returns the dispatch that drives `reader` from the request lifecycle. + + Collection has to happen while a request is in flight: see the module + docstring of ``_agent_engine_metric_exporter`` for why. Traces and logs are + not flushed here -- on Agent Engine the ``AdkApp`` in + ``vertexai.agent_engines`` already force-flushes them per request. + """ + + async def dispatch( + request: Request, call_next: RequestResponseEndpoint + ) -> Response: + # Never let a metrics failure break the request it rides on. + try: + if reader.note_request_start(): + _ = reader.submit_collect() # point 2 -- fire-and-forget, no await. + except Exception: # pylint: disable=broad-exception-caught + logger.exception("Metrics request-start hook failed") + # BaseHTTPMiddleware always hands back a streaming response, so it exposes + # body_iterator (the base Response type does not). + try: + response = cast("StreamingResponse", await call_next(request)) + except BaseException: + # call_next failed: the draining body iterator below is never installed, + # so balance note_request_start here or _in_flight leaks for good. + await _drain_metrics(reader) + raise + # BaseHTTPMiddleware's body_iterator is an async generator (supports + # aclose); the public type is only an AsyncIterable, so narrow it here. + original_iterator = cast( + "AsyncGenerator[bytes, None]", response.body_iterator + ) + + async def iterate_then_flush() -> AsyncIterator[bytes]: + try: + async with contextlib.aclosing(original_iterator) as body: + async for chunk in body: + yield chunk + finally: + # Drain metrics after the body streams (while the request still has + # CPU) and before the connection closes, so the export completes + # without adding latency to the response. + await _drain_metrics(reader) + + response.body_iterator = iterate_then_flush() + return response + + return dispatch + + +async def _drain_metrics(reader: _RequestDrivenMetricReader) -> None: + """Drain collect on request end: awaited so the export completes before close.""" + try: + if reader.note_request_end(): + future = reader.submit_collect() + if future is not None: + await asyncio.wrap_future(future) + except Exception: # pylint: disable=broad-exception-caught + logger.exception("Failed to flush metrics on request end") + + +def telemetry_user_agent_headers() -> dict[str, str] | None: + """Returns the Vertex Agent Engine User-Agent header, if telemetry is on.""" + if not os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY"): + return None + from google.cloud.aiplatform import version as aip_version + + otlp_http_version: ModuleType | None + try: + from opentelemetry.exporter.otlp.proto.http import version as otlp_http_version + except (ImportError, AttributeError): + otlp_http_version = None + + user_agent = f"Vertex-Agent-Engine/{aip_version.__version__}" + if otlp_http_version: + user_agent += f" OTel-OTLP-Exporter-Python/{otlp_http_version.__version__}" + return {"User-Agent": user_agent} + + +@functools.cache +def _get_agent_engine_metrics_setup() -> MetricsState | None: + """Builds the request-driven metric state on Agent Engine, memoized. + + Returns the reader plus the span processor that drives it (to wire into the + OTel providers and later drain from the request path), or None when: + + 1. a ``MeterProvider`` is already installed by something else -- our reader + would not land on the active provider, so we defer to that setup; or + 2. we are not running on Agent Engine + (``GOOGLE_CLOUD_AGENT_ENGINE_ID`` unset); or + 3. the GCP metric exporter is unavailable / setup fails. + + Cached so the "already installed" check is evaluated once -- before ADK sets + its own ``MeterProvider`` in ``maybe_set_otel_providers`` -- and the same + handles are returned to ``get_gcp_exporters`` (which wires them onto the + providers) and to the middleware install below. + """ + if not os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID"): + return None + if isinstance(metrics.get_meter_provider(), MeterProvider): + logger.warning( + "A MeterProvider is already installed; skipping request-driven metric" + " export. On Agent Engine's request-billed runtime metrics may be" + " dropped between requests." + ) + return None + try: + from ._agent_engine_metric_exporter import build_request_driven_metrics + from .google_cloud import _get_gcp_otlp_metric_exporter + + exporter = _get_gcp_otlp_metric_exporter() + if exporter is None: + return None + return build_request_driven_metrics(exporter) + except Exception: # pylint: disable=broad-exception-caught + logger.warning( + "Failed to set up request-driven metric export on Agent Engine.", + exc_info=True, + ) + return None + + +def maybe_install_request_metrics_middleware( + app: fastapi.FastAPI, *, otel_to_cloud: bool +) -> None: + """Installs the request-path metric flushing middleware, if applicable. + + On Agent Engine's request-billed runtime CPU is throttled the instant a + request ends, starving background metric export. The GCP exporter setup + builds a request-driven metric reader there (wired onto the MeterProvider); + drive it from the request path here. No-op off Agent Engine. + + Args: + app: The app to install the middleware on. + otel_to_cloud: Whether to setup telemetry export to GCP. + """ + if not otel_to_cloud: + return + + metrics_state = _get_agent_engine_metrics_setup() + if metrics_state is None: + return + app.add_middleware( + BaseHTTPMiddleware, + dispatch=_metrics_flushing_dispatch(metrics_state.reader), + ) diff --git a/src/google/adk/telemetry/_agent_engine_metric_exporter.py b/src/google/adk/telemetry/_agent_engine_metric_exporter.py new file mode 100644 index 00000000000..1577d6d358e --- /dev/null +++ b/src/google/adk/telemetry/_agent_engine_metric_exporter.py @@ -0,0 +1,489 @@ +# 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. + +r"""Request-driven, sleepless metric export. + +A bespoke OpenTelemetry metric reader (plus the span processor that drives it) +that exports metrics from an agent running on the Vertex AI Agent Runtime +(request-based billing) *without adding latency to any request* and dropping +data only in a rare, well-defined tail case (see below). + +The FastAPI/Starlette middleware that also drives this reader from the request +lifecycle -- and flushes traces and logs on the same request path -- lives in +``google.adk.telemetry._agent_engine``. + +The problem +----------- +The Agent Runtime throttles CPU the instant a request finishes (request-based +billing). A normal metric pipeline exports on a background timer thread; between +requests that thread gets no CPU, so its periodic export is starved and metric +points are dropped. + +The residual loss case +---------------------- +This is not fully lossless. A request shorter than the floor that drains right +after a collect -- and is the *last* request before the process goes idle -- +loses its points: a collect now is muted by the floor (I2), and the next +collect never comes. Illustrated at the end of the timeline section. + +Three constraints shape the solution: + +- I1 -- export only while serving. A collect+export must run while a request is + in flight (the only time CPU is guaranteed). +- I2 -- never collect more often than the floor. Two collects closer than the + floor (5s default) are rejected by the collection path. +- I3 -- never collect too rarely. A single export carries at most ~200 points, + so collect at least once per configured period. + +Design +------ +There is no background ticker and no ``time.sleep``. The reader subclasses the +base ``MetricReader`` directly, so no daemon thread is started; ``collect()`` is +invoked only from the request lifecycle (middleware) and from +``generate_content`` span starts (span processor). The configured export period +is demoted from a hard schedule to a guidepost grid -- hint times used to decide +whether an event-driven collect is warranted. Every collect runs on the request +path, so it always has CPU (I1); the floor is enforced by skipping a collect +that would land too soon (I2); guideposts force a collect during sustained load +so points can't pile up (I3). + +Timeline legend:: + + [====] a request being served (holds CPU until its connection closes) + C a collect: drain the SDK aggregation + export, run on the request + path + P a guidepost: a periodic "would-be" collect time -- a hint, never a + real scheduled call + x a collect skipped because it would violate the floor (min spacing) + FLOOR minimum spacing between two collects (default 3s) + PERIOD spacing of the guidepost grid (default 60s) + +Baseline -- a request collects when it drains to zero:: + + req1 [==========] + +- in_flight 1->0 -> C + +Overlap batches into one collect:: + + req1 [======] + req2 [========] + req3 [==========] + +- in_flight -> 0 only here -> C (one collect) + +Point 2 -- under continuous overlap, a guidepost fires at the next start:: + + req1 [====] + req2 [=======] + req3 [==========] + req4 [==========] + P +- guidepost crossed while req3 in flight; req4 arrives -> C + +Point 3 -- a guidepost too close to the last collect is muted (grid += PERIOD). + +Point 4 -- a lone very long request uses its own inference spans: once +1.5xPERIOD has elapsed since the last collect, the next ``generate_content`` +span start collects. + +Loss case -- the last request is too short to collect, and nothing follows:: + + req1 [==========] + +- in_flight 1->0 -> C + req2 [=] (drains < FLOOR after C) + +- in_flight 1->0, but a collect here is muted by the + floor (I2); its points wait for the next collect + ... process goes idle, no later request -> C never + comes, req2's points are lost + +Threading model +--------------- +The blocking export never runs on the event loop. The reader owns a +single-worker ``ThreadPoolExecutor``; ``submit_collect()`` runs the collect on +it. The drain collect is the only awaited path (via ``asyncio.wrap_future``), +after the response body has streamed, so the export completes before the +connection closes. Start/``generate_content`` collects are fire-and-forget. All +shared state is guarded by one ``threading.Lock``. + +Besides the tail loss case above, metrics-on / tracing-off leaves point 4 +inactive; a lone ultra-long request under that config can accumulate. Both are +rare and accepted. +""" + +# This module deliberately mirrors OTel SDK internals (the private aggregation / +# temporality preferences, the instrumentation-suppression key, and the +# _receive_metrics contract), so private access is expected throughout. +# pyright: reportPrivateUsage=false +from __future__ import annotations + +from collections.abc import Callable +import concurrent.futures +import dataclasses +import logging +import os +import threading +import time + +from opentelemetry import context as otel_context +from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY +from opentelemetry.context.context import Context +from opentelemetry.sdk.environment_variables import OTEL_METRIC_EXPORT_INTERVAL +from opentelemetry.sdk.environment_variables import OTEL_METRIC_EXPORT_TIMEOUT +from opentelemetry.sdk.metrics import export as metrics_export +from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.sdk.trace import Span +from opentelemetry.sdk.trace import SpanProcessor + +logger = logging.getLogger("google_adk." + __name__) + +# Env-var name for the hard floor on collect spacing (I2), in milliseconds. +GOOGLE_CLOUD_AGENT_ENGINE_METRICS_COLLECTION_INTERVAL_FLOOR_MS = ( + "GOOGLE_CLOUD_AGENT_ENGINE_METRICS_COLLECTION_INTERVAL_FLOOR_MS" +) + +# The single minimum spacing between two metric exports to +# telemetry.googleapis.com, shared by every ADK metric reader: the floor (I2) +# for the request-driven reader here, and the export interval of the periodic +# reader in `google_cloud`. +# +# The backend currently accepts points sent more frequently than this, but that +# is only there to absorb drift (a reader firing slightly early); it is not a +# supported rate and must not be relied on. Exporting faster than this interval +# risks points being rejected or throttled -- keep new readers at or above it. +MIN_EXPORT_INTERVAL_MS = 5000.0 +# Semantic-convention attribute carrying the GenAI operation name. +_GEN_AI_OPERATION_NAME = "gen_ai.operation.name" + + +def _env_float(name: str, default: float) -> float: + """Reads a float env var, falling back to a default on missing/invalid.""" + raw = os.environ.get(name) + if raw is None: + return default + try: + return float(raw) + except ValueError: + logger.warning( + "Found invalid value for %s=%r, using default %s", name, raw, default + ) + return default + + +def _floor_seconds() -> float: + """Returns the min collect spacing in seconds (from env, default 5.0s).""" + return ( + _env_float( + GOOGLE_CLOUD_AGENT_ENGINE_METRICS_COLLECTION_INTERVAL_FLOOR_MS, + MIN_EXPORT_INTERVAL_MS, + ) + / 1000.0 + ) + + +class _RequestDrivenMetricReader(metrics_export.MetricReader): + """A `MetricReader` whose collects are driven by the request lifecycle.""" + + def __init__( + self, + exporter: metrics_export.MetricExporter, + *, + export_interval_millis: float | None = None, + export_timeout_millis: float | None = None, + floor_millis: float | None = None, + now: Callable[[], float] = time.monotonic, + ): + # Defer temporality/aggregation to the wrapped exporter, exactly as the + # SDK's PeriodicExportingMetricReader does. + super().__init__( + preferred_temporality=exporter._preferred_temporality, # pylint: disable=protected-access + preferred_aggregation=exporter._preferred_aggregation, # pylint: disable=protected-access + ) + self._exporter: metrics_export.MetricExporter = exporter + # Held whenever calling the wrapped exporter, matching the SDK's contract + # that MetricExporter.export() is never called concurrently. + self._export_lock: threading.Lock = threading.Lock() + self._now: Callable[[], float] = now + + if export_interval_millis is None: + export_interval_millis = _env_float(OTEL_METRIC_EXPORT_INTERVAL, 60000.0) + if export_timeout_millis is None: + export_timeout_millis = _env_float(OTEL_METRIC_EXPORT_TIMEOUT, 30000.0) + self._export_timeout_millis: float = export_timeout_millis + self._period_s: float = export_interval_millis / 1000.0 + self._floor_s: float = ( + (floor_millis / 1000.0) + if floor_millis is not None + else _floor_seconds() + ) + + # All fields below are guarded by _lock. + self._lock: threading.Lock = threading.Lock() + self._in_flight: int = 0 + self._last_collect: float | None = None + # Start of the current busy period (stamped when in-flight goes 0 -> 1). The + # point-4 "overdue" reference for the current stretch of activity, so a + # collect from a long-past busy period can't make a short request look + # overdue (see _overdue_15). + self._busy_start: float | None = None + self._collecting: bool = False + self._next_due: float = self._now() + self._period_s + self._shutdown: bool = False + + # One worker => collects are serialized (the >= floor timestamp guarantee + # relies on this). + self._executor: concurrent.futures.ThreadPoolExecutor = ( + concurrent.futures.ThreadPoolExecutor( + max_workers=1, thread_name_prefix="metrics-collect" + ) + ) + # Keep references to in-flight collect futures so they aren't GC'd. + self._inflight_collects: set[concurrent.futures.Future[None]] = set() + + # --- Decision helpers (all called under _lock). ------------------------- + + def _due(self, now: float) -> bool: + return now >= self._next_due + + def _floor_ok(self, now: float) -> bool: + return ( + self._last_collect is None + or (now - self._last_collect) >= self._floor_s + ) + + def _overdue_15(self, now: float) -> bool: + # Point 4 keeps a lone long-running request from piling up between collects. + # "Overdue" is measured over the *current* busy period: from the last + # collect in it, or -- if none yet -- from when the period began + # (_busy_start). A collect from an earlier busy period is ignored, so a + # short request arriving after a long idle gap does not look overdue; and a + # fresh period is not overdue until 1.5*PERIOD into it. This stops point 4 + # firing on a short request's inference span, where it would stamp the floor + # and then mute the request-end drain that carries the request's points. + if self._busy_start is None: + return False + ref = self._busy_start + if self._last_collect is not None: + ref = max(ref, self._last_collect) + return (now - ref) >= 1.5 * self._period_s + + def _arm(self, now: float) -> None: + """Commits to a collect: marks it in flight and advances the guidepost grid. + + Does NOT stamp _last_collect -- that is done by the worker at the actual + collect, so I2 constrains real collect spacing rather than decision spacing. + + Args: + now: The current monotonic time, in seconds. + """ + self._collecting = True + # Advance the guidepost grid to the first tick strictly after `now`. + if now >= self._next_due: + missed = (now - self._next_due) // self._period_s + 1 + self._next_due += missed * self._period_s + + # --- Hooks: fast, synchronous, return "should I collect now?". ---------- + + def note_request_start(self) -> bool: + """Middleware, on request enter. _in_flight is always maintained.""" + with self._lock: + now = self._now() + overlap = self._in_flight >= 1 + if not overlap: + self._busy_start = now # a fresh busy period; reset the point-4 ref. + self._in_flight += 1 + if self._collecting: + return False + if overlap and self._due(now): + if self._floor_ok(now): # point 2 -- collect at the just-started req. + self._arm(now) + return True + # point 3 -- guidepost too close to last collect -> mute it. + self._next_due += self._period_s + return False + + def note_request_end(self) -> bool: + """Middleware, after the response body is fully sent.""" + with self._lock: + now = self._now() + self._in_flight -= 1 + if self._in_flight < 0: + self._in_flight = 0 + if self._in_flight == 0 and not self._collecting and self._floor_ok(now): + self._arm(now) # baseline force-flush (not gated on a guidepost). + return True + return False + + def note_generate_content_start(self) -> bool: + """Span processor, on a generate_content span start (point 4).""" + with self._lock: + now = self._now() + if ( + not self._collecting + and self._in_flight >= 1 + and self._overdue_15(now) + and self._floor_ok(now) + ): + self._arm(now) + return True + return False + + # --- Collect execution. ------------------------------------------------- + + def collect_now(self) -> None: + """Runs a committed collect (on the single-worker executor).""" + with self._lock: + self._last_collect = self._now() # actual collect time. + try: + self.collect(timeout_millis=self._export_timeout_millis) + except Exception: # pylint: disable=broad-exception-caught + # Runs on the executor; a fire-and-forget collect has no one to observe + # its Future, so swallow-and-log rather than let the failure vanish. + logger.exception("Exception during request-driven metric collect") + finally: + with self._lock: + self._collecting = False + + def submit_collect(self) -> concurrent.futures.Future[None] | None: + """Schedules a collect on the executor; returns its Future (or None). + + Returns None (and clears the _collecting guard) if the work can't be + scheduled, e.g. during shutdown, so the guard can't wedge shut. + + Returns: + The scheduled collect's Future, or None if it could not be scheduled. + """ + with self._lock: + shutting_down = self._shutdown + if shutting_down: + with self._lock: + self._collecting = False + return None + try: + future = self._executor.submit(self.collect_now) + except RuntimeError: + with self._lock: + self._collecting = False + return None + self._inflight_collects.add(future) + future.add_done_callback(self._inflight_collects.discard) + return future + + def _receive_metrics( + self, + metrics_data: metrics_export.MetricsData, + timeout_millis: float = 10_000, + **kwargs: object, + ) -> None: + del kwargs # unused + token = otel_context.attach( + otel_context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) + ) + try: + with self._export_lock: + _ = self._exporter.export(metrics_data, timeout_millis=timeout_millis) + except Exception: # pylint: disable=broad-exception-caught + logger.exception("Exception while exporting metrics") + finally: + otel_context.detach(token) + + def shutdown(self, timeout_millis: float = 30_000, **kwargs: object) -> None: + with self._lock: + if self._shutdown: + return + self._shutdown = True + # Drain any in-flight collects, then a best-effort final collect. + self._executor.shutdown(wait=True) + try: + self.collect(timeout_millis=self._export_timeout_millis) + except Exception: # pylint: disable=broad-exception-caught + logger.exception("Exception during final metric collect on shutdown") + self._exporter.shutdown(timeout_millis=timeout_millis, **kwargs) + + +@dataclasses.dataclass(frozen=True) +class MetricsState: + """The metric-export handles the app wires up. + + The middleware that drives the reader lives separately, in + ``google.adk.telemetry._agent_engine``; build it from ``reader``. + """ + + reader: _RequestDrivenMetricReader # installed on the MeterProvider + span_processor: SpanProcessor # installed on the TracerProvider + + +def build_request_driven_metrics( + exporter: metrics_export.MetricExporter, +) -> MetricsState: + """Builds the request-driven reader and the span processor that drives it. + + Unlike a plain periodic reader, the returned reader collects only when driven + from the request lifecycle. This does NOT set any global provider: the caller + installs ``reader`` on a ``MeterProvider`` and ``span_processor`` on the + ``TracerProvider`` (see ``google.adk.telemetry.google_cloud``), and drives the reader + from the request path via ``google.adk.telemetry._agent_engine``. + + The export timeout comes from the OTel env the SDK reads + (``OTEL_METRIC_EXPORT_TIMEOUT`` default 30000); the collect floor comes from + ``GOOGLE_CLOUD_AGENT_ENGINE_METRICS_COLLECTION_INTERVAL_FLOOR_MS`` (default + ``MIN_EXPORT_INTERVAL_MS``). Each falls back to its default on a missing/invalid value. + + Args: + exporter: The metric exporter the reader drains into on each collect. + + Returns: + The reader and the span processor that drive request-based metric export. + """ + reader = _RequestDrivenMetricReader(exporter) + return MetricsState( + reader=reader, + span_processor=_metrics_flushing_span_processor(reader), + ) + + +def _metrics_flushing_span_processor( + reader: _RequestDrivenMetricReader, + *, + operation: str = "generate_content", +) -> SpanProcessor: + """Returns a span processor that collects on `generate_content` span starts.""" + + class _MetricsFlushingSpanProcessor(SpanProcessor): + """Fires a fire-and-forget collect on each matching span start (point 4).""" + + def on_start( + self, span: Span, parent_context: Context | None = None + ) -> None: + del parent_context # unused + # on_start runs inside span creation (ADK's inference path); a metrics + # failure must not break the span it observes. + try: + attributes = span.attributes or {} + name = span.name or "" + if attributes.get( + _GEN_AI_OPERATION_NAME + ) == operation or name.startswith(operation): + if reader.note_generate_content_start(): + _ = reader.submit_collect() # fire-and-forget. + except Exception: # pylint: disable=broad-exception-caught + logger.exception("Metrics span-start hook failed") + + def on_end(self, span: ReadableSpan) -> None: + del span # unused + + def shutdown(self) -> None: + pass + + def force_flush(self, timeout_millis: int = 30000) -> bool: + return True + + return _MetricsFlushingSpanProcessor() diff --git a/src/google/adk/telemetry/google_cloud.py b/src/google/adk/telemetry/google_cloud.py index a0c8f7db398..aa3f6e4895f 100644 --- a/src/google/adk/telemetry/google_cloud.py +++ b/src/google/adk/telemetry/google_cloud.py @@ -36,10 +36,15 @@ from opentelemetry.sdk.trace import SpanProcessor from opentelemetry.sdk.trace.export import BatchSpanProcessor +from ._agent_engine import _get_agent_engine_metrics_setup +from ._agent_engine import telemetry_user_agent_headers +from ._agent_engine_metric_exporter import MIN_EXPORT_INTERVAL_MS from .setup import OTelHooks if TYPE_CHECKING: from google.auth.credentials import Credentials + from google.auth.transport.requests import AuthorizedSession + from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter logger = logging.getLogger("google_adk." + __name__) @@ -57,6 +62,12 @@ _DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT = ( "https://telemetry.mtls.googleapis.com/v1/traces" ) +_DEFAULT_TELEMETRY_METRICS_ENDPOINT = ( + "https://telemetry.googleapis.com/v1/metrics" +) +_DEFAULT_MTLS_TELEMETRY_METRICS_ENDPOINT = ( + "https://telemetry.mtls.googleapis.com/v1/metrics" +) class _MtlsEndpoint(enum.Enum): @@ -118,17 +129,18 @@ def get_gcp_exporters( metric_readers: list[MetricReader] = [] if enable_cloud_metrics: - exporter = _get_gcp_metrics_exporter(project_id) - if exporter: - metric_readers.append(exporter) + if reader := _get_gcp_metrics_exporter((credentials, project_id)): + metric_readers.append(reader) + if agent_engine_metrics := _get_agent_engine_metrics_setup(): + span_processors.append(agent_engine_metrics.span_processor) log_record_processors: list[LogRecordProcessor] = [] if enable_cloud_logging: - exporter = _get_gcp_logs_exporter( + logs_exporter = _get_gcp_logs_exporter( project_id=project_id, ) - if exporter: - log_record_processors.append(exporter) + if logs_exporter: + log_record_processors.append(logs_exporter) return OTelHooks( span_processors=span_processors, @@ -145,55 +157,114 @@ def _get_gcp_span_exporter(credentials: Credentials) -> SpanProcessor: session = AuthorizedSession(credentials=credentials) - use_client_cert = _use_client_cert_effective() - if use_client_cert: - client_cert_source = ( - mtls.default_client_cert_source() - if mtls.has_default_client_cert_source() - else None - ) - session.configure_mtls_channel() - endpoint = _get_api_endpoint(client_cert_source) - else: - endpoint = _DEFAULT_TELEMETRY_TRACES_ENPOINT - - headers = None - if os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY"): - from google.cloud.aiplatform import version as aip_version - - try: - from opentelemetry.exporter.otlp.proto.http import version as otlp_http_version - except (ImportError, AttributeError): - otlp_http_version = None - - user_agent = f"Vertex-Agent-Engine/{aip_version.__version__}" - if otlp_http_version: - user_agent += ( - f" OTel-OTLP-Exporter-Python/{otlp_http_version.__version__}" - ) - headers = {"User-Agent": user_agent} + endpoint = _get_telemetry_endpoint( + session, + _DEFAULT_TELEMETRY_TRACES_ENPOINT, + _DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT, + ) return BatchSpanProcessor( OTLPSpanExporter( session=session, endpoint=endpoint, - headers=headers, + headers=telemetry_user_agent_headers(), ) ) -def _get_gcp_metrics_exporter(project_id: str) -> MetricReader: - from opentelemetry.exporter.cloud_monitoring import CloudMonitoringMetricsExporter +def _get_gcp_otlp_metric_exporter( + google_auth: tuple[Credentials, str] | None = None, +) -> OTLPMetricExporter | None: + """Returns a raw OTLP push metric exporter to telemetry.googleapis.com. + + This is the default GCP metric exporter (over Cloud Monitoring). It returns a + bare push exporter so any metric reader can drain it -- a periodic reader for + the local ``adk web`` path, or the request-driven reader on Agent Engine's + request-billed runtime. Returns None if the OTLP exporter package is + unavailable. + + Args: + google_auth: optional custom credentials and project_id. + google.auth.default() is used when this is omitted. + """ + try: + from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter + except (ImportError, AttributeError): + logger.warning( + "opentelemetry-exporter-otlp-proto-http is not installed; request-path" + " metric export is disabled." + ) + return None + + from google.auth.transport.requests import AuthorizedSession + + credentials, _ = ( + google_auth if google_auth is not None else google.auth.default() + ) + session = AuthorizedSession(credentials=credentials) + endpoint = _get_telemetry_endpoint( + session, + _DEFAULT_TELEMETRY_METRICS_ENDPOINT, + _DEFAULT_MTLS_TELEMETRY_METRICS_ENDPOINT, + ) + return OTLPMetricExporter( + session=session, + endpoint=endpoint, + headers=telemetry_user_agent_headers(), + ) + + +def _get_telemetry_endpoint( + session: AuthorizedSession, default_endpoint: str, mtls_endpoint: str +) -> str: + """Configures the session for mTLS if enabled and returns the endpoint. + Args: + session: The AuthorizedSession to (maybe) configure for mTLS in place. + default_endpoint: The plain telemetry.googleapis.com endpoint. + mtls_endpoint: The telemetry.mtls.googleapis.com endpoint. + + Returns: + The effective endpoint to export to. + """ + if not _use_client_cert_effective(): + return default_endpoint + client_cert_source = ( + mtls.default_client_cert_source() + if mtls.has_default_client_cert_source() + else None + ) + session.configure_mtls_channel() + return _get_api_endpoint( + client_cert_source, + default_endpoint=default_endpoint, + mtls_endpoint=mtls_endpoint, + ) + + +def _get_gcp_metrics_exporter( + google_auth: tuple[Credentials, str], +) -> MetricReader | None: + """Returns the metric reader to install, or None if metrics are unavailable. + + On Agent Engine this is the request-driven reader (background export is + starved by the request-billed runtime); elsewhere a periodic reader over the + default OTLP metric exporter. + """ + if agent_engine_metrics := _get_agent_engine_metrics_setup(): + return agent_engine_metrics.reader + exporter = _get_gcp_otlp_metric_exporter(google_auth=google_auth) + if exporter is None: + return None return PeriodicExportingMetricReader( - CloudMonitoringMetricsExporter(project_id=project_id), - export_interval_millis=5000, + exporter, + export_interval_millis=MIN_EXPORT_INTERVAL_MS, ) def _get_gcp_logs_exporter( project_id: str, -) -> LogRecordProcessor: +) -> LogRecordProcessor | None: if os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID"): return _get_agent_engine_logs_exporter( project_id=project_id, @@ -283,12 +354,16 @@ def get_gcp_resource(project_id: Optional[str] = None) -> Resource: def _get_api_endpoint( client_cert_source: Callable[[], tuple[bytes, bytes]] | None = None, + default_endpoint: str = _DEFAULT_TELEMETRY_TRACES_ENPOINT, + mtls_endpoint: str = _DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT, ) -> str: """Returns API endpoint based on mTLS configuration and cert availability. Args: client_cert_source: A callable that returns the client certificate and key, or None. + default_endpoint: The endpoint to use without mTLS. + mtls_endpoint: The endpoint to use with mTLS. Returns: str: The API endpoint to be used. @@ -311,9 +386,9 @@ def _get_api_endpoint( if (use_mtls_endpoint is _MtlsEndpoint.ALWAYS) or ( use_mtls_endpoint is _MtlsEndpoint.AUTO and client_cert_source ): - return _DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT + return mtls_endpoint - return _DEFAULT_TELEMETRY_TRACES_ENPOINT + return default_endpoint def _use_client_cert_effective() -> bool: @@ -343,7 +418,7 @@ def _use_client_cert_effective() -> bool: def _get_agent_engine_logs_exporter( *, project_id: str, -): +) -> LogRecordProcessor | None: """Configures logging for Agent Engine. Args: @@ -361,7 +436,7 @@ def _get_agent_engine_logs_exporter( "proceeding with logging disabled because not all packages for" " logging have been installed" ) - return + return None class _SimpleLogRecordProcessor(SimpleLogRecordProcessor): diff --git a/src/google/adk/telemetry/setup.py b/src/google/adk/telemetry/setup.py index 08e5a2ec1cd..645ebef4cc2 100644 --- a/src/google/adk/telemetry/setup.py +++ b/src/google/adk/telemetry/setup.py @@ -102,6 +102,8 @@ def maybe_set_otel_providers( MeterProvider( metric_readers=metric_readers, resource=otel_resource, + # Not collecting on exit to avoid points being collected too close together. + shutdown_on_exit=False, ) ) diff --git a/tests/unittests/telemetry/test_agent_engine.py b/tests/unittests/telemetry/test_agent_engine.py index 986d3aa988c..05811c20f53 100644 --- a/tests/unittests/telemetry/test_agent_engine.py +++ b/tests/unittests/telemetry/test_agent_engine.py @@ -12,53 +12,72 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for trace context propagated from request headers.""" +"""Unit tests for Agent Engine telemetry. + +Covers trace context propagated from request headers and the request-path +metric flushing middleware. The middleware drives the request-driven metric +reader from the request lifecycle: a fire-and-forget collect at request start +and an awaited drain collect after the response body has streamed. Traces and +logs are not flushed here (the Agent Engine AdkApp does that). These tests use +a spy reader; no real time, no network. +""" + +# pylint: disable=protected-access,redefined-outer-name +# pyright: reportPrivateUsage=false from __future__ import annotations +import asyncio +from collections.abc import AsyncIterator +import inspect +from types import SimpleNamespace +from unittest import mock + import fastapi +from google.adk.telemetry import _agent_engine from google.adk.telemetry._agent_engine import get_propagated_context from google.adk.telemetry._agent_engine import TopSpanProcessor from opentelemetry import baggage from opentelemetry import context +from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.sdk.trace import ReadableSpan from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter import pytest -_AE_TRACEPARENT_HEADER = 'Google-Agent-Engine-Traceparent' -_TRACEPARENT_HEADER = 'traceparent' -_SUPPORT_ID_ATTRIBUTE = 'supportID' -_SUPPORT_ID_VALUE = 'support-id-value' -_TOP_SPAN = 'invocation' -_CHILD_SPAN = 'child' +_AE_TRACEPARENT_HEADER = "Google-Agent-Engine-Traceparent" +_TRACEPARENT_HEADER = "traceparent" +_SUPPORT_ID_ATTRIBUTE = "supportID" +_SUPPORT_ID_VALUE = "support-id-value" +_TOP_SPAN = "invocation" +_CHILD_SPAN = "child" -_TRACE_ID_HEX = '4bf92f3577b34da6a3ce929d0e0e4736' -_REMOTE_SPAN_ID_HEX = '00f067aa0ba902b7' -_WELL_FORMED_TRACEPARENT = f'00-{_TRACE_ID_HEX}-{_REMOTE_SPAN_ID_HEX}-01' +_TRACE_ID_HEX = "4bf92f3577b34da6a3ce929d0e0e4736" +_REMOTE_SPAN_ID_HEX = "00f067aa0ba902b7" +_WELL_FORMED_TRACEPARENT = f"00-{_TRACE_ID_HEX}-{_REMOTE_SPAN_ID_HEX}-01" # Values the trace context propagator refuses, either because they do not # match the wire format or because the ids they carry are not usable. _REJECTED_TRACEPARENT_VALUES = [ - 'x', - '00-abc-zz-01', - '', - '00', - '-', - f'00-{_TRACE_ID_HEX}-{_REMOTE_SPAN_ID_HEX}', + "x", + "00-abc-zz-01", + "", + "00", + "-", + f"00-{_TRACE_ID_HEX}-{_REMOTE_SPAN_ID_HEX}", f'00-{"0" * 32}-{_REMOTE_SPAN_ID_HEX}-01', - f'ff-{_TRACE_ID_HEX}-{_REMOTE_SPAN_ID_HEX}-01', + f"ff-{_TRACE_ID_HEX}-{_REMOTE_SPAN_ID_HEX}-01", ] def _request(**headers: str) -> fastapi.Request: """Builds a minimal request carrying the given headers.""" return fastapi.Request({ - 'type': 'http', - 'method': 'POST', - 'path': '/', - 'headers': [ + "type": "http", + "method": "POST", + "path": "/", + "headers": [ (name.lower().encode(), value.encode()) for name, value in headers.items() ], @@ -84,7 +103,7 @@ def _record_spans(ctx: context.Context) -> dict[str, ReadableSpan]: return {span.name: span for span in exporter.get_finished_spans()} -@pytest.mark.parametrize('header_value', _REJECTED_TRACEPARENT_VALUES) +@pytest.mark.parametrize("header_value", _REJECTED_TRACEPARENT_VALUES) def test_rejected_header_still_produces_child_spans(header_value): """A caller-supplied header must not be able to break span creation.""" spans = _record_spans( @@ -94,7 +113,7 @@ def test_rejected_header_still_produces_child_spans(header_value): assert set(spans) == {_TOP_SPAN, _CHILD_SPAN} -@pytest.mark.parametrize('header_value', _REJECTED_TRACEPARENT_VALUES) +@pytest.mark.parametrize("header_value", _REJECTED_TRACEPARENT_VALUES) def test_rejected_header_is_not_stored_in_baggage(header_value): """Only a header the propagator accepted is worth carrying in baggage.""" ctx = get_propagated_context( @@ -104,7 +123,7 @@ def test_rejected_header_is_not_stored_in_baggage(header_value): assert _TRACEPARENT_HEADER not in baggage.get_all(context=ctx) -@pytest.mark.parametrize('baggage_value', _REJECTED_TRACEPARENT_VALUES) +@pytest.mark.parametrize("baggage_value", _REJECTED_TRACEPARENT_VALUES) def test_rejected_value_in_baggage_still_produces_child_spans(baggage_value): """The processor runs on every span, so it cannot trust baggage contents.""" spans = _record_spans(baggage.set_baggage(_TRACEPARENT_HEADER, baggage_value)) @@ -145,7 +164,7 @@ def test_first_span_is_parentless_when_header_is_rejected(): spans = _record_spans( get_propagated_context( _request(**{ - _AE_TRACEPARENT_HEADER: 'x', + _AE_TRACEPARENT_HEADER: "x", _TRACEPARENT_HEADER: _SUPPORT_ID_VALUE, }) ) @@ -154,3 +173,279 @@ def test_first_span_is_parentless_when_header_is_rejected(): assert spans[_TOP_SPAN].parent is None assert spans[_TOP_SPAN].attributes[_SUPPORT_ID_ATTRIBUTE] == _SUPPORT_ID_VALUE assert _SUPPORT_ID_ATTRIBUTE not in spans[_CHILD_SPAN].attributes + + +class _SpyReader: + """Records the order of hook/submit calls made by the middleware.""" + + def __init__(self) -> None: + self.events: list[str] = [] + + def note_request_start(self) -> bool: + self.events.append("start") + return True + + def note_request_end(self) -> bool: + self.events.append("end") + return True + + def note_generate_content_start(self) -> bool: + self.events.append("generate_content") + return False + + def submit_collect(self) -> None: + self.events.append("submit") + return None + + +class _FakeResponse: + """A minimal ASGI-ish response exposing a consumable body_iterator.""" + + def __init__(self, chunks: list[bytes]): + async def _gen() -> AsyncIterator[bytes]: + for chunk in chunks: + yield chunk + + self.body_iterator: AsyncIterator[bytes] = _gen() + + +def test_middleware_glue() -> None: + """note_request_start precedes call_next; end drain only after the body.""" + spy = _SpyReader() + dispatch = _agent_engine._metrics_flushing_dispatch(spy) + + async def _drive() -> _SpyReader: + response = _FakeResponse([b"a", b"b"]) + + async def call_next(request: object) -> _FakeResponse: + del request + spy.events.append("call_next") + return response + + wrapped = await dispatch(object(), call_next) + + # Body not consumed yet: request end drain must not have fired. + assert spy.events == ["start", "submit", "call_next"] + + consumed = [chunk async for chunk in wrapped.body_iterator] + assert consumed == [b"a", b"b"] + return spy + + result = asyncio.run(_drive()) + assert "end" in result.events + assert ( + result.events.index("start") + < result.events.index("call_next") + < result.events.index("end") + ) + + +def test_metrics_drained_on_request_end() -> None: + """The reader is drained (end + submit) after the body streams, once.""" + spy = _SpyReader() + dispatch = _agent_engine._metrics_flushing_dispatch(spy) + + async def _drive() -> None: + response = _FakeResponse([b"x"]) + + async def call_next(request: object) -> _FakeResponse: + del request + return response + + wrapped = await dispatch(object(), call_next) + # Before the body is consumed, no request-end drain. + assert "end" not in spy.events + _ = [chunk async for chunk in wrapped.body_iterator] + + asyncio.run(_drive()) + assert spy.events.count("end") == 1 + + +def test_drain_failure_does_not_break_response() -> None: + """A reader that raises on drain never breaks the draining response.""" + + class _BoomReader(_SpyReader): + + def note_request_end(self) -> bool: + raise RuntimeError("boom") + + spy = _BoomReader() + dispatch = _agent_engine._metrics_flushing_dispatch(spy) + + async def _drive() -> list[bytes]: + response = _FakeResponse([b"a", b"b"]) + + async def call_next(request: object) -> _FakeResponse: + del request + return response + + wrapped = await dispatch(object(), call_next) + return [chunk async for chunk in wrapped.body_iterator] + + consumed = asyncio.run(_drive()) + assert consumed == [b"a", b"b"] # body streamed despite drain failure. + + +def test_call_next_exception_drains_and_reraises() -> None: + """If call_next raises, note_request_end still runs (no in_flight leak).""" + spy = _SpyReader() + dispatch = _agent_engine._metrics_flushing_dispatch(spy) + + async def _drive() -> None: + async def call_next(request: object) -> _FakeResponse: + del request + raise RuntimeError("boom") + + with pytest.raises(RuntimeError, match="boom"): + await dispatch(object(), call_next) + + asyncio.run(_drive()) + # start balanced by end even though the body iterator never installed. + assert spy.events == ["start", "submit", "end", "submit"] + + +@pytest.mark.parametrize( + "otel_to_cloud, metrics_state, expected_middleware", + [ + # GCP telemetry setup never ran: the reader is on no MeterProvider. + (False, "state", 0), + # Not on Agent Engine (or setup failed): nothing to drive. + (True, None, 0), + (True, "state", 1), + ], +) +def test_maybe_install_request_metrics_middleware( + otel_to_cloud: bool, + metrics_state: str | None, + expected_middleware: int, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The middleware is installed only with both a reader and cloud telemetry.""" + state = ( + SimpleNamespace(reader=_SpyReader(), span_processor=None) + if metrics_state + else None + ) + monkeypatch.setattr( + "google.adk.telemetry._agent_engine._get_agent_engine_metrics_setup", + lambda: state, + ) + app = fastapi.FastAPI() + + _agent_engine.maybe_install_request_metrics_middleware( + app, otel_to_cloud=otel_to_cloud + ) + + assert len(app.user_middleware) == expected_middleware + + +@pytest.fixture(autouse=True) +def _clear_agent_engine_metrics_cache(): + """The memoized agent-engine metrics builder must not leak across tests.""" + _agent_engine._get_agent_engine_metrics_setup.cache_clear() + yield + _agent_engine._get_agent_engine_metrics_setup.cache_clear() + + +def test_agent_engine_metrics_skipped_off_agent_engine( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Without GOOGLE_CLOUD_AGENT_ENGINE_ID, no metric state is built.""" + monkeypatch.delenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", raising=False) + + assert _agent_engine._get_agent_engine_metrics_setup() is None + + +def test_agent_engine_metrics_skipped_when_meter_provider_installed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """If a real MeterProvider is already installed, defer to it (return None).""" + monkeypatch.setenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", "123") + monkeypatch.setattr( + "opentelemetry.metrics.get_meter_provider", + lambda: MeterProvider(), + ) + + assert _agent_engine._get_agent_engine_metrics_setup() is None + + +def test_agent_engine_metrics_built_on_agent_engine( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """On Agent Engine (no MeterProvider yet), the metric state is built.""" + monkeypatch.setenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", "123") + monkeypatch.setattr( + "opentelemetry.metrics.get_meter_provider", + lambda: mock.MagicMock(), # not an SDK MeterProvider. + ) + fake_state = mock.MagicMock(name="metrics_state") + monkeypatch.setattr( + "google.adk.telemetry.google_cloud._get_gcp_otlp_metric_exporter", + lambda **_: mock.MagicMock(name="exporter"), + ) + monkeypatch.setattr( + "google.adk.telemetry._agent_engine_metric_exporter.build_request_driven_metrics", + lambda exporter: fake_state, + ) + + assert _agent_engine._get_agent_engine_metrics_setup() is fake_state + + +def test_agent_engine_metrics_memoized( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The result is cached: the 'already installed' check runs only once.""" + monkeypatch.setenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", "123") + monkeypatch.setattr( + "opentelemetry.metrics.get_meter_provider", + lambda: mock.MagicMock(), + ) + fake_state = mock.MagicMock(name="metrics_state") + calls = {"n": 0} + + def _build(exporter): + del exporter + calls["n"] += 1 + return fake_state + + monkeypatch.setattr( + "google.adk.telemetry.google_cloud._get_gcp_otlp_metric_exporter", + lambda **_: mock.MagicMock(name="exporter"), + ) + monkeypatch.setattr( + "google.adk.telemetry._agent_engine_metric_exporter.build_request_driven_metrics", + _build, + ) + + first = _agent_engine._get_agent_engine_metrics_setup() + second = _agent_engine._get_agent_engine_metrics_setup() + + assert first is second is fake_state + assert calls["n"] == 1 + + +def test_agent_engine_metrics_none_when_exporter_unavailable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A missing GCP metric exporter yields None, not an error.""" + monkeypatch.setenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", "123") + monkeypatch.setattr( + "opentelemetry.metrics.get_meter_provider", + lambda: mock.MagicMock(), + ) + monkeypatch.setattr( + "google.adk.telemetry.google_cloud._get_gcp_otlp_metric_exporter", + lambda **_: None, + ) + + assert _agent_engine._get_agent_engine_metrics_setup() is None + + +def test_agent_engine_metrics_builder_takes_no_args() -> None: + """@functools.cache keys on args, so the one cache entry shared between the + exporter-setup and middleware-install call sites only holds if the builder is + nullary. Guard against a param sneaking in and silently breaking export.""" + sig = inspect.signature( + _agent_engine._get_agent_engine_metrics_setup.__wrapped__ + ) + assert not sig.parameters diff --git a/tests/unittests/telemetry/test_agent_engine_metric_exporter.py b/tests/unittests/telemetry/test_agent_engine_metric_exporter.py new file mode 100644 index 00000000000..ac2fc78fcb1 --- /dev/null +++ b/tests/unittests/telemetry/test_agent_engine_metric_exporter.py @@ -0,0 +1,287 @@ +# 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. + +"""Unit tests for the reader in `_agent_engine_metric_exporter`. + +Deterministic: no real time, no network. A fake monotonic clock is injected into +the reader; collects are driven inline (call a `note_*` hook and, when it +returns +True, run the collect synchronously so the fake clock stamps the export). The +scenarios mirror the diagrams in the module docstring (baseline drain, overlap +batching, the four guidepost points, the sub-floor skip). Two blanket invariants +are asserted over every scenario: + + I2 -- consecutive collects are >= FLOOR apart. + I1 -- every collect lands inside some [request_start, request_end] window. +""" + +# pylint: disable=protected-access,redefined-outer-name +# This is a unit test of a private module, so private access is expected. +# pyright: reportPrivateUsage=false + +from google.adk.telemetry import _agent_engine_metric_exporter as _metrics +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import MetricExporter +from opentelemetry.sdk.metrics.export import MetricExportResult +from opentelemetry.sdk.metrics.export import MetricsData +import pytest + +_PERIOD_S = 10.0 +_FLOOR_S = 3.0 + + +class _RecordingExporter(MetricExporter): + """Records the fake-clock time of every export.""" + + def __init__(self, clock: list[float]): + super().__init__() + self._clock: list[float] = clock + self.times: list[float] = [] + + def export( + self, + metrics_data: MetricsData, + timeout_millis: float = 10_000, + **kwargs: object, + ) -> MetricExportResult: + del metrics_data, timeout_millis, kwargs # unused + self.times.append(self._clock[0]) + return MetricExportResult.SUCCESS + + def force_flush(self, timeout_millis: float = 10_000) -> bool: + del timeout_millis # unused + return True + + def shutdown(self, timeout_millis: float = 30_000, **kwargs: object) -> None: + del timeout_millis, kwargs # unused + + +class _Harness: + """Drives a reader with a fake clock and records collects + request windows.""" + + def __init__(self, period_s: float = _PERIOD_S, floor_s: float = _FLOOR_S): + self.t: list[float] = [0.0] + self.exporter: _RecordingExporter = _RecordingExporter(self.t) + self.reader: _metrics._RequestDrivenMetricReader = ( + _metrics._RequestDrivenMetricReader( + self.exporter, + export_interval_millis=period_s * 1000.0, + floor_millis=floor_s * 1000.0, + now=lambda: self.t[0], + ) + ) + self.meter_provider: MeterProvider = MeterProvider( + metric_readers=[self.reader] + ) + # A cumulative counter with a recorded value so every collect has data to + # export (an empty collect exports nothing). + self.meter_provider.get_meter("test").create_counter("c").add(1) + self._open: dict[str, float] = {} + self.windows: list[tuple[float, float]] = [] + + def at(self, when: float) -> "_Harness": + self.t[0] = float(when) + return self + + def start(self, rid: str) -> None: + self._open[rid] = self.t[0] + if self.reader.note_request_start(): + self.reader.collect_now() + + def end(self, rid: str) -> None: + self.windows.append((self._open.pop(rid), self.t[0])) + if self.reader.note_request_end(): + self.reader.collect_now() + + def generate_content(self) -> None: + if self.reader.note_generate_content_start(): + self.reader.collect_now() + + @property + def collects(self) -> list[float]: + return list(self.exporter.times) + + def close(self) -> None: + self.meter_provider.shutdown() + + +# --- Scenario builders (each returns a driven harness). -------------------- + + +def _scenario_baseline_drain() -> _Harness: + """an isolated request collects when it drains to zero.""" + h = _Harness() + h.at(0).start("r1") + h.at(5).end("r1") + assert h.collects == [5.0] + return h + + +def _scenario_overlap_batched() -> _Harness: + """A burst of overlapping requests produces a single collect.""" + h = _Harness() + h.at(0).start("r1") + h.at(1).start("r2") + h.at(2).start("r3") + h.at(3).end("r1") + h.at(4).end("r2") + h.at(5).end("r3") + assert h.collects == [5.0] # one collect for all three. + return h + + +def _scenario_guidepost_consumed_by_drain() -> _Harness: + """A guidepost inside a lone request is swept by its drain.""" + h = _Harness() + h.at(0).start("r1") + h.at(12).end("r1") # crosses the guidepost at 10, but only drains here. + assert h.collects == [12.0] # the guidepost never fired on its own. + return h + + +def _scenario_guidepost_fires_at_start() -> _Harness: + """Under continuous overlap, a guidepost fires at next start.""" + h = _Harness() + h.at(0).start("r1") + h.at(2).start("r2") + h.at(4).start("r3") + h.at(11).start("r4") # guidepost (10) crossed, overlap -> collect at start. + h.at(12).end("r1") + h.at(13).end("r2") + h.at(14).end("r3") + h.at(16).end("r4") # baseline drain collect. + assert h.collects == [11.0, 16.0] + return h + + +def _scenario_guidepost_muted() -> _Harness: + """A guidepost within FLOOR of the last collect is muted.""" + h = _Harness() + h.at(0).start("r1") + h.at(9).end("r1") # drain collect at 9. + h.at(9).start("r2") + h.at(10).start("r3") # guidepost due, but 10-9 < FLOOR -> muted, no collect. + h.at(11).end("r2") + h.at(12).end("r3") # next collect is this drain. + assert h.collects == [9.0, 12.0] + return h + + +def _scenario_generate_content_backstop() -> _Harness: + """A lone long request collects off its generate_content spans.""" + h = _Harness() + h.at(0).start("r1") + h.at(5).generate_content() # 5s into busy period (<1.5*PERIOD) -> no collect. + h.at(10).generate_content() # 10s into busy period (<15) -> no collect. + h.at(21).generate_content() # 21s into busy period (>=15) -> collect at 21. + h.at(30).generate_content() # 9s since last collect -> no collect. + h.at(37).generate_content() # 16s since last collect (>=15) -> collect at 37. + h.at(40).end("r1") # drain collect at 40 (>= FLOOR after 37). + assert h.collects == [21.0, 37.0, 40.0] + return h + + +def _scenario_short_first_request_not_preempted() -> _Harness: + """A short first request's drain carries its points; no premature gen collect. + + Regression for the empty-metrics bug: point 4 used to treat "no collect yet" + as overdue, so the first inference span of the very first request fired a + collect *before* the request's metrics were recorded. That collect stamped the + floor and muted the request-end drain (< FLOOR later) that carries the points, + so nothing useful was ever exported. The collect must land at the drain (t=4), + not at the generation (t=2). + + Returns: + The driven harness, for the shared invariant checks. + """ + h = _Harness() + h.at(0).start("r1") + h.at(2).generate_content() # first span of a short first req -> no collect. + h.at(4).end("r1") # drain (would be muted if a collect had fired at t=2). + assert h.collects == [4.0] + return h + + +def _scenario_subfloor_skip() -> _Harness: + """A sub-floor request draining right after a collect is skipped.""" + h = _Harness() + h.at(0).start("r1") + h.at(5).end("r1") # collect at 5. + h.at(6).start("r2") + h.at(6.5).end("r2") # 6.5-5 < FLOOR -> skipped; its points ride the next. + h.at(9).start("r3") + h.at(9).end("r3") # 9-5 >= FLOOR -> collect at 9 (sweeps r2's points). + assert h.collects == [5.0, 9.0] + return h + + +_SCENARIOS = { + "baseline_drain": _scenario_baseline_drain, + "overlap_batched": _scenario_overlap_batched, + "guidepost_consumed_by_drain": _scenario_guidepost_consumed_by_drain, + "guidepost_fires_at_start": _scenario_guidepost_fires_at_start, + "guidepost_muted": _scenario_guidepost_muted, + "generate_content_backstop": _scenario_generate_content_backstop, + "short_first_request_not_preempted": ( + _scenario_short_first_request_not_preempted + ), + "subfloor_skip": _scenario_subfloor_skip, +} + + +@pytest.mark.parametrize("name", list(_SCENARIOS)) +def test_scenario_invariants(name: str) -> None: + """Every scenario honors I1 (in-flight) and I2 (floor spacing).""" + h = _SCENARIOS[name]() + try: + collects = h.collects + assert collects, "scenario produced no collects" + + # I2: consecutive collects are >= FLOOR apart. + for a, b in zip(collects, collects[1:]): + assert b - a >= _FLOOR_S, f"{name}: floor violated: {collects}" + + # I1: each collect lands inside some [start, end] request window. + for c in collects: + assert any( + start <= c <= end for start, end in h.windows + ), f"{name}: collect {c} outside all windows {h.windows}" + finally: + h.close() + + +def test_floor_seconds_default(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv( + _metrics.GOOGLE_CLOUD_AGENT_ENGINE_METRICS_COLLECTION_INTERVAL_FLOOR_MS, + raising=False, + ) + assert _metrics._floor_seconds() == _metrics.MIN_EXPORT_INTERVAL_MS / 1000.0 + + +def test_floor_seconds_env_override(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv( + _metrics.GOOGLE_CLOUD_AGENT_ENGINE_METRICS_COLLECTION_INTERVAL_FLOOR_MS, + "1500", + ) + assert _metrics._floor_seconds() == 1.5 + + +def test_floor_seconds_invalid_falls_back( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv( + _metrics.GOOGLE_CLOUD_AGENT_ENGINE_METRICS_COLLECTION_INTERVAL_FLOOR_MS, + "not-a-number", + ) + assert _metrics._floor_seconds() == _metrics.MIN_EXPORT_INTERVAL_MS / 1000.0 diff --git a/tests/unittests/telemetry/test_google_cloud.py b/tests/unittests/telemetry/test_google_cloud.py index ac1aba1971b..709a71e2a3d 100644 --- a/tests/unittests/telemetry/test_google_cloud.py +++ b/tests/unittests/telemetry/test_google_cloud.py @@ -16,10 +16,17 @@ from typing import Optional from unittest import mock +from google.adk.telemetry import _agent_engine from google.adk.telemetry import google_cloud +from google.adk.telemetry._agent_engine import telemetry_user_agent_headers +from google.adk.telemetry._agent_engine_metric_exporter import MIN_EXPORT_INTERVAL_MS +from google.adk.telemetry.google_cloud import _DEFAULT_MTLS_TELEMETRY_METRICS_ENDPOINT from google.adk.telemetry.google_cloud import _DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT +from google.adk.telemetry.google_cloud import _DEFAULT_TELEMETRY_METRICS_ENDPOINT from google.adk.telemetry.google_cloud import _DEFAULT_TELEMETRY_TRACES_ENPOINT from google.adk.telemetry.google_cloud import _get_api_endpoint +from google.adk.telemetry.google_cloud import _get_gcp_metrics_exporter +from google.adk.telemetry.google_cloud import _get_gcp_otlp_metric_exporter from google.adk.telemetry.google_cloud import _get_gcp_span_exporter from google.adk.telemetry.google_cloud import _use_client_cert_effective from google.adk.telemetry.google_cloud import get_gcp_exporters @@ -28,6 +35,7 @@ from google.auth.transport import mtls from google.auth.transport import requests from opentelemetry.exporter.otlp.proto.http import trace_exporter +from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader import pytest @@ -58,7 +66,7 @@ def test_get_gcp_exporters( ) monkeypatch.setattr( "google.adk.telemetry.google_cloud._get_gcp_metrics_exporter", - lambda project_id: mock.MagicMock(), + lambda google_auth: mock.MagicMock(), ) monkeypatch.setattr( "google.adk.telemetry.google_cloud._get_gcp_logs_exporter", @@ -194,6 +202,34 @@ def test_get_api_endpoint( assert _get_api_endpoint(cert_source) == expected +@pytest.mark.parametrize( + "env_val, cert_source, expected", + [ + ("auto", lambda: b"cert", _DEFAULT_MTLS_TELEMETRY_METRICS_ENDPOINT), + ("auto", None, _DEFAULT_TELEMETRY_METRICS_ENDPOINT), + ("always", None, _DEFAULT_MTLS_TELEMETRY_METRICS_ENDPOINT), + ("never", lambda: b"cert", _DEFAULT_TELEMETRY_METRICS_ENDPOINT), + ], +) +def test_get_api_endpoint_for_metrics( + env_val, + cert_source, + expected, + monkeypatch: pytest.MonkeyPatch, +): + """The same mTLS matrix, with the endpoints overridden for metrics.""" + monkeypatch.setenv("GOOGLE_API_USE_MTLS_ENDPOINT", env_val) + + assert ( + _get_api_endpoint( + cert_source, + default_endpoint=_DEFAULT_TELEMETRY_METRICS_ENDPOINT, + mtls_endpoint=_DEFAULT_MTLS_TELEMETRY_METRICS_ENDPOINT, + ) + == expected + ) + + @mock.patch.object(requests, "AuthorizedSession", autospec=True) @mock.patch( "opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter", @@ -236,3 +272,200 @@ def test_get_gcp_span_exporter_mtls( endpoint=_DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT, headers=None, ) + + +@mock.patch.object(requests, "AuthorizedSession", autospec=True) +@mock.patch( + "opentelemetry.exporter.otlp.proto.http.metric_exporter.OTLPMetricExporter", + autospec=True, +) +@mock.patch( + "google.adk.telemetry.google_cloud._use_client_cert_effective", + autospec=True, +) +@mock.patch( + "google.auth.transport.mtls.has_default_client_cert_source", autospec=True +) +@mock.patch( + "google.auth.transport.mtls.default_client_cert_source", autospec=True +) +def test_get_gcp_otlp_metric_exporter_mtls( + mock_default_cert: mock.MagicMock, + mock_has_cert: mock.MagicMock, + mock_use_cert: mock.MagicMock, + mock_exporter: mock.MagicMock, + mock_session: mock.MagicMock, +): + """Metrics take the mTLS branch onto the *metrics* endpoint, not traces'.""" + credentials = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + mock_use_cert.return_value = True + mock_has_cert.return_value = True + mock_default_cert.return_value = b"cert" + + _get_gcp_otlp_metric_exporter(google_auth=(credentials, "project-id")) + + mock_session.assert_called_once_with(credentials=credentials) + mock_session.return_value.configure_mtls_channel.assert_called_once() + mock_exporter.assert_called_once_with( + session=mock_session.return_value, + endpoint=_DEFAULT_MTLS_TELEMETRY_METRICS_ENDPOINT, + headers=None, + ) + + +@mock.patch.object(requests, "AuthorizedSession", autospec=True) +@mock.patch( + "opentelemetry.exporter.otlp.proto.http.metric_exporter.OTLPMetricExporter", + autospec=True, +) +@mock.patch( + "google.adk.telemetry.google_cloud._use_client_cert_effective", + autospec=True, +) +def test_get_gcp_otlp_metric_exporter_no_mtls( + mock_use_cert: mock.MagicMock, + mock_exporter: mock.MagicMock, + mock_session: mock.MagicMock, +): + """Without a client cert, export goes to the plain metrics endpoint.""" + credentials = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + mock_use_cert.return_value = False + + _get_gcp_otlp_metric_exporter(google_auth=(credentials, "project-id")) + + mock_session.return_value.configure_mtls_channel.assert_not_called() + mock_exporter.assert_called_once_with( + session=mock_session.return_value, + endpoint=_DEFAULT_TELEMETRY_METRICS_ENDPOINT, + headers=None, + ) + + +@mock.patch.object(requests, "AuthorizedSession", autospec=True) +@mock.patch( + "opentelemetry.exporter.otlp.proto.http.metric_exporter.OTLPMetricExporter", + autospec=True, +) +@mock.patch( + "google.adk.telemetry.google_cloud._use_client_cert_effective", + autospec=True, +) +def test_get_gcp_otlp_metric_exporter_sends_agent_engine_user_agent( + mock_use_cert: mock.MagicMock, + mock_exporter: mock.MagicMock, + mock_session: mock.MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """Agent Engine attributes metric traffic via the User-Agent header.""" + credentials = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + mock_use_cert.return_value = False + monkeypatch.setenv("GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY", "1") + + _get_gcp_otlp_metric_exporter(google_auth=(credentials, "project-id")) + + headers = mock_exporter.call_args.kwargs["headers"] + assert headers == telemetry_user_agent_headers() + assert headers["User-Agent"].startswith("Vertex-Agent-Engine/") + + +def test_get_gcp_otlp_metric_exporter_uses_default_credentials( + monkeypatch: pytest.MonkeyPatch, +): + """Omitting google_auth falls back to google.auth.default().""" + credentials = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + monkeypatch.setattr( + "google.auth.default", lambda: (credentials, "project-id") + ) + session = mock.MagicMock(name="session") + monkeypatch.setattr( + "google.auth.transport.requests.AuthorizedSession", + lambda credentials: session, + ) + monkeypatch.setattr( + "google.adk.telemetry.google_cloud._use_client_cert_effective", + lambda: False, + ) + exporter = mock.MagicMock(name="exporter") + monkeypatch.setattr( + "opentelemetry.exporter.otlp.proto.http.metric_exporter.OTLPMetricExporter", + lambda **kwargs: exporter, + ) + + assert _get_gcp_otlp_metric_exporter() is exporter + + +def test_get_gcp_metrics_exporter_wraps_otlp_in_periodic_reader( + monkeypatch: pytest.MonkeyPatch, +): + """Off Agent Engine, metrics go through a 5s periodic reader over OTLP.""" + exporter = mock.MagicMock(name="exporter") + monkeypatch.setattr( + "google.adk.telemetry.google_cloud._get_gcp_otlp_metric_exporter", + lambda google_auth: exporter, + ) + captured = {} + + def _reader(exp, export_interval_millis): + captured["exporter"] = exp + captured["interval"] = export_interval_millis + return mock.MagicMock(spec=PeriodicExportingMetricReader) + + monkeypatch.setattr( + "google.adk.telemetry.google_cloud.PeriodicExportingMetricReader", _reader + ) + + reader = _get_gcp_metrics_exporter(("credentials", "project-id")) + + assert reader is not None + assert captured == {"exporter": exporter, "interval": MIN_EXPORT_INTERVAL_MS} + + +def test_get_gcp_metrics_exporter_none_when_otlp_unavailable( + monkeypatch: pytest.MonkeyPatch, +): + """A missing OTLP exporter package disables metrics instead of raising.""" + monkeypatch.setattr( + "google.adk.telemetry.google_cloud._get_gcp_otlp_metric_exporter", + lambda google_auth: None, + ) + + assert _get_gcp_metrics_exporter(("credentials", "project-id")) is None + + +@pytest.fixture(autouse=True) +def _clear_agent_engine_metrics_cache(): + """The memoized agent-engine metrics builder must not leak across tests.""" + _agent_engine._get_agent_engine_metrics_setup.cache_clear() + yield + _agent_engine._get_agent_engine_metrics_setup.cache_clear() + + +def test_agent_engine_uses_only_request_driven_reader( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """On Agent Engine there must be exactly one metric reader: two exporters + would double-report every point.""" + monkeypatch.delenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", raising=False) + monkeypatch.setattr("google.auth.default", lambda: ("", "project-id")) + fake_state = mock.MagicMock(name="metrics_state") + monkeypatch.setattr( + "google.adk.telemetry.google_cloud._get_agent_engine_metrics_setup", + lambda: fake_state, + ) + monkeypatch.setattr( + "google.adk.telemetry.google_cloud._get_gcp_otlp_metric_exporter", + lambda google_auth=None: mock.MagicMock(name="otlp_exporter"), + ) + + otel_hooks = get_gcp_exporters(enable_cloud_metrics=True) + + assert otel_hooks.metric_readers == [fake_state.reader] + assert otel_hooks.span_processors == [fake_state.span_processor]