fix: avoid OTel "Failed to detach context" error on early run_async close

src/google/adk/telemetry/_instrumentation.py

@@ -68,8 +68,24 @@ def record_invocation(
     Nothing; the span (if any) is active for the duration of the block.
   """
   if resolve_schema_version() < SCHEMA_VERSION_SEMCONV_ALIGNED:
-    with tracing.tracer.start_as_current_span("invocation"):
+    # This context manager wraps runners._run_node_async, an async generator.
+    # When a caller stops iterating early, that generator is finalized
+    # (GeneratorExit / CancelledError) in a different execution context than the
+    # one where the span was attached. start_as_current_span's automatic
+    # detach() would then raise "Token was created in a different Context" --
+    # OpenTelemetry swallows it but logs it at ERROR on every early-terminated
+    # run. Manage the span/context explicitly and detach only on normal
+    # completion; always end the span so trace data stays complete.
+    span = tracing.tracer.start_span("invocation")
+    token = context_api.attach(trace.set_span_in_context(span))
+    completed = False
+    try:
       yield
+      completed = True
+    finally:
+      if completed:
+        context_api.detach(token)
+      span.end()
     return
 
   from . import node_tracing

tests/unittests/telemetry/test_instrumentation.py

@@ -14,12 +14,20 @@
 
 # pylint: disable=protected-access
 
+import asyncio
+import logging
 import time
 from unittest import mock
 
 from google.adk.telemetry import _instrumentation
 from google.adk.telemetry import _metrics
+from google.adk.telemetry import tracing
 from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
+    InMemorySpanExporter,
+)
 import pytest
 
 
@@ -106,3 +114,102 @@ async def test_record_tool_execution_forwards_detected_error_type():
   mock_record.assert_called_once()
   assert mock_record.call_args.kwargs["error"] is None
   assert mock_record.call_args.kwargs["error_type"] == "MCP_TOOL_ERROR"
+
+
+def _install_v1_tracer(
+    monkeypatch: pytest.MonkeyPatch,
+) -> InMemorySpanExporter:
+  """Installs an in-memory SDK tracer and forces telemetry schema v1."""
+  # Schema v1 (the legacy ``invocation`` span) is the default off Agent Engine.
+  monkeypatch.delenv("ADK_TELEMETRY_SCHEMA_VERSION_OPT_IN", raising=False)
+  monkeypatch.delenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", raising=False)
+
+  exporter = InMemorySpanExporter()
+  provider = TracerProvider()
+  provider.add_span_processor(SimpleSpanProcessor(exporter))
+  real_tracer = provider.get_tracer(__name__)
+  monkeypatch.setattr(tracing.tracer, "start_span", real_tracer.start_span)
+  return exporter
+
+
+async def _event_queue():
+  for i in range(10):
+    await asyncio.sleep(0)
+    yield i
+
+
+def _make_run_node():
+  """Returns an async generator mirroring runners._run_node_async."""
+
+  async def _run_node_async():
+    with _instrumentation.record_invocation(
+        entrypoint_node=None, conversation_id="conversation-id"
+    ):
+      async for event in _event_queue():
+        yield event
+
+  return _run_node_async
+
+
+def test_record_invocation_no_detach_error_on_early_close(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+):
+  """Early caller stop must not log an OTel 'Failed to detach context' ERROR.
+
+  Regression test: when a caller stops iterating runners._run_node_async early,
+  the async generator is finalized in a different execution context than the one
+  where the invocation span was attached. Automatically detaching there raises
+  "Token was created in a different Context" (swallowed by OpenTelemetry but
+  logged at ERROR on every early-terminated run).
+  """
+  exporter = _install_v1_tracer(monkeypatch)
+  run_node = _make_run_node()
+
+  async def _caller_early_stop():
+    async for event in run_node():
+      if event == 2:
+        return  # leave generator open -> closed later in a different context
+
+  caplog.set_level(logging.ERROR, logger="opentelemetry.context")
+  # asyncio.run finalizes the still-open async generator via
+  # loop.shutdown_asyncgens(), reproducing the different-context close.
+  asyncio.run(_caller_early_stop())
+
+  detach_errors = [
+      record
+      for record in caplog.records
+      if "Failed to detach context" in record.getMessage()
+  ]
+  assert not detach_errors, f"unexpected detach errors: {detach_errors}"
+
+  # The invocation span is still ended, so trace data stays complete.
+  spans = exporter.get_finished_spans()
+  assert [span.name for span in spans] == ["invocation"]
+  assert spans[0].end_time is not None
+
+
+def test_record_invocation_full_consumption_still_records_span(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+):
+  """Full consumption keeps working: span recorded, no detach error."""
+  exporter = _install_v1_tracer(monkeypatch)
+  run_node = _make_run_node()
+
+  async def _caller_full():
+    count = 0
+    async for _ in run_node():
+      count += 1
+    return count
+
+  caplog.set_level(logging.ERROR, logger="opentelemetry.context")
+  count = asyncio.run(_caller_full())
+
+  assert count == 10
+  detach_errors = [
+      record
+      for record in caplog.records
+      if "Failed to detach context" in record.getMessage()
+  ]
+  assert not detach_errors
+  spans = exporter.get_finished_spans()
+  assert [span.name for span in spans] == ["invocation"]