chore: import upstream snapshot with attribution

2026-07-13 13:25:13 +08:00
commit ec2b666284
2231 changed files with 491535 additions and 0 deletions
@@ -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