diff --git a/src/google/adk/cli/api_server.py b/src/google/adk/cli/api_server.py index 94aac65ae4..54708b2811 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 f096717dc2..ad033df795 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 99e6b55d50..97afc18819 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 0000000000..1577d6d358 --- /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 a0c8f7db39..aa3f6e4895 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 08e5a2ec1c..645ebef4cc 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 986d3aa988..05811c20f5 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 0000000000..ac2fc78fcb --- /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 ac1aba1971..709a71e2a3 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]