diff --git a/src/google/adk/telemetry/_instrumentation.py b/src/google/adk/telemetry/_instrumentation.py index 93cab277fb..39df7f49ea 100644 --- a/src/google/adk/telemetry/_instrumentation.py +++ b/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 diff --git a/tests/unittests/telemetry/test_instrumentation.py b/tests/unittests/telemetry/test_instrumentation.py index bc0838e55a..219128864a 100644 --- a/tests/unittests/telemetry/test_instrumentation.py +++ b/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"]