chore: import upstream snapshot with attribution

tests/unittests/telemetry/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/unittests/telemetry/functional_test_helpers.py

@@ -0,0 +1,737 @@
+# 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.
+
+"""Shared infrastructure for the telemetry functional tests.
+
+This module hosts:
+
+* The ``SpanDigest`` / ``LogDigest`` types used to build a deterministic
+  comparison shape for in-memory spans + log records.
+* ``install_telemetry`` which patches an in-memory tracer + log exporter
+  onto ADK's globals.
+* The canonical agent / tool / mock-LLM scenario shared across the
+  ``test_functional.py``, ``test_node_functional.py`` and
+  ``test_web_ui_functional.py`` test suites.
+* The ``FunctionalTestCase`` carrier used to parametrize tests against the
+  hand-written expected shapes in ``functional_test_cases.py`` /
+  ``functional_node_test_cases.py``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+from collections.abc import Iterator
+from contextlib import aclosing
+from contextlib import contextmanager
+from dataclasses import dataclass
+from dataclasses import field
+from enum import Enum
+import gc
+import inspect
+import json
+import sys
+from types import CodeType
+from typing import Literal
+from typing import NamedTuple
+from typing import TYPE_CHECKING
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.models.llm_response import LlmResponse
+from google.adk.runners import InMemoryRunner
+from google.adk.telemetry import _metrics
+from google.adk.telemetry import node_tracing
+from google.adk.telemetry import tracing
+from google.adk.tools.function_tool import FunctionTool
+from google.adk.workflow._base_node import START
+from google.adk.workflow._workflow import Workflow
+from google.genai.types import Content
+from google.genai.types import FinishReason
+from google.genai.types import Part
+from opentelemetry.sdk._logs import LoggerProvider
+from opentelemetry.sdk._logs.export import SimpleLogRecordProcessor
+from opentelemetry.sdk.metrics import MeterProvider
+from opentelemetry.sdk.metrics.export import HistogramDataPoint
+from opentelemetry.sdk.metrics.export import InMemoryMetricReader
+from opentelemetry.sdk.metrics.export import NumberDataPoint
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+import pytest
+
+if TYPE_CHECKING:
+  from google.adk.events.event import Event
+  from opentelemetry.sdk.trace import ReadableSpan
+  from opentelemetry.util.types import AttributeValue
+  from opentelemetry.sdk._logs import ReadableLogRecord
+  from opentelemetry.sdk._logs.export import InMemoryLogRecordExporter
+  from opentelemetry.sdk.metrics.export import MetricsData
+  from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+from ..testing_utils import MockModel
+from ..testing_utils import TestInMemoryRunner
+
+# ---------------------------------------------------------------------------
+# Env var + semconv constants.
+# ---------------------------------------------------------------------------
+
+OTEL_OPT_IN = "OTEL_SEMCONV_STABILITY_OPT_IN"
+CAPTURE_CONTENT = "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"
+EXPERIMENTAL_OPT_IN = "gen_ai_latest_experimental"
+ADK_TELEMETRY_SCHEMA_VERSION_OPT_IN = "ADK_TELEMETRY_SCHEMA_VERSION_OPT_IN"
+
+# Stable semconv event names.
+GEN_AI_SYSTEM_MESSAGE_EVENT = "gen_ai.system.message"
+GEN_AI_USER_MESSAGE_EVENT = "gen_ai.user.message"
+GEN_AI_CHOICE_EVENT = "gen_ai.choice"
+
+# Experimental semconv event name.
+GEN_AI_COMPLETION_DETAILS_EVENT = "gen_ai.client.inference.operation.details"
+
+# Difficult to extract, non deterministic attribute keys.
+# We check only for their presence, instead of their values.
+NON_DETERMINISTIC_ATTRIBUTE_KEYS: frozenset[str] = frozenset({
+    "gcp.vertex.agent.event_id",
+    "gen_ai.tool.call.id",
+    "gcp.vertex.agent.associated_event_ids",
+    "gen_ai.conversation.id",
+    "gcp.vertex.agent.invocation_id",
+    "gcp.vertex.agent.session_id",
+})
+
+# Span attribute keys whose values are JSON-serialized strings.
+# These are parsed back into Python objects before comparison so that JSON
+# property ordering doesn't drive test stability.
+JSON_ATTRIBUTE_KEYS: frozenset[str] = frozenset({
+    "gen_ai.input.messages",
+    "gen_ai.output.messages",
+    "gen_ai.system_instructions",
+    "gen_ai.tool.definitions",
+})
+
+# Sentinel used for non deterministic fields that we still want to assert as
+# being present.
+PRESENT = "PRESENT"
+
+
+# ---------------------------------------------------------------------------
+# Digests.
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class LogDigest:
+  """A deterministic digest of a ``ReadableLogRecord``.
+
+  ``attributes`` and ``body`` are normalized via ``_normalize`` so test
+  expectations can be written using plain Python literals (lists/dicts).
+  """
+
+  event_name: str
+  body: object = None
+  attributes: dict[str, object] = field(default_factory=dict)
+
+  @classmethod
+  def from_log(cls, log: ReadableLogRecord) -> LogDigest:
+    attrs: dict[str, object] = {}
+    for k, v in (log.log_record.attributes or {}).items():
+      if k in NON_DETERMINISTIC_ATTRIBUTE_KEYS:
+        attrs[k] = PRESENT
+      else:
+        attrs[k] = _normalize(v)
+    return cls(
+        event_name=log.log_record.event_name or "",
+        body=_normalize(log.log_record.body),
+        attributes=attrs,
+    )
+
+
+@dataclass(frozen=True)
+class SpanDigest:
+  """A deterministic digest of a span in the in-memory span tree.
+
+  In addition to the span's own name + attributes + child spans, each
+  digest also carries the ``LogDigest`` records that were emitted while
+  the span was the active span (matched by ``log_record.span_id``).
+  """
+
+  name: str
+  attributes: dict[str, AttributeValue]
+  children: list[SpanDigest] = field(default_factory=list)
+  logs: list[LogDigest] = field(default_factory=list)
+
+  @classmethod
+  def from_span(cls, span: ReadableSpan) -> SpanDigest:
+    """Builds a single ``SpanDigest`` (no children, no logs) from a span.
+
+    Attribute values are normalized so that:
+    * Non-deterministic keys collapse to the ``PRESENT`` sentinel.
+    * JSON-serialized attribute values are parsed into Python objects.
+    * All other values pass through ``_normalize`` (tuples → lists,
+      enums → ``.value``, ``None`` dict entries dropped).
+    """
+    determinized_attributes: dict[str, AttributeValue] = {}
+    for attr_key, attr_val in (span.attributes or {}).items():
+      if attr_key in NON_DETERMINISTIC_ATTRIBUTE_KEYS:
+        determinized_attributes[attr_key] = PRESENT
+      elif attr_key in JSON_ATTRIBUTE_KEYS and isinstance(attr_val, str):
+        determinized_attributes[attr_key] = _normalize(json.loads(attr_val))
+      else:
+        determinized_attributes[attr_key] = _normalize(attr_val)
+    return cls(name=span.name, attributes=determinized_attributes)
+
+  @classmethod
+  def build(
+      cls,
+      spans: tuple[ReadableSpan, ...],
+      logs: tuple[ReadableLogRecord, ...] = (),
+  ) -> SpanDigest:
+    """Builds the in-memory span tree, attaching logs by span id.
+
+    Used for clear diffs with pytest assertions.
+    """
+    digest_by_id: dict[int, SpanDigest] = {}
+    for span in spans:
+      if span.context is None:
+        continue
+      digest_by_id[span.context.span_id] = cls.from_span(span)
+
+    # Attach each log to its enclosing span (matched by span_id).
+    for log in logs:
+      span_id = log.log_record.span_id
+      if span_id is None or span_id == 0:
+        continue
+      digest = digest_by_id.get(span_id)
+      if digest is None:
+        continue
+      digest.logs.append(LogDigest.from_log(log))
+
+    root: SpanDigest | None = None
+    for span in spans:
+      if span.context is None:
+        continue
+      digest = digest_by_id[span.context.span_id]
+      if span.parent and span.parent.span_id in digest_by_id:
+        parent_digest = digest_by_id[span.parent.span_id]
+        parent_digest.children.append(digest)
+      else:
+        if root is not None:
+          raise ValueError("Multiple root spans found.")
+        root = digest
+
+    # Sort for deterministic comparisons.
+    for digest in digest_by_id.values():
+      digest.children.sort(key=lambda s: s.name)
+      digest.logs[:] = sorted_log_digests(digest.logs)
+
+    if root is None:
+      raise ValueError("No root span found in the provided spans.")
+    return root
+
+  def all_logs(self) -> list[LogDigest]:
+    """Returns all log digests in the tree, sorted deterministically."""
+    collected: list[LogDigest] = []
+
+    def _walk(node: SpanDigest) -> None:
+      collected.extend(node.logs)
+      for child in node.children:
+        _walk(child)
+
+    _walk(self)
+    return sorted_log_digests(collected)
+
+
+def sorted_log_digests(logs: list[LogDigest]) -> list[LogDigest]:
+  """Returns ``logs`` sorted in a stable, content-derived order."""
+  return sorted(
+      logs,
+      key=lambda log: (
+          log.event_name,
+          json.dumps(log.body, sort_keys=True, default=str),
+          json.dumps(log.attributes, sort_keys=True, default=str),
+      ),
+  )
+
+
+class _NonDeterministic:
+  """Sentinel for a metric value that is non-deterministic (e.g. wall-clock)."""
+
+  __slots__ = ()
+
+  def __repr__(self) -> str:
+    return "NON_DETERMINISTIC"
+
+
+# Marks a recorded metric value that cannot be pinned (e.g. ``*.duration``
+# wall-clock timings); used in place of the actual value on both sides.
+NON_DETERMINISTIC = _NonDeterministic()
+
+
+@dataclass(frozen=True)
+class MetricPoint:
+  """A single recorded metric data point."""
+
+  attributes: dict[str, AttributeValue]
+  value: object
+
+  def __hash__(self) -> int:
+    return hash(
+        (json.dumps(self.attributes, sort_keys=True, default=str), self.value)
+    )
+
+
+class HistogramSpec(NamedTuple):
+  """Locates one ADK metric histogram so a test can redirect it.
+
+  ``module`` is the module holding the histogram, ``attr`` the global on it to
+  monkeypatch, and ``metric_name`` the instrument name it is recreated under.
+  """
+
+  module: object
+  attr: str
+  metric_name: str
+
+
+# Histograms recorded by ADK. Each test redirects these onto an in-memory
+# reader so the recorded points can be asserted.
+_PATCHED_HISTOGRAMS: tuple[HistogramSpec, ...] = (
+    HistogramSpec(
+        module=_metrics,
+        attr="_agent_invocation_duration",
+        metric_name="gen_ai.invoke_agent.duration",
+    ),
+    HistogramSpec(
+        module=_metrics,
+        attr="_tool_execution_duration",
+        metric_name="gen_ai.execute_tool.duration",
+    ),
+    HistogramSpec(
+        module=_metrics,
+        attr="_client_operation_duration",
+        metric_name="gen_ai.client.operation.duration",
+    ),
+    HistogramSpec(
+        module=_metrics,
+        attr="_client_token_usage",
+        metric_name="gen_ai.client.token.usage",
+    ),
+    HistogramSpec(
+        module=_metrics,
+        attr="_workflow_invocation_duration",
+        metric_name="gen_ai.invoke_workflow.duration",
+    ),
+    HistogramSpec(
+        module=_metrics,
+        attr="_invoke_agent_inference_calls",
+        metric_name="gen_ai.invoke_agent.inference_calls",
+    ),
+    HistogramSpec(
+        module=_metrics,
+        attr="_invoke_agent_tool_calls",
+        metric_name="gen_ai.invoke_agent.tool_calls",
+    ),
+)
+
+
+def _grouped_metric_points(
+    metrics_data: MetricsData,
+) -> dict[str, frozenset[MetricPoint]]:
+  """Groups every recorded point by metric name as an order-free frozenset."""
+  grouped: dict[str, set[MetricPoint]] = {}
+  for resource_metric in metrics_data.resource_metrics:
+    for scope_metric in resource_metric.scope_metrics:
+      for metric in scope_metric.metrics:
+        for dp in metric.data.data_points:
+          # Sum histograms expose ``.sum``; gauge / counter points expose
+          # ``.value``. isinstance (not hasattr) keeps the typing precise.
+          if isinstance(dp, HistogramDataPoint):
+            value = dp.sum
+          elif isinstance(dp, NumberDataPoint):
+            value = dp.value
+          else:
+            value = NON_DETERMINISTIC
+          # ``*.duration`` histograms record wall-clock timings, which are
+          # non-deterministic; replace them so expectations need not pin a
+          # timing.
+          if metric.name.endswith(".duration"):
+            value = NON_DETERMINISTIC
+          grouped.setdefault(metric.name, set()).add(
+              MetricPoint(attributes=dict(dp.attributes), value=value)
+          )
+  return {name: frozenset(points) for name, points in grouped.items()}
+
+
+@dataclass(frozen=True)
+class TelemetryDigest:
+  """The full telemetry surface produced by one scenario run.
+
+  Bundles the root span tree (with per-span logs attached) and every recorded
+  metric point grouped by metric name. Points are held in a frozenset per
+  group so equality is independent of recording / authoring order. Test cases
+  hand-write the expected instance; ``build`` produces the actual one.
+  """
+
+  root_span: SpanDigest
+  metric_points: dict[str, frozenset[MetricPoint]]
+
+  @classmethod
+  def build(
+      cls,
+      spans: tuple[ReadableSpan, ...],
+      logs: tuple[ReadableLogRecord, ...],
+      metrics_data: MetricsData,
+  ) -> TelemetryDigest:
+    """Builds the actual digest from in-memory spans, logs and metrics."""
+    return cls(
+        root_span=SpanDigest.build(spans, logs),
+        metric_points=_grouped_metric_points(metrics_data),
+    )
+
+
+def _normalize(value: object) -> object:
+  """Normalizes a value for stable equality.
+
+  * Tuples become lists (OTel coerces sequences to tuples on attributes).
+  * Enums become their ``.value``.
+  * Dict entries whose value is ``None`` are dropped (these are inserted by
+    pydantic ``model_dump`` for unset fields and would dominate diffs).
+  """
+  if isinstance(value, Enum):
+    return value.value
+  if isinstance(value, tuple):
+    return [_normalize(v) for v in value]
+  if isinstance(value, list):
+    return [_normalize(v) for v in value]
+  if isinstance(value, dict):
+    return {k: _normalize(v) for k, v in value.items() if v is not None}
+  return value
+
+
+# ---------------------------------------------------------------------------
+# Telemetry plumbing.
+# ---------------------------------------------------------------------------
+
+
+def install_telemetry(
+    monkeypatch: pytest.MonkeyPatch,
+    span_exporter: InMemorySpanExporter,
+    log_exporter: InMemoryLogRecordExporter,
+    metric_reader: InMemoryMetricReader,
+) -> None:
+  """Installs an in-memory tracer + log exporter + metric reader.
+
+  Spans, logs and metric points emitted by ADK during the test are written
+  into the provided exporters / reader. All three MUST be passed in so each
+  test makes the choice of sink explicit (e.g. ``InMemoryLogRecordExporter``
+  vs ``WebUILogExporter``).
+  """
+  tracer_provider = TracerProvider()
+  tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
+  real_tracer = tracer_provider.get_tracer(__name__)
+
+  monkeypatch.setattr(
+      tracing.tracer,
+      "start_as_current_span",
+      real_tracer.start_as_current_span,
+  )
+  monkeypatch.setattr(
+      tracing.tracer,
+      "start_span",
+      real_tracer.start_span,
+  )
+  monkeypatch.setattr(
+      node_tracing.tracer,
+      "start_as_current_span",
+      real_tracer.start_as_current_span,
+  )
+  monkeypatch.setattr(
+      node_tracing.tracer,
+      "start_span",
+      real_tracer.start_span,
+  )
+
+  logger_provider = LoggerProvider()
+  logger_provider.add_log_record_processor(
+      SimpleLogRecordProcessor(log_exporter)
+  )
+  real_logger = logger_provider.get_logger(__name__)
+  monkeypatch.setattr(tracing.otel_logger, "emit", real_logger.emit)
+
+  meter_provider = MeterProvider(metric_readers=[metric_reader])
+  meter = meter_provider.get_meter("functional_test_meter")
+  for spec in _PATCHED_HISTOGRAMS:
+    monkeypatch.setattr(
+        spec.module, spec.attr, meter.create_histogram(spec.metric_name)
+    )
+
+
+# ---------------------------------------------------------------------------
+# Canonical agent / tool / mock-LLM scenario.
+# ---------------------------------------------------------------------------
+
+USER_PROMPT = "hello"
+AGENT_NAME = "some_root_agent"
+AGENT_DESCRIPTION = "A sample root agent."
+BASE_INSTRUCTION = "you are helpful"
+# ADK auto-appends agent identity info to the system instruction when the
+# agent is invoked as the root of an InMemoryRunner directly.
+FULL_SYSTEM_INSTRUCTION = (
+    f"{BASE_INSTRUCTION}\n\n"
+    f'You are an agent. Your internal name is "{AGENT_NAME}".'
+    f' The description about you is "{AGENT_DESCRIPTION}".'
+)
+FINAL_TEXT = "text response"
+TOOL_NAME = "some_tool"
+TOOL_DESCRIPTION = "A sample tool."
+TOOL_ARGS = {"arg1": "val1"}
+TOOL_RESULT_PREFIX = "processed "
+TOOL_RESULT = f"{TOOL_RESULT_PREFIX}{TOOL_ARGS['arg1']}"
+
+# The node scenario uses a workflow node whose output drives the agent's
+# input. The workflow itself wraps the same agent.
+WORKFLOW_NAME = "my_workflow"
+# The root workflow invokes a nested workflow whose sole node produces the
+# input for the agent. The nested workflow exercises the `gen_ai.workflow.nested`
+# span attribute + metric dimension (only nested workflows carry it).
+NESTED_WORKFLOW_NAME = "my_nested_workflow"
+NODE_NAME = "some_node"
+NODE_RESULT = "some result"
+NODE_USER_ID = "some_user"
+NODE_APP_NAME = "some_app"
+
+
+def _make_llm_response(part: Part) -> LlmResponse:
+  return LlmResponse(
+      content=Content(role="model", parts=[part]),
+      finish_reason=FinishReason.STOP,
+  )
+
+
+def build_test_agent(*, failing: bool = False) -> Agent:
+  """Builds the canonical 1-tool, 2-LLM-turn agent."""
+  mock_model = MockModel.create(
+      responses=[
+          _make_llm_response(
+              Part.from_function_call(name=TOOL_NAME, args=TOOL_ARGS)
+          ),
+          _make_llm_response(Part.from_text(text=FINAL_TEXT)),
+      ]
+  )
+
+  def some_tool(arg1: str) -> str:
+    """A sample tool."""
+    if failing:
+      raise ValueError("This tool always fails")
+
+    return f"{TOOL_RESULT_PREFIX}{arg1}"
+
+  return Agent(
+      name=AGENT_NAME,
+      description=AGENT_DESCRIPTION,
+      instruction=BASE_INSTRUCTION,
+      model=mock_model,
+      tools=[FunctionTool(some_tool)],
+  )
+
+
+def build_test_runner(*, failing: bool = False) -> TestInMemoryRunner:
+  """Builds a runner around the canonical agent (no workflow wrapper)."""
+  return TestInMemoryRunner(node=build_test_agent(failing=failing))
+
+
+def build_test_workflow(*, failing: bool = False) -> Workflow:
+  """Builds the canonical Workflow: a nested workflow feeding the agent."""
+  test_agent = build_test_agent(failing=failing)
+
+  async def some_node(ctx, node_input):
+    return NODE_RESULT
+
+  # Trivial workflow to test o11y of nested workflows
+  nested_workflow = Workflow(
+      name=NESTED_WORKFLOW_NAME,
+      edges=[(START, some_node)],
+  )
+
+  return Workflow(
+      name=WORKFLOW_NAME,
+      edges=[(START, nested_workflow, test_agent)],
+  )
+
+
+async def run_node_scenario(
+    *, failing: bool = False, event_sink: list[Event] | None = None
+) -> list[Event]:
+  """Runs the workflow scenario to completion, draining the event stream.
+
+  If ``event_sink`` is provided, collected events are appended to it as they
+  are drained. This lets callers inspect the events that were emitted before
+  an exception propagates (e.g. when ``failing=True``).
+  """
+  workflow = build_test_workflow(failing=failing)
+  runner = InMemoryRunner(app_name=NODE_APP_NAME, node=workflow)
+  session = await runner.session_service.create_session(
+      app_name=NODE_APP_NAME, user_id=NODE_USER_ID
+  )
+  content = Content(parts=[Part.from_text(text=USER_PROMPT)], role="user")
+
+  collected_events: list[Event] = event_sink if event_sink is not None else []
+
+  async with aclosing(
+      runner.run_async(
+          user_id=NODE_USER_ID,
+          session_id=session.id,
+          new_message=content,
+      )
+  ) as agen:
+    async for event in agen:
+      collected_events.append(event)
+
+  return collected_events
+
+
+async def run_agent_scenario(runner: TestInMemoryRunner) -> None:
+  """Runs the non-node scenario to completion, draining the event stream."""
+  async with aclosing(
+      runner.run_async_with_new_session_agen(
+          Content(parts=[Part.from_text(text=USER_PROMPT)], role="user")
+      )
+  ) as agen:
+    async for _ in agen:
+      pass
+
+
+# ---------------------------------------------------------------------------
+# Parametrization carrier.
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class FunctionalTestCase:
+  """One row of the (semconv, capture-content, schema-version) matrix."""
+
+  test_id: str
+  semconv_opt_in: str | None
+  capture_content: str | None
+  schema_version: Literal[1, 2]
+  expected: TelemetryDigest
+
+  def apply_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
+    """Applies the per-case env vars for semconv + content capture.
+
+    Always pins ``ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS=false`` so the tool
+    span attributes remain deterministic across all cases.
+    """
+    if self.semconv_opt_in is None:
+      monkeypatch.delenv(OTEL_OPT_IN, raising=False)
+    else:
+      monkeypatch.setenv(OTEL_OPT_IN, self.semconv_opt_in)
+    if self.capture_content is None:
+      monkeypatch.delenv(CAPTURE_CONTENT, raising=False)
+    else:
+      monkeypatch.setenv(CAPTURE_CONTENT, self.capture_content)
+    monkeypatch.setenv(
+        ADK_TELEMETRY_SCHEMA_VERSION_OPT_IN, str(self.schema_version)
+    )
+    monkeypatch.setenv("ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS", "false")
+
+
+# ---------------------------------------------------------------------------
+# aclosing wrapping assertions.
+# ---------------------------------------------------------------------------
+
+
+@contextmanager
+def aclosing_wrapping_assertions() -> Iterator[None]:
+  """Context manager that asserts every async generator is wrapped in ``aclosing``.
+
+  The check uses ``gc.get_referrers`` on every async generator first
+  iterated within the block, which is expensive (~5 seconds per
+  scenario). Run this once per scenario rather than per parametrized
+  test case.
+
+  On exit the original ``sys`` async-gen hooks are restored.
+  """
+  prev_firstiter, prev_finalizer = sys.get_asyncgen_hooks()
+
+  def wrapped_firstiter(coro: AsyncGenerator[object, object]):
+    if _is_async_context_manager():
+      if prev_firstiter:
+        prev_firstiter(coro)
+      return
+
+    assert any(
+        isinstance(referrer, aclosing)
+        or isinstance(indirect_referrer, aclosing)
+        for referrer in gc.get_referrers(coro)
+        # Some coroutines have a layer of indirection in Python 3.10
+        for indirect_referrer in gc.get_referrers(referrer)
+    ), _no_aclosing_assertion_error(coro)
+
+    if prev_firstiter:
+      prev_firstiter(coro)
+
+  sys.set_asyncgen_hooks(wrapped_firstiter, prev_finalizer)
+  try:
+    yield
+  finally:
+    sys.set_asyncgen_hooks(prev_firstiter, prev_finalizer)
+
+
+def _no_aclosing_assertion_error(coro: AsyncGenerator[object, object]) -> str:
+  first_iter_loc = ""
+  definition_loc = ""
+
+  if (f := inspect.currentframe()) and (f := f.f_back) and (f := f.f_back):
+    first_iter_loc = f'file "{f.f_code.co_filename}" line "{f.f_lineno}"'
+  if (ag_code := getattr(coro, "ag_code", None)) and isinstance(
+      ag_code, CodeType
+  ):
+    definition_loc = (
+        f'file "{ag_code.co_filename}" line "{ag_code.co_firstlineno}"'
+    )
+
+  header_str = f'Async generator "{coro.__name__}" is not wrapped in aclosing'
+  first_iter_str = (
+      f"first iterated in {first_iter_loc}" if first_iter_loc else ""
+  )
+  definition_str = f"defined in {definition_loc}" if definition_loc else ""
+  instruction_str = """
+Wrap the iteration in the following code snippet before iterating:
+
+async with contextlib.aclosing(...) as agen:
+  async for ... as agen:
+     ...
+"""
+
+  return "\n".join(
+      part
+      for part in [
+          header_str,
+          first_iter_str,
+          definition_str,
+          instruction_str,
+      ]
+      if part
+  )
+
+
+def _is_async_context_manager() -> bool:
+  """Checks if this function was invoked by contextlib.asynccontextmanager."""
+  frame = inspect.currentframe()
+  while frame:
+    if (
+        frame.f_code.co_name == "__aenter__"
+        and "contextlib" in frame.f_code.co_filename
+    ):
+      return True
+    frame = frame.f_back
+  return False

tests/unittests/telemetry/test_functional.py

@@ -0,0 +1,331 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.telemetry import tracing
+from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams
+from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
+from google.genai.types import Part
+from mcp import ClientSession as McpClientSession
+from mcp import StdioServerParameters
+from mcp.types import ListToolsResult
+from mcp.types import PaginatedRequestParams
+from mcp.types import Tool as McpTool
+from opentelemetry.instrumentation.google_genai import GoogleGenAiSdkInstrumentor
+from opentelemetry.sdk._logs.export import InMemoryLogRecordExporter
+from opentelemetry.sdk.metrics.export import InMemoryMetricReader
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+import pytest
+from typing_extensions import override
+
+from ..testing_utils import MockModel
+from ..testing_utils import TestInMemoryRunner
+from .functional_test_cases import ALL_CASES
+from .functional_test_cases import EXPECTED_EXPERIMENTAL_SPAN_AND_EVENT_WITH_MCP
+from .functional_test_helpers import aclosing_wrapping_assertions
+from .functional_test_helpers import build_test_runner
+from .functional_test_helpers import CAPTURE_CONTENT
+from .functional_test_helpers import EXPERIMENTAL_OPT_IN
+from .functional_test_helpers import FunctionalTestCase
+from .functional_test_helpers import install_telemetry
+from .functional_test_helpers import OTEL_OPT_IN
+from .functional_test_helpers import run_agent_scenario
+from .functional_test_helpers import SpanDigest
+from .functional_test_helpers import TelemetryDigest
+
+
+@pytest.mark.parametrize("case", ALL_CASES, ids=lambda c: c.test_id)
+@pytest.mark.asyncio
+async def test_telemetry_schema(
+    case: FunctionalTestCase,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+  """Tests creation of spans/logs/metrics in an E2E runner invocation.
+
+  Asserts the entire telemetry schema (spans + attributes + per-span logs +
+  recorded metric points) matches the hand-written expected shape for the
+  given semconv + content-capture configuration.
+  """
+  case.apply_env(monkeypatch)
+
+  span_exporter = InMemorySpanExporter()
+  log_exporter = InMemoryLogRecordExporter()
+  metric_reader = InMemoryMetricReader()
+  install_telemetry(monkeypatch, span_exporter, log_exporter, metric_reader)
+
+  await run_agent_scenario(build_test_runner())
+
+  digest = TelemetryDigest.build(
+      span_exporter.get_finished_spans(),
+      log_exporter.get_finished_logs(),
+      metric_reader.get_metrics_data(),
+  )
+  assert digest == case.expected
+
+
+@pytest.mark.asyncio
+async def test_async_generators_wrapped_in_aclosing(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+  """Asserts each async generator iterated by the scenario is wrapped in ``aclosing``.
+
+  Necessary because instrumentation utilizes contextvars, which run into
+  "ContextVar was created in a different Context" errors when a given
+  coroutine gets indeterminately suspended.
+
+  Kept as a single non-parametrized test because the underlying
+  ``gc.get_referrers`` walk is expensive (~5 seconds per scenario).
+  """
+  install_telemetry(
+      monkeypatch,
+      InMemorySpanExporter(),
+      InMemoryLogRecordExporter(),
+      InMemoryMetricReader(),
+  )
+
+  with aclosing_wrapping_assertions():
+    await run_agent_scenario(build_test_runner())
+
+
+@pytest.mark.asyncio
+async def test_exception_preserves_attributes(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+  """Test when an exception occurs during tool execution, span attributes are still present on spans where they are expected."""
+
+  span_exporter = InMemorySpanExporter()
+  install_telemetry(
+      monkeypatch,
+      span_exporter,
+      InMemoryLogRecordExporter(),
+      InMemoryMetricReader(),
+  )
+
+  with pytest.raises(ValueError, match="This tool always fails"):
+    _ = await run_agent_scenario(build_test_runner(failing=True))
+
+  spans = span_exporter.get_finished_spans()
+
+  assert len(spans) > 1
+  assert all(
+      span.attributes is not None and len(span.attributes) > 0
+      for span in spans
+      if span.name != "invocation"  # not expected to have attributes
+  )
+
+
+@pytest.mark.asyncio
+async def test_no_generate_content_for_gemini_model_when_already_instrumented(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+  """Tests that generate_content span is not created if already instrumented."""
+  span_exporter = InMemorySpanExporter()
+  install_telemetry(
+      monkeypatch,
+      span_exporter,
+      InMemoryLogRecordExporter(),
+      InMemoryMetricReader(),
+  )
+
+  monkeypatch.setattr(
+      tracing,
+      "_instrumented_with_opentelemetry_instrumentation_google_genai",
+      lambda: True,
+  )
+  monkeypatch.setattr(
+      tracing,
+      "_is_gemini_agent",
+      lambda _: True,
+  )
+
+  _ = await run_agent_scenario(build_test_runner())
+
+  spans = span_exporter.get_finished_spans()
+  assert not any(span.name.startswith("generate_content") for span in spans)
+
+
+def test_instrumented_with_opentelemetry_instrumentation_google_genai():
+  instrumentor = GoogleGenAiSdkInstrumentor()
+
+  assert (
+      not tracing._instrumented_with_opentelemetry_instrumentation_google_genai()
+  )
+  try:
+    instrumentor.instrument()
+    assert (
+        tracing._instrumented_with_opentelemetry_instrumentation_google_genai()
+    )
+  finally:
+    instrumentor.uninstrument()
+  assert (
+      not tracing._instrumented_with_opentelemetry_instrumentation_google_genai()
+  )
+
+
+# ---------------------------------------------------------------------------
+# MCP integration: telemetry adds zero ``list_tools()`` calls of its own.
+#
+# The standard ADK ↔ MCP integration path is:
+#
+#   Agent(tools=[McpToolset(...)])
+#     → McpToolset.get_tools()  ─ calls list_tools() ONCE, caches MCPTool list
+#     → BaseLlmFlow loop calls each MCPTool.process_llm_request, which
+#       materializes the tool's FunctionDeclaration into
+#       llm_request.config.tools.
+#
+# By the time the experimental semconv builder reads
+# ``llm_request.config.tools``, MCP tools are ALREADY ``types.Tool``
+# entries with ``function_declarations``. Because the builder is fully
+# synchronous (it never calls ``list_tools()`` itself), the MCP server is
+# queried EXACTLY ONCE per agent invocation regardless of which semconv
+# (or capture mode) is active. These tests pin that contract AND verify
+# the resolved tool definitions surface intact in the experimental
+# telemetry.
+#
+# A ``_FakeMcpSession`` substitutes the live ``McpClientSession`` so the
+# test doesn't need a running MCP server. ``McpToolset.create_session``
+# is patched to hand it out instead of dialing ``StdioServerParameters``.
+# ---------------------------------------------------------------------------
+
+
+class _FakeMcpSession(McpClientSession):
+  """Minimal ``McpClientSession`` stand-in with a counted ``list_tools()``.
+
+  Subclasses ``McpClientSession`` (and skips its real ``__init__``) so that
+  every ``isinstance(x, McpClientSession)`` check in ADK and in the MCP
+  Python client passes, without needing to wire up the underlying anyio
+  memory streams + peer process.
+  """
+
+  def __init__(  # pyright: ignore[reportMissingSuperCall]
+      self, *, tools: list[McpTool]
+  ) -> None:
+    # Deliberately skip ``McpClientSession.__init__``: the real one wants
+    # live anyio streams + a peer process. ``isinstance`` checks still
+    # succeed, which is all ADK's MCP plumbing requires.
+    self._tools: list[McpTool] = tools
+    self.list_tools_call_count: int = 0
+
+  @override
+  async def list_tools(
+      self,
+      cursor: str | None = None,
+      *,
+      params: PaginatedRequestParams | None = None,
+  ) -> ListToolsResult:
+    self.list_tools_call_count += 1
+    return ListToolsResult(tools=list(self._tools))
+
+
+def _make_fake_mcp_toolset(
+    monkeypatch: pytest.MonkeyPatch, fake_session: _FakeMcpSession
+) -> McpToolset:
+  """Returns an ``McpToolset`` whose session manager hands out ``fake_session``.
+
+  Patches the toolset's ``MCPSessionManager`` so:
+    * ``create_session`` returns the fake (no socket / subprocess).
+    * ``close`` is a no-op (the fake holds no resources).
+
+  Connection params are nominally a stdio command but never actually
+  invoked because ``create_session`` is overridden.
+  """
+  toolset = McpToolset(
+      connection_params=StdioConnectionParams(
+          server_params=StdioServerParameters(command="unused-by-test"),
+      )
+  )
+
+  async def _create_session(*_args, **_kwargs):  # pyright: ignore[reportUnknownParameterType, reportMissingParameterType]
+    return fake_session
+
+  async def _close(*_args, **_kwargs):  # pyright: ignore[reportUnknownParameterType, reportMissingParameterType]
+    return None
+
+  monkeypatch.setattr(
+      toolset._mcp_session_manager, "create_session", _create_session  # pyright: ignore[reportPrivateUsage, reportUnknownArgumentType]
+  )
+  monkeypatch.setattr(toolset._mcp_session_manager, "close", _close)  # pyright: ignore[reportPrivateUsage, reportUnknownArgumentType]
+  return toolset
+
+
+def _build_mcp_test_runner(toolset: McpToolset) -> TestInMemoryRunner:
+  """Builds a single-turn agent runner whose only tool source is ``toolset``.
+
+  Single-turn (one ``Part.from_text`` response) so the assertion on
+  ``list_tools_call_count`` is unambiguous: exactly one agent invocation
+  is performed.
+  """
+  mock_model = MockModel.create(
+      responses=[Part.from_text(text="text response")]
+  )
+  test_agent = Agent(
+      name="some_root_agent",
+      description="A sample root agent.",
+      instruction="you are helpful",
+      model=mock_model,
+      tools=[toolset],
+  )
+  return TestInMemoryRunner(node=test_agent)
+
+
+@pytest.mark.asyncio
+async def test_mcp_list_tools_called_once_under_experimental_semconv(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+  """Experimental semconv: exactly one ``list_tools()`` call per invocation.
+
+  By the time the experimental semconv builder inspects
+  ``llm_request.config.tools``, ``McpToolset`` has already materialized
+  each MCP tool into a ``FunctionDeclaration`` — so the synchronous
+  builder never has to (and never does) talk to the MCP server. The
+  MCP-resolved tool definition still surfaces in the experimental
+  telemetry intact, sourced from the ``FunctionDeclaration`` rather than
+  from a fresh ``list_tools()`` call.
+  """
+  monkeypatch.setenv(OTEL_OPT_IN, EXPERIMENTAL_OPT_IN)
+  monkeypatch.setenv(CAPTURE_CONTENT, "span_and_event")
+  monkeypatch.setenv("ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS", "false")
+
+  span_exporter = InMemorySpanExporter()
+  log_exporter = InMemoryLogRecordExporter()
+  install_telemetry(
+      monkeypatch, span_exporter, log_exporter, InMemoryMetricReader()
+  )
+
+  fake_session = _FakeMcpSession(
+      tools=[
+          McpTool(
+              name="mcp_echo",
+              description="Echoes back its input.",
+              inputSchema={
+                  "type": "object",
+                  "properties": {"text": {"type": "string"}},
+                  "required": ["text"],
+              },
+          )
+      ]
+  )
+  toolset = _make_fake_mcp_toolset(monkeypatch, fake_session)
+
+  await run_agent_scenario(_build_mcp_test_runner(toolset))
+
+  assert fake_session.list_tools_call_count == 1
+
+  digest = SpanDigest.build(
+      span_exporter.get_finished_spans(),
+      log_exporter.get_finished_logs(),
+  )
+  assert digest == EXPECTED_EXPERIMENTAL_SPAN_AND_EVENT_WITH_MCP

tests/unittests/telemetry/test_google_cloud.py

@@ -0,0 +1,238 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from typing import Optional
+from unittest import mock
+
+from google.adk.telemetry import google_cloud
+from google.adk.telemetry.google_cloud import _DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT
+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_span_exporter
+from google.adk.telemetry.google_cloud import _use_client_cert_effective
+from google.adk.telemetry.google_cloud import get_gcp_exporters
+from google.adk.telemetry.google_cloud import get_gcp_resource
+import google.auth.credentials
+from google.auth.transport import mtls
+from google.auth.transport import requests
+from opentelemetry.exporter.otlp.proto.http import trace_exporter
+import pytest
+
+
+@pytest.mark.parametrize("enable_cloud_tracing", [True, False])
+@pytest.mark.parametrize("enable_cloud_metrics", [True, False])
+@pytest.mark.parametrize("enable_cloud_logging", [True, False])
+def test_get_gcp_exporters(
+    enable_cloud_tracing: bool,
+    enable_cloud_metrics: bool,
+    enable_cloud_logging: bool,
+    monkeypatch: pytest.MonkeyPatch,
+):
+  """
+  Test initializing correct providers in setup_otel
+  when enabling telemetry via Google O11y.
+  """
+  # Arrange.
+  # Mocking google.auth.default to improve the test time.
+  auth_mock = mock.MagicMock()
+  auth_mock.return_value = ("", "project-id")
+  monkeypatch.setattr(
+      "google.auth.default",
+      auth_mock,
+  )
+  monkeypatch.setattr(
+      "google.adk.telemetry.google_cloud._get_gcp_span_exporter",
+      lambda credentials: mock.MagicMock(),
+  )
+  monkeypatch.setattr(
+      "google.adk.telemetry.google_cloud._get_gcp_metrics_exporter",
+      lambda project_id: mock.MagicMock(),
+  )
+  monkeypatch.setattr(
+      "google.adk.telemetry.google_cloud._get_gcp_logs_exporter",
+      lambda project_id: mock.MagicMock(),
+  )
+
+  # Act.
+  otel_hooks = get_gcp_exporters(
+      enable_cloud_tracing=enable_cloud_tracing,
+      enable_cloud_metrics=enable_cloud_metrics,
+      enable_cloud_logging=enable_cloud_logging,
+  )
+
+  # Assert.
+  # If given telemetry type was enabled,
+  # the corresponding provider should be set.
+  assert len(otel_hooks.span_processors) == (1 if enable_cloud_tracing else 0)
+  assert len(otel_hooks.metric_readers) == (1 if enable_cloud_metrics else 0)
+  assert len(otel_hooks.log_record_processors) == (
+      1 if enable_cloud_logging else 0
+  )
+
+
+@pytest.mark.parametrize("project_id_in_arg", ["project_id_in_arg", None])
+@pytest.mark.parametrize("project_id_on_env", ["project_id_on_env", None])
+def test_get_gcp_resource(
+    project_id_in_arg: Optional[str],
+    project_id_on_env: Optional[str],
+    monkeypatch: pytest.MonkeyPatch,
+):
+  # Arrange.
+  if project_id_on_env is not None:
+    monkeypatch.setenv(
+        "OTEL_RESOURCE_ATTRIBUTES", f"gcp.project_id={project_id_on_env}"
+    )
+
+  # Act.
+  otel_resource = get_gcp_resource(project_id_in_arg)
+
+  # Assert.
+  expected_project_id = (
+      project_id_on_env
+      if project_id_on_env is not None
+      else project_id_in_arg
+      if project_id_in_arg is not None
+      else None
+  )
+  assert otel_resource is not None
+  assert (
+      otel_resource.attributes.get("gcp.project_id", None)
+      == expected_project_id
+  )
+
+
+def test_get_gcp_resource_sets_standard_cloud_resource_id(
+    monkeypatch: pytest.MonkeyPatch,
+):
+  # Arrange.
+  monkeypatch.setenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", "1234567890")
+  monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "us-central1")
+
+  # Act.
+  otel_resource = get_gcp_resource("my-project")
+
+  # Assert.
+  # The Agent Engine dashboard filters on the OTel-standard key.
+  assert otel_resource.attributes.get("cloud.resource_id") == (
+      "//aiplatform.googleapis.com/projects/my-project"
+      "/locations/us-central1/reasoningEngines/1234567890"
+  )
+  assert "cloud.resource.id" not in otel_resource.attributes
+
+
+@mock.patch.object(mtls, "should_use_client_cert", autospec=True)
+def test_use_client_cert_effective_from_mtls(mock_should_use):
+  mock_should_use.return_value = True
+  assert _use_client_cert_effective()
+
+  mock_should_use.return_value = False
+  assert not _use_client_cert_effective()
+
+
+def test_use_client_cert_effective_from_env(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+):
+  with mock.patch.object(
+      mtls,
+      "should_use_client_cert",
+      autospec=True,
+      side_effect=AttributeError,
+  ):
+    monkeypatch.setenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "true")
+    assert _use_client_cert_effective()
+
+    monkeypatch.setenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "false")
+    assert not _use_client_cert_effective()
+
+    # Test invalid value defaults to False
+    monkeypatch.setenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "maybe")
+    assert not _use_client_cert_effective()
+    assert (
+        "Environment variable `GOOGLE_API_USE_CLIENT_CERTIFICATE` must be"
+        " either `true` or `false`"
+        in caplog.text
+    )
+
+
+@pytest.mark.parametrize(
+    "env_val, cert_source, expected",
+    [
+        ("auto", lambda: b"cert", _DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT),
+        ("auto", None, _DEFAULT_TELEMETRY_TRACES_ENPOINT),
+        ("always", None, _DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT),
+        ("never", lambda: b"cert", _DEFAULT_TELEMETRY_TRACES_ENPOINT),
+        ("invalid", None, _DEFAULT_TELEMETRY_TRACES_ENPOINT),
+    ],
+)
+def test_get_api_endpoint(
+    env_val,
+    cert_source,
+    expected,
+    monkeypatch: pytest.MonkeyPatch,
+    caplog: pytest.LogCaptureFixture,
+):
+  monkeypatch.setenv("GOOGLE_API_USE_MTLS_ENDPOINT", env_val)
+  if env_val == "invalid":
+    assert _get_api_endpoint(cert_source) == expected
+    assert (
+        "Environment variable `GOOGLE_API_USE_MTLS_ENDPOINT` must be one of"
+        in caplog.text
+    )
+  else:
+    assert _get_api_endpoint(cert_source) == expected
+
+
+@mock.patch.object(requests, "AuthorizedSession", autospec=True)
+@mock.patch(
+    "opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter",
+    autospec=True,
+)
+@mock.patch(
+    "google.adk.telemetry.google_cloud.BatchSpanProcessor", 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_span_exporter_mtls(
+    mock_default_cert: mock.MagicMock,
+    mock_has_cert: mock.MagicMock,
+    mock_use_cert: mock.MagicMock,
+    mock_batch: mock.MagicMock,
+    mock_exporter: mock.MagicMock,
+    mock_session: mock.MagicMock,
+):
+  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_span_exporter(credentials)
+
+  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_TRACES_ENPOINT,
+      headers=None,
+  )

tests/unittests/telemetry/test_instrumentation.py

@@ -0,0 +1,82 @@
+# 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.
+
+# pylint: disable=protected-access
+
+import time
+from unittest import mock
+
+from google.adk.telemetry import _metrics
+from opentelemetry import trace
+
+
+def test_get_elapsed_s_span_none():
+  """Tests fallback when span is None."""
+  start_time = 10.0
+  with mock.patch("time.monotonic", return_value=12.0):
+    elapsed = _metrics.get_elapsed_s(None, start_time)
+  assert elapsed == 2.0  # 12 - 10
+
+
+def test_get_elapsed_s_span_valid():
+  """Tests duration calculation with valid span times."""
+  mock_span = mock.MagicMock(spec=trace.Span)
+  mock_span.start_time = 1000000000  # 1s in ns
+  mock_span.end_time = 2000000000  # 2s in ns
+  elapsed = _metrics.get_elapsed_s(mock_span, time.monotonic())
+  assert elapsed == 1.0  # (2 - 1) s
+
+
+def test_get_elapsed_s_span_missing_start():
+  """Tests fallback when start_time is missing."""
+  mock_span = mock.MagicMock(spec=trace.Span)
+  del mock_span.start_time
+  mock_span.end_time = 2000000000
+  start_time = 10.0
+  with mock.patch("time.monotonic", return_value=12.0):
+    elapsed = _metrics.get_elapsed_s(mock_span, start_time)
+  assert elapsed == 2.0
+
+
+def test_get_elapsed_s_span_missing_end():
+  """Tests fallback when end_time is missing."""
+  mock_span = mock.MagicMock(spec=trace.Span)
+  mock_span.start_time = 1000000000
+  del mock_span.end_time
+  start_time = 10.0
+  with mock.patch("time.monotonic", return_value=12.0):
+    elapsed = _metrics.get_elapsed_s(mock_span, start_time)
+  assert elapsed == 2.0
+
+
+def test_get_elapsed_s_span_non_int_start():
+  """Tests fallback when start_time is not an integer."""
+  mock_span = mock.MagicMock(spec=trace.Span)
+  mock_span.start_time = 1000000000.0
+  mock_span.end_time = 2000000000
+  start_time = 10.0
+  with mock.patch("time.monotonic", return_value=12.0):
+    elapsed = _metrics.get_elapsed_s(mock_span, start_time)
+  assert elapsed == 2.0
+
+
+def test_get_elapsed_s_span_non_int_end():
+  """Tests fallback when end_time is not an integer."""
+  mock_span = mock.MagicMock(spec=trace.Span)
+  mock_span.start_time = 1000000000
+  mock_span.end_time = 2000000000.0
+  start_time = 10.0
+  with mock.patch("time.monotonic", return_value=12.0):
+    elapsed = _metrics.get_elapsed_s(mock_span, start_time)
+  assert elapsed == 2.0

tests/unittests/telemetry/test_metrics.py

@@ -0,0 +1,261 @@
+# 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.
+
+# pylint: disable=protected-access
+
+from unittest import mock
+
+from google.adk.telemetry import _metrics
+from google.genai import types
+from opentelemetry import metrics
+import pytest
+
+
+@pytest.fixture(name="mock_meter_setup")
+def _mock_meter_setup(monkeypatch):
+  """Sets up mock meter and histograms for testing."""
+  mock_meter = mock.MagicMock()
+  agent_duration_hist = mock.MagicMock(spec=metrics.Histogram)
+  workflow_duration_hist = mock.MagicMock(spec=metrics.Histogram)
+  tool_duration_hist = mock.MagicMock(spec=metrics.Histogram)
+  client_duration_hist = mock.MagicMock(spec=metrics.Histogram)
+  client_token_usage_hist = mock.MagicMock(spec=metrics.Histogram)
+
+  agent_duration_hist.name = "agent_invocation_duration"
+  workflow_duration_hist.name = "workflow_invocation_duration"
+  tool_duration_hist.name = "tool_execution_duration"
+  client_duration_hist.name = "client_operation_duration"
+  client_token_usage_hist.name = "client_token_usage"
+
+  def create_histogram_side_effect(name, **_kwargs):
+    if name == "gen_ai.invoke_agent.duration":
+      return agent_duration_hist
+    elif name == "gen_ai.invoke_workflow.duration":
+      return workflow_duration_hist
+    elif name == "gen_ai.execute_tool.duration":
+      return tool_duration_hist
+    elif name == "gen_ai.client.operation.duration":
+      return client_duration_hist
+    elif name == "gen_ai.client.token.usage":
+      return client_token_usage_hist
+    raise ValueError(f"Unknown metric name: {name}")
+
+  mock_meter.create_histogram.side_effect = create_histogram_side_effect
+
+  # Re-initialize the module-level variables in _metrics with mocked histograms
+  monkeypatch.setattr(_metrics, "meter", mock_meter)
+  monkeypatch.setattr(
+      _metrics, "_agent_invocation_duration", agent_duration_hist
+  )
+  monkeypatch.setattr(
+      _metrics, "_workflow_invocation_duration", workflow_duration_hist
+  )
+  monkeypatch.setattr(_metrics, "_tool_execution_duration", tool_duration_hist)
+  monkeypatch.setattr(
+      _metrics, "_client_operation_duration", client_duration_hist
+  )
+  monkeypatch.setattr(_metrics, "_client_token_usage", client_token_usage_hist)
+
+  return {
+      "meter": mock_meter,
+      "agent_duration": agent_duration_hist,
+      "workflow_duration": workflow_duration_hist,
+      "tool_duration": tool_duration_hist,
+      "client_duration": client_duration_hist,
+      "client_token_usage": client_token_usage_hist,
+  }
+
+
+def test_record_agent_invocation_duration(mock_meter_setup):
+  """Tests record_agent_invocation_duration records correctly."""
+  _metrics.record_agent_invocation_duration(
+      "test_agent",
+      1.0,
+  )
+  agent_duration_hist = mock_meter_setup["agent_duration"]
+  agent_duration_hist.record.assert_called_once()
+  args, kwargs = agent_duration_hist.record.call_args
+  assert args[0] == 1.0
+  want_attributes = {"gen_ai.agent.name": "test_agent"}
+  assert kwargs["attributes"] == want_attributes
+
+
+def test_record_agent_invocation_duration_with_error(mock_meter_setup):
+  """Tests record_agent_invocation_duration records error correctly."""
+  test_error = ValueError("agent failed")
+  _metrics.record_agent_invocation_duration(
+      "test_agent",
+      1.0,
+      error=test_error,
+  )
+  agent_duration_hist = mock_meter_setup["agent_duration"]
+  agent_duration_hist.record.assert_called_once()
+  _, kwargs = agent_duration_hist.record.call_args
+  assert kwargs["attributes"]["error.type"] == "ValueError"
+
+
+def test_record_workflow_invocation_duration_root(mock_meter_setup):
+  """Tests record_workflow_invocation_duration omits nested for the root."""
+  _metrics.record_workflow_invocation_duration(
+      workflow_name="my_workflow",
+      elapsed_s=1.0,
+      nested=False,
+  )
+  hist = mock_meter_setup["workflow_duration"]
+  hist.record.assert_called_once()
+  args, kwargs = hist.record.call_args
+  assert args[0] == 1.0
+  assert kwargs["attributes"] == {
+      "gen_ai.operation.name": "invoke_workflow",
+      "gen_ai.workflow.name": "my_workflow",
+  }
+
+
+def test_record_workflow_invocation_duration_nested_with_error(
+    mock_meter_setup,
+):
+  """Tests record_workflow_invocation_duration records nested + error."""
+  _metrics.record_workflow_invocation_duration(
+      workflow_name="nested_workflow",
+      elapsed_s=2.0,
+      nested=True,
+      error=ValueError("boom"),
+  )
+  hist = mock_meter_setup["workflow_duration"]
+  hist.record.assert_called_once()
+  _, kwargs = hist.record.call_args
+  assert kwargs["attributes"]["gen_ai.workflow.nested"] is True
+  assert kwargs["attributes"]["error.type"] == "ValueError"
+
+
+def test_record_tool_execution_duration(mock_meter_setup):
+  """Tests record_tool_execution_duration records correctly."""
+  _metrics.record_tool_execution_duration(
+      "test_tool",
+      "test_tool_type",
+      "test_agent",
+      0.5,
+  )
+  tool_duration_hist = mock_meter_setup["tool_duration"]
+  tool_duration_hist.record.assert_called_once()
+  args, kwargs = tool_duration_hist.record.call_args
+  assert args[0] == 0.5
+  want_attributes = {
+      "gen_ai.agent.name": "test_agent",
+      "gen_ai.tool.name": "test_tool",
+      "gen_ai.tool.type": "test_tool_type",
+  }
+  assert kwargs["attributes"] == want_attributes
+
+
+def test_record_tool_execution_duration_with_error(mock_meter_setup):
+  """Tests record_tool_execution_duration records error correctly."""
+  test_error = ValueError("tool failed")
+  _metrics.record_tool_execution_duration(
+      "test_tool",
+      "test_tool_type",
+      "test_agent",
+      0.5,
+      error=test_error,
+  )
+  tool_duration_hist = mock_meter_setup["tool_duration"]
+  tool_duration_hist.record.assert_called_once()
+  _, kwargs = tool_duration_hist.record.call_args
+  assert kwargs["attributes"]["error.type"] == "ValueError"
+
+
+def test_record_client_operation_duration(mock_meter_setup):
+  """Tests record_client_operation_duration records correctly."""
+  llm_request = mock.MagicMock(
+      contents=[types.Content(parts=[types.Part(text="hello")])]
+  )
+  response = mock.MagicMock(
+      content=types.Content(parts=[types.Part(text="hello response")])
+  )
+  _metrics.record_client_operation_duration(
+      agent_name="test_agent",
+      elapsed_s=0.1,
+      llm_request=llm_request,
+      responses=[response],
+  )
+  client_duration_hist = mock_meter_setup["client_duration"]
+  client_duration_hist.record.assert_called_once()
+  args, kwargs = client_duration_hist.record.call_args
+  assert args[0] == 0.1
+  want_attributes = {
+      "gen_ai.agent.name": "test_agent",
+      "gen_ai.operation.name": "generate_content",
+      "gen_ai.provider.name": "gemini",
+      "gen_ai.request.model": llm_request.model,
+      "gen_ai.response.model": response.model_version,
+  }
+  assert kwargs["attributes"] == want_attributes
+
+
+def test_record_client_token_usage(mock_meter_setup):
+  """Tests record_client_token_usage records correctly under different usage conditions."""
+  llm_request = mock.MagicMock(
+      contents=[types.Content(parts=[types.Part(text="hello")])],
+      model="test-model",
+  )
+  response = mock.MagicMock(
+      content=types.Content(parts=[types.Part(text="hello response")]),
+      model_version="test-model-v1",
+      usage_metadata=types.GenerateContentResponseUsageMetadata(
+          prompt_token_count=20,
+          candidates_token_count=30,
+          tool_use_prompt_token_count=5,
+          thoughts_token_count=10,
+      ),
+  )
+  _metrics.record_client_token_usage(
+      agent_name="test_agent",
+      llm_request=llm_request,
+      responses=[response],
+  )
+  client_token_usage_hist = mock_meter_setup["client_token_usage"]
+  assert client_token_usage_hist.record.call_count == 2
+
+  base_attributes = {
+      "gen_ai.agent.name": "test_agent",
+      "gen_ai.operation.name": "generate_content",
+      "gen_ai.provider.name": "gemini",
+      "gen_ai.request.model": "test-model",
+      "gen_ai.response.model": "test-model-v1",
+  }
+
+  input_call = None
+  output_call = None
+
+  for args, kwargs in client_token_usage_hist.record.call_args_list:
+    token_type = kwargs.get("attributes", {}).get("gen_ai.token.type")
+    if token_type == "input":
+      input_call = (args, kwargs)
+    elif token_type == "output":
+      output_call = (args, kwargs)
+
+  assert input_call is not None, "Missing 'input' token usage record"
+  assert output_call is not None, "Missing 'output' token usage record"
+
+  # Verify input tokens (prompt_token_count + tool_use_prompt_token_count)
+  assert input_call[0][0] == 25
+  assert input_call[1]["attributes"] == base_attributes | {
+      "gen_ai.token.type": "input"
+  }
+
+  # Verify output tokens (candidates_token_count + thoughts_token_count)
+  assert output_call[0][0] == 40
+  assert output_call[1]["attributes"] == base_attributes | {
+      "gen_ai.token.type": "output"
+  }

tests/unittests/telemetry/test_node_functional.py

@@ -0,0 +1,197 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from google.adk.telemetry import tracing
+from opentelemetry.sdk._logs.export import InMemoryLogRecordExporter
+from opentelemetry.sdk.metrics.export import InMemoryMetricReader
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+import pytest
+
+from .functional_node_test_cases import ALL_NODE_CASES
+from .functional_test_helpers import aclosing_wrapping_assertions
+from .functional_test_helpers import install_telemetry
+from .functional_test_helpers import run_node_scenario
+from .functional_test_helpers import TelemetryDigest
+
+if TYPE_CHECKING:
+  from google.adk.events.event import Event
+  from opentelemetry.sdk.trace import ReadableSpan
+
+  from .functional_test_helpers import FunctionalTestCase
+
+
+@pytest.mark.parametrize('case', ALL_NODE_CASES, ids=lambda c: c.test_id)
+@pytest.mark.asyncio
+async def test_telemetry_schema(
+    case: FunctionalTestCase,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+  """Tests creation of multiple spans/logs in an E2E runner invocation with a
+
+  workflow.
+
+  Asserts the entire telemetry schema (spans + attributes + per-span logs)
+  matches the hand-written expected shape for the given semconv +
+  content-capture configuration.
+  """
+  case.apply_env(monkeypatch)
+  span_exporter = InMemorySpanExporter()
+  log_exporter = InMemoryLogRecordExporter()
+  metric_reader = InMemoryMetricReader()
+  install_telemetry(monkeypatch, span_exporter, log_exporter, metric_reader)
+
+  events = await run_node_scenario()
+  spans = span_exporter.get_finished_spans()
+  digest = TelemetryDigest.build(
+      spans, log_exporter.get_finished_logs(), metric_reader.get_metrics_data()
+  )
+
+  assert digest == case.expected
+  _verify_associated_events(spans, events)
+
+
+@pytest.mark.asyncio
+async def test_async_generators_wrapped_in_aclosing(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+  """Asserts each async generator iterated by the scenario is wrapped in ``aclosing``.
+
+  Necessary because instrumentation utilizes contextvars, which run into
+  "ContextVar was created in a different Context" errors when a given
+  coroutine gets indeterminately suspended.
+
+  Kept as a single non-parametrized test because the underlying
+  ``gc.get_referrers`` walk is expensive (~5 seconds per scenario).
+  """
+  install_telemetry(
+      monkeypatch,
+      InMemorySpanExporter(),
+      InMemoryLogRecordExporter(),
+      InMemoryMetricReader(),
+  )
+
+  with aclosing_wrapping_assertions():
+    _ = await run_node_scenario()
+
+
+def _verify_associated_events(
+    spans: tuple[ReadableSpan, ...], events: list[Event]
+):
+  def _nodelike_name(span: ReadableSpan) -> str:
+    for prefix in ['invoke_node ', 'invoke_workflow ', 'invoke_agent ']:
+      if span.name.startswith(prefix):
+        return span.name.replace(prefix, '')
+    return ''
+
+  def _emitting_node_name(event: Event) -> str:
+    # Strip out
+    # 1. Path except for the last node (everything before "/")
+    # 2. Retry count (everything after "@")
+    return event.node_info.path.split('/')[-1].split('@')[0]
+
+  events_by_id = {event.id: event for event in events}
+  for span in spans:
+    if not span.attributes:
+      continue
+    associated_ids = span.attributes.get(
+        'gcp.vertex.agent.associated_event_ids', None
+    )
+    if associated_ids is None:
+      continue
+    assert isinstance(associated_ids, tuple)
+    assert len(associated_ids) > 0, f'Span name {span.name} emitted no events'
+    for event_id in associated_ids:
+      event = events_by_id[str(event_id)]
+      assert _nodelike_name(span) == _emitting_node_name(event)
+
+
+@pytest.mark.asyncio
+async def test_exception_preserves_attributes(
+    monkeypatch: pytest.MonkeyPatch,
+):
+  """Test when an exception occurs during tool execution, span attributes are still present on spans where they are expected."""
+
+  span_exporter = InMemorySpanExporter()
+  install_telemetry(
+      monkeypatch,
+      span_exporter,
+      InMemoryLogRecordExporter(),
+      InMemoryMetricReader(),
+  )
+
+  captured_events: list[Event] = []
+  with pytest.raises(ValueError, match='This tool always fails'):
+    await run_node_scenario(failing=True, event_sink=captured_events)
+
+  # Assert
+  spans = span_exporter.get_finished_spans()
+  _verify_associated_events(spans, captured_events)
+  spans_by_name = {span.name: span for span in spans}
+
+  assert 'execute_tool some_tool' in spans_by_name
+  tool_span = spans_by_name['execute_tool some_tool']
+
+  attrs = dict(tool_span.attributes)
+  # Dynamic ID
+  tool_call_id = attrs.get('gen_ai.tool.call.id')
+
+  assert dict(tool_span.attributes) == {
+      'gen_ai.operation.name': 'execute_tool',
+      'gen_ai.tool.name': 'some_tool',
+      'gen_ai.tool.description': 'A sample tool.',
+      'gen_ai.tool.type': 'FunctionTool',
+      'error.type': 'ValueError',
+      'gcp.vertex.agent.llm_request': '{}',
+      'gcp.vertex.agent.llm_response': '{}',
+      'gcp.vertex.agent.tool_call_args': '{"arg1": "val1"}',
+      'gen_ai.tool.call.id': tool_call_id,
+      'gcp.vertex.agent.tool_response': '{"result": "<not specified>"}',
+  }
+
+
+@pytest.mark.asyncio
+async def test_no_generate_content_for_gemini_model_when_already_instrumented(
+    monkeypatch: pytest.MonkeyPatch,
+):
+  """Tests that generate_content span is not created if already instrumented."""
+
+  span_exporter = InMemorySpanExporter()
+  install_telemetry(
+      monkeypatch,
+      span_exporter,
+      InMemoryLogRecordExporter(),
+      InMemoryMetricReader(),
+  )
+
+  # Arrange
+  monkeypatch.setattr(
+      tracing,
+      '_instrumented_with_opentelemetry_instrumentation_google_genai',
+      lambda: True,
+  )
+  monkeypatch.setattr(
+      tracing,
+      '_is_gemini_agent',
+      lambda _: True,
+  )
+
+  _ = await run_node_scenario()
+
+  # Assert
+  spans = span_exporter.get_finished_spans()
+  assert not any(span.name.startswith('generate_content') for span in spans)

tests/unittests/telemetry/test_setup.py

@@ -0,0 +1,119 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from unittest import mock
+
+from google.adk.telemetry.setup import maybe_set_otel_providers
+import pytest
+
+
+@pytest.fixture
+def mock_os_environ():
+  initial_env = os.environ.copy()
+  with mock.patch.dict(os.environ, initial_env, clear=False) as m:
+    yield m
+
+
+@pytest.mark.parametrize(
+    "env_vars, should_setup_trace, should_setup_metrics, should_setup_logs",
+    [
+        (
+            {"OTEL_EXPORTER_OTLP_TRACES_ENDPOINT": "some-endpoint"},
+            True,
+            False,
+            False,
+        ),
+        (
+            {"OTEL_EXPORTER_OTLP_METRICS_ENDPOINT": "some-endpoint"},
+            False,
+            True,
+            False,
+        ),
+        (
+            {"OTEL_EXPORTER_OTLP_LOGS_ENDPOINT": "some-endpoint"},
+            False,
+            False,
+            True,
+        ),
+        (
+            {
+                "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT": "some-endpoint",
+                "OTEL_EXPORTER_OTLP_METRICS_ENDPOINT": "some-endpoint",
+                "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT": "some-endpoint",
+            },
+            True,
+            True,
+            True,
+        ),
+        (
+            {"OTEL_EXPORTER_OTLP_ENDPOINT": "some-endpoint"},
+            True,
+            True,
+            True,
+        ),
+    ],
+)
+def test_maybe_set_otel_providers(
+    env_vars: dict[str, str],
+    should_setup_trace: bool,
+    should_setup_metrics: bool,
+    should_setup_logs: bool,
+    monkeypatch: pytest.MonkeyPatch,
+    mock_os_environ,  # pylint: disable=unused-argument,redefined-outer-name
+):
+  """
+  Test initializing correct providers in setup_otel
+  when providing OTel env variables.
+  """
+  # Arrange.
+  for k, v in env_vars.items():
+    monkeypatch.setenv(k, v)
+  trace_provider_mock = mock.MagicMock()
+  monkeypatch.setattr(
+      "opentelemetry.trace.set_tracer_provider",
+      trace_provider_mock,
+  )
+  meter_provider_mock = mock.MagicMock()
+  monkeypatch.setattr(
+      "opentelemetry.metrics.set_meter_provider",
+      meter_provider_mock,
+  )
+  logs_provider_mock = mock.MagicMock()
+  monkeypatch.setattr(
+      "opentelemetry._logs.set_logger_provider",
+      logs_provider_mock,
+  )
+  monkeypatch.setattr(
+      "google.adk.telemetry.setup._get_otel_span_exporter",
+      lambda: mock.MagicMock(),
+  )
+  monkeypatch.setattr(
+      "google.adk.telemetry.setup._get_otel_metrics_exporter",
+      lambda: mock.MagicMock(),
+  )
+  monkeypatch.setattr(
+      "google.adk.telemetry.setup._get_otel_logs_exporter",
+      lambda: mock.MagicMock(),
+  )
+
+  # Act.
+  maybe_set_otel_providers()
+
+  # Assert.
+  # If given telemetry type was enabled,
+  # the corresponding provider should be set.
+  assert trace_provider_mock.call_count == (1 if should_setup_trace else 0)
+  assert meter_provider_mock.call_count == (1 if should_setup_metrics else 0)
+  assert logs_provider_mock.call_count == (1 if should_setup_logs else 0)

tests/unittests/telemetry/test_sqlite_span_exporter.py

@@ -0,0 +1,462 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from google.adk.telemetry.sqlite_span_exporter import SqliteSpanExporter
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExportResult
+from opentelemetry.trace import SpanContext
+from opentelemetry.trace import TraceFlags
+from opentelemetry.trace import TraceState
+
+
+def _create_span(
+    *,
+    span_id: int = 0x00000000000ABC12,
+    trace_id: int = 0x000000000000000000000000000DEF45,
+    parent_span_id: int | None = None,
+    name: str = "test_span",
+    attributes: dict | None = None,
+    start_time: int = 1000,
+    end_time: int = 2000,
+) -> ReadableSpan:
+  """Helper to create ReadableSpan instances for testing."""
+  context = SpanContext(
+      trace_id=trace_id,
+      span_id=span_id,
+      is_remote=False,
+      trace_flags=TraceFlags(TraceFlags.SAMPLED),
+      trace_state=TraceState(),
+  )
+
+  parent = None
+  if parent_span_id is not None:
+    parent = SpanContext(
+        trace_id=trace_id,
+        span_id=parent_span_id,
+        is_remote=False,
+        trace_flags=TraceFlags(TraceFlags.SAMPLED),
+        trace_state=TraceState(),
+    )
+
+  return ReadableSpan(
+      name=name,
+      context=context,
+      parent=parent,
+      attributes=attributes or {},
+      start_time=start_time,
+      end_time=end_time,
+  )
+
+
+def test_export_single_span_returns_success(tmp_path):
+  db_path = tmp_path / "test.db"
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+
+  span = _create_span(
+      name="test_operation",
+      attributes={"gcp.vertex.agent.session_id": "session-123"},
+  )
+
+  result = exporter.export([span])
+
+  assert result == SpanExportResult.SUCCESS
+  assert db_path.exists()
+
+
+def test_export_empty_list_returns_success(tmp_path):
+  db_path = tmp_path / "test.db"
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+
+  result = exporter.export([])
+
+  assert result == SpanExportResult.SUCCESS
+
+
+def test_get_all_spans_for_session_returns_matching_spans(tmp_path):
+  db_path = tmp_path / "test.db"
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+
+  span1 = _create_span(
+      span_id=0x111,
+      trace_id=0xAAA111,  # Different trace for session-123
+      attributes={"gcp.vertex.agent.session_id": "session-123"},
+      name="span1",
+  )
+  span2 = _create_span(
+      span_id=0x222,
+      trace_id=0xAAA222,  # Different trace for session-123
+      attributes={"gcp.vertex.agent.session_id": "session-123"},
+      name="span2",
+  )
+  span3 = _create_span(
+      span_id=0x333,
+      trace_id=0xBBB333,  # Different trace for session-456
+      attributes={"gcp.vertex.agent.session_id": "session-456"},
+      name="span3",
+  )
+
+  exporter.export([span1, span2, span3])
+
+  result = exporter.get_all_spans_for_session("session-123")
+
+  assert len(result) == 2
+  names = [span.name for span in result]
+  assert "span1" in names
+  assert "span2" in names
+  assert "span3" not in names
+
+
+def test_get_all_spans_for_session_includes_sibling_spans_without_session_id(
+    tmp_path,
+):
+  db_path = tmp_path / "test.db"
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+
+  # Parent span without session_id (e.g., invocation span)
+  parent_span = _create_span(
+      span_id=0x100,
+      trace_id=0xAAA,
+      name="invocation",
+      attributes={},  # No session_id
+  )
+
+  # Child span with session_id
+  child_span = _create_span(
+      span_id=0x200,
+      trace_id=0xAAA,  # Same trace
+      parent_span_id=0x100,
+      name="call_llm",
+      attributes={"gcp.vertex.agent.session_id": "session-789"},
+  )
+
+  # Sibling span without session_id (should be included)
+  sibling_span = _create_span(
+      span_id=0x300,
+      trace_id=0xAAA,  # Same trace
+      parent_span_id=0x100,
+      name="tool_call",
+      attributes={},  # No session_id
+  )
+
+  # Unrelated span with different trace_id (should not be included)
+  unrelated_span = _create_span(
+      span_id=0x400,
+      trace_id=0xBBB,  # Different trace
+      name="unrelated",
+      attributes={},
+  )
+
+  exporter.export([parent_span, child_span, sibling_span, unrelated_span])
+
+  result = exporter.get_all_spans_for_session("session-789")
+
+  assert len(result) == 3
+  names = [span.name for span in result]
+  assert "invocation" in names
+  assert "call_llm" in names
+  assert "tool_call" in names
+  assert "unrelated" not in names
+
+
+def test_get_all_spans_for_unknown_session_returns_empty_list(tmp_path):
+  db_path = tmp_path / "test.db"
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+
+  span = _create_span(
+      attributes={"gcp.vertex.agent.session_id": "session-123"},
+  )
+  exporter.export([span])
+
+  result = exporter.get_all_spans_for_session("unknown-session")
+
+  assert result == []
+
+
+def test_round_trip_preserves_span_attributes(tmp_path):
+  db_path = tmp_path / "test.db"
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+
+  original_attributes = {
+      "gcp.vertex.agent.session_id": "session-123",
+      "gcp.vertex.agent.invocation_id": "invocation-456",
+      "gen_ai.conversation.id": "conv-789",
+      "custom.attribute": "test_value",
+      "numeric.value": 42,
+      "boolean.value": True,
+      "list.value": [1, 2, 3],
+      "dict.value": {"nested": "data"},
+  }
+
+  original_span = _create_span(
+      span_id=0x12345678,
+      trace_id=0xABCDEF123456789,
+      name="test_operation",
+      attributes=original_attributes,
+      start_time=1000000,
+      end_time=2000000,
+  )
+
+  exporter.export([original_span])
+
+  retrieved_spans = exporter.get_all_spans_for_session("session-123")
+
+  assert len(retrieved_spans) == 1
+  retrieved = retrieved_spans[0]
+
+  assert retrieved.name == "test_operation"
+  assert retrieved.context.span_id == 0x12345678
+  assert retrieved.context.trace_id == 0xABCDEF123456789
+  assert retrieved.start_time == 1000000
+  assert retrieved.end_time == 2000000
+  assert retrieved.attributes == original_attributes
+
+
+def test_spans_with_parent_context_exported_correctly(tmp_path):
+  db_path = tmp_path / "test.db"
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+
+  parent_span = _create_span(
+      span_id=0xAAA,
+      trace_id=0x123,
+      name="parent",
+      attributes={"gcp.vertex.agent.session_id": "session-001"},
+  )
+
+  child_span = _create_span(
+      span_id=0xBBB,
+      trace_id=0x123,
+      parent_span_id=0xAAA,
+      name="child",
+      attributes={"gcp.vertex.agent.session_id": "session-001"},
+  )
+
+  exporter.export([parent_span, child_span])
+
+  retrieved_spans = exporter.get_all_spans_for_session("session-001")
+
+  assert len(retrieved_spans) == 2
+
+  # Find child span in results
+  child = next(s for s in retrieved_spans if s.name == "child")
+  assert child.parent is not None
+  assert child.parent.span_id == 0xAAA
+  assert child.parent.trace_id == 0x123
+
+  # Find parent span in results
+  parent = next(s for s in retrieved_spans if s.name == "parent")
+  assert parent.parent is None
+
+
+def test_shutdown_closes_connection(tmp_path):
+  db_path = tmp_path / "test.db"
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+
+  # Create a span to ensure connection is open
+  span = _create_span()
+  exporter.export([span])
+
+  # Verify connection exists
+  assert exporter._conn is not None
+
+  exporter.shutdown()
+
+  # Verify connection is closed
+  assert exporter._conn is None
+
+
+def test_force_flush_returns_true(tmp_path):
+  db_path = tmp_path / "test.db"
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+
+  result = exporter.force_flush()
+
+  assert result is True
+
+  # Also test with timeout parameter
+  result_with_timeout = exporter.force_flush(timeout_millis=5000)
+  assert result_with_timeout is True
+
+
+def test_export_handles_spans_with_none_attributes(tmp_path):
+  db_path = tmp_path / "test.db"
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+
+  span = _create_span(attributes=None)
+
+  result = exporter.export([span])
+
+  assert result == SpanExportResult.SUCCESS
+
+  # Verify the span was stored correctly
+  rows = exporter._query("SELECT attributes_json FROM spans", [])
+  assert len(rows) == 1
+  attributes_json = rows[0]["attributes_json"]
+  assert json.loads(attributes_json) == {}
+
+
+def test_duplicate_span_id_replaces_previous_row(tmp_path):
+  db_path = tmp_path / "test.db"
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+
+  # Export first version of span
+  span1 = _create_span(
+      span_id=0x999,
+      name="first_version",
+      attributes={"version": 1, "gcp.vertex.agent.session_id": "session-dup"},
+  )
+  exporter.export([span1])
+
+  # Export second version with same span_id
+  span2 = _create_span(
+      span_id=0x999,
+      name="second_version",
+      attributes={"version": 2, "gcp.vertex.agent.session_id": "session-dup"},
+  )
+  exporter.export([span2])
+
+  # Verify only one row exists with updated data
+  retrieved_spans = exporter.get_all_spans_for_session("session-dup")
+  assert len(retrieved_spans) == 1
+  assert retrieved_spans[0].name == "second_version"
+  assert retrieved_spans[0].attributes["version"] == 2
+
+
+def test_non_serializable_attributes_use_fallback(tmp_path):
+  db_path = tmp_path / "test.db"
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+
+  # Create a non-serializable object
+  class NonSerializable:
+    pass
+
+  attributes = {
+      "gcp.vertex.agent.session_id": "session-nonser",
+      "normal_attr": "value",
+      "non_serializable": NonSerializable(),
+  }
+
+  span = _create_span(attributes=attributes)
+
+  result = exporter.export([span])
+
+  assert result == SpanExportResult.SUCCESS
+
+  # Verify the span was stored and non-serializable attribute has fallback
+  retrieved_spans = exporter.get_all_spans_for_session("session-nonser")
+  assert len(retrieved_spans) == 1
+  assert retrieved_spans[0].attributes["normal_attr"] == "value"
+  assert (
+      retrieved_spans[0].attributes["non_serializable"] == "<not serializable>"
+  )
+
+
+def test_export_multiple_spans_in_batch(tmp_path):
+  db_path = tmp_path / "test.db"
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+
+  spans = [
+      _create_span(
+          span_id=i,
+          name=f"span_{i}",
+          attributes={"gcp.vertex.agent.session_id": "batch-session"},
+      )
+      for i in range(10)
+  ]
+
+  result = exporter.export(spans)
+
+  assert result == SpanExportResult.SUCCESS
+
+  retrieved_spans = exporter.get_all_spans_for_session("batch-session")
+  assert len(retrieved_spans) == 10
+  names = {span.name for span in retrieved_spans}
+  assert names == {f"span_{i}" for i in range(10)}
+
+
+def test_export_with_alternative_session_id_attribute(tmp_path):
+  db_path = tmp_path / "test.db"
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+
+  # Test using gen_ai.conversation.id as fallback for session_id
+  span = _create_span(
+      attributes={"gen_ai.conversation.id": "conv-session-123"},
+  )
+
+  exporter.export([span])
+
+  # Should be queryable by the conversation id
+  result = exporter.get_all_spans_for_session("conv-session-123")
+
+  assert len(result) == 1
+  assert result[0].attributes["gen_ai.conversation.id"] == "conv-session-123"
+
+
+def test_deserialize_handles_invalid_json(tmp_path):
+  db_path = tmp_path / "test.db"
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+
+  # Manually insert a row with invalid JSON
+  conn = exporter._get_connection()
+  conn.execute(
+      "INSERT INTO spans (span_id, trace_id, name, attributes_json) VALUES (?,"
+      " ?, ?, ?)",
+      ("abc123", "def456", "test", "not valid json"),
+  )
+  conn.commit()
+
+  # Try to retrieve the span - should not raise, but attributes should be empty
+  rows = exporter._query("SELECT * FROM spans", [])
+  span = exporter._row_to_readable_span(rows[0])
+
+  assert span.name == "test"
+  assert span.attributes == {}
+
+
+def test_get_spans_ordered_by_start_time(tmp_path):
+  db_path = tmp_path / "test.db"
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+
+  # Create spans with different start times
+  spans = [
+      _create_span(
+          span_id=0x300,
+          start_time=3000,
+          attributes={"gcp.vertex.agent.session_id": "session-order"},
+      ),
+      _create_span(
+          span_id=0x100,
+          start_time=1000,
+          attributes={"gcp.vertex.agent.session_id": "session-order"},
+      ),
+      _create_span(
+          span_id=0x200,
+          start_time=2000,
+          attributes={"gcp.vertex.agent.session_id": "session-order"},
+      ),
+  ]
+
+  exporter.export(spans)
+
+  result = exporter.get_all_spans_for_session("session-order")
+
+  # Verify spans are ordered by start_time
+  assert len(result) == 3
+  assert result[0].context.span_id == 0x100
+  assert result[1].context.span_id == 0x200
+  assert result[2].context.span_id == 0x300

tests/unittests/telemetry/test_token_usage.py

@@ -0,0 +1,221 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.adk.telemetry import _token_usage
+from google.genai import types
+import pytest
+
+
+@pytest.fixture(name="usage_metadata")
+def fixture_usage_metadata() -> types.GenerateContentResponseUsageMetadata:
+  """Provides a baseline GenerateContentResponseUsageMetadata fixture with all token counts initialized to None."""
+  m = types.GenerateContentResponseUsageMetadata()
+  m.prompt_token_count = None
+  m.tool_use_prompt_token_count = None
+  m.candidates_token_count = None
+  m.thoughts_token_count = None
+  m.cached_content_token_count = None
+  return m
+
+
+def test_input_token_count_all_present(
+    usage_metadata: types.GenerateContentResponseUsageMetadata,
+):
+  """Tests input_token_count when all components are present."""
+  usage_metadata.prompt_token_count = 10
+  usage_metadata.tool_use_prompt_token_count = 5
+  token_usage = _token_usage.TokenUsage(usage_metadata)
+  assert token_usage.input_token_count == 15
+
+
+def test_input_token_count_only_prompt(
+    usage_metadata: types.GenerateContentResponseUsageMetadata,
+):
+  """Tests input_token_count when only prompt_token_count is present."""
+  usage_metadata.prompt_token_count = 10
+  usage_metadata.tool_use_prompt_token_count = None
+  token_usage = _token_usage.TokenUsage(usage_metadata)
+  assert token_usage.input_token_count == 10
+
+
+def test_input_token_count_only_tool(
+    usage_metadata: types.GenerateContentResponseUsageMetadata,
+):
+  """Tests input_token_count when only tool_use_prompt_token_count is present."""
+  usage_metadata.prompt_token_count = None
+  usage_metadata.tool_use_prompt_token_count = 5
+  token_usage = _token_usage.TokenUsage(usage_metadata)
+  assert token_usage.input_token_count == 5
+
+
+def test_input_token_count_none(
+    usage_metadata: types.GenerateContentResponseUsageMetadata,
+):
+  """Tests input_token_count when all components are None."""
+  usage_metadata.prompt_token_count = None
+  usage_metadata.tool_use_prompt_token_count = None
+  token_usage = _token_usage.TokenUsage(usage_metadata)
+  assert token_usage.input_token_count is None
+
+
+def test_input_token_count_zero(
+    usage_metadata: types.GenerateContentResponseUsageMetadata,
+):
+  """Tests input_token_count when all components are zero."""
+  usage_metadata.prompt_token_count = 0
+  usage_metadata.tool_use_prompt_token_count = 0
+  token_usage = _token_usage.TokenUsage(usage_metadata)
+  assert token_usage.input_token_count == 0
+
+
+def test_input_token_count_metadata_none():
+  """Tests input_token_count when usage_metadata is None."""
+  token_usage = _token_usage.TokenUsage(None)
+  assert token_usage.input_token_count is None
+
+
+def test_input_token_count_missing_tool_use_attr():
+  """Tests input_token_count when tool_use_prompt_token_count is missing."""
+  token_usage = _token_usage.TokenUsage(
+      types.GenerateContentResponseUsageMetadata(prompt_token_count=10)
+  )
+  assert token_usage.input_token_count == 10
+
+
+def test_output_token_count_all_present(
+    usage_metadata: types.GenerateContentResponseUsageMetadata,
+):
+  """Tests output_token_count when all components are present."""
+  usage_metadata.candidates_token_count = 20
+  usage_metadata.thoughts_token_count = 8
+  token_usage = _token_usage.TokenUsage(usage_metadata)
+  assert token_usage.output_token_count == 28
+
+
+def test_output_token_count_only_candidates(
+    usage_metadata: types.GenerateContentResponseUsageMetadata,
+):
+  """Tests output_token_count when only candidates_token_count is present."""
+  usage_metadata.candidates_token_count = 20
+  usage_metadata.thoughts_token_count = None
+  token_usage = _token_usage.TokenUsage(usage_metadata)
+  assert token_usage.output_token_count == 20
+
+
+def test_output_token_count_only_thoughts(
+    usage_metadata: types.GenerateContentResponseUsageMetadata,
+):
+  """Tests output_token_count when only thoughts_token_count is present."""
+  usage_metadata.candidates_token_count = None
+  usage_metadata.thoughts_token_count = 8
+  token_usage = _token_usage.TokenUsage(usage_metadata)
+  assert token_usage.output_token_count == 8
+
+
+def test_output_token_count_none(
+    usage_metadata: types.GenerateContentResponseUsageMetadata,
+):
+  """Tests output_token_count when all components are None."""
+  usage_metadata.candidates_token_count = None
+  usage_metadata.thoughts_token_count = None
+  token_usage = _token_usage.TokenUsage(usage_metadata)
+  assert token_usage.output_token_count is None
+
+
+def test_output_token_count_zero(
+    usage_metadata: types.GenerateContentResponseUsageMetadata,
+):
+  """Tests output_token_count when all components are zero."""
+  usage_metadata.candidates_token_count = 0
+  usage_metadata.thoughts_token_count = 0
+  token_usage = _token_usage.TokenUsage(usage_metadata)
+  assert token_usage.output_token_count == 0
+
+
+def test_output_token_count_metadata_none():
+  """Tests output_token_count when usage_metadata is None."""
+  token_usage = _token_usage.TokenUsage(None)
+  assert token_usage.output_token_count is None
+
+
+def test_to_attributes_full(
+    usage_metadata: types.GenerateContentResponseUsageMetadata,
+):
+  """Tests to_attributes with all attributes present."""
+  usage_metadata.prompt_token_count = 10
+  usage_metadata.tool_use_prompt_token_count = 5
+  usage_metadata.candidates_token_count = 20
+  usage_metadata.thoughts_token_count = 8
+  usage_metadata.cached_content_token_count = 100
+
+  token_usage = _token_usage.TokenUsage(usage_metadata)
+  attrs = token_usage.to_attributes()
+  assert attrs[_token_usage.GEN_AI_USAGE_INPUT_TOKENS] == 15
+  assert attrs[_token_usage.GEN_AI_USAGE_OUTPUT_TOKENS] == 28
+  assert attrs[_token_usage.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS] == 100
+  assert attrs[_token_usage.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS] == 8
+
+
+def test_to_attributes_partial(
+    usage_metadata: types.GenerateContentResponseUsageMetadata,
+):
+  """Tests to_attributes with only some attributes present."""
+  usage_metadata.prompt_token_count = 10
+  usage_metadata.tool_use_prompt_token_count = None
+  usage_metadata.candidates_token_count = None
+  usage_metadata.thoughts_token_count = None
+  usage_metadata.cached_content_token_count = None
+
+  token_usage = _token_usage.TokenUsage(usage_metadata)
+  attrs = token_usage.to_attributes()
+  assert attrs[_token_usage.GEN_AI_USAGE_INPUT_TOKENS] == 10
+  assert _token_usage.GEN_AI_USAGE_OUTPUT_TOKENS not in attrs
+  assert _token_usage.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS not in attrs
+  assert _token_usage.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS not in attrs
+
+
+def test_to_attributes_metadata_none():
+  """Tests to_attributes when usage_metadata is None."""
+  token_usage = _token_usage.TokenUsage(None)
+  assert token_usage.to_attributes() == {}
+
+
+def test_to_attributes_with_zeros(
+    usage_metadata: types.GenerateContentResponseUsageMetadata,
+):
+  """Tests to_attributes when all attributes are zero."""
+  usage_metadata.prompt_token_count = 0
+  usage_metadata.tool_use_prompt_token_count = 0
+  usage_metadata.candidates_token_count = 0
+  usage_metadata.thoughts_token_count = 0
+  usage_metadata.cached_content_token_count = 0
+
+  token_usage = _token_usage.TokenUsage(usage_metadata)
+  attrs = token_usage.to_attributes()
+  assert attrs[_token_usage.GEN_AI_USAGE_INPUT_TOKENS] == 0
+  assert attrs[_token_usage.GEN_AI_USAGE_OUTPUT_TOKENS] == 0
+  assert attrs[_token_usage.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS] == 0
+  assert attrs[_token_usage.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS] == 0
+
+
+def test_to_attributes_missing_optional_attrs():
+  """Tests to_attributes when optional attributes are missing from metadata object."""
+  token_usage = _token_usage.TokenUsage(
+      types.GenerateContentResponseUsageMetadata(
+          prompt_token_count=10, candidates_token_count=20
+      )
+  )
+  attrs = token_usage.to_attributes()
+  assert attrs[_token_usage.GEN_AI_USAGE_INPUT_TOKENS] == 10
+  assert attrs[_token_usage.GEN_AI_USAGE_OUTPUT_TOKENS] == 20