diff --git a/tests/tracing/test_inmemory_exporter.py b/tests/tracing/test_inmemory_exporter.py new file mode 100644 index 000000000..90a3aeaad --- /dev/null +++ b/tests/tracing/test_inmemory_exporter.py @@ -0,0 +1,178 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# 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 asyncio + +import pytest +from opentelemetry import trace +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider, sampling + +from veadk.tracing.telemetry.exporters.inmemory_exporter import ( + _InMemoryExporter, + _InMemorySpanProcessor, +) + + +def test_processor_preserves_nested_agent_trace_context(): + exporter = _InMemoryExporter() + provider = TracerProvider() + provider.add_span_processor(_InMemorySpanProcessor(exporter)) + tracer = provider.get_tracer(__name__) + initial_span = trace.get_current_span() + + invocation_span = tracer.start_span("invocation root") + try: + assert trace.get_current_span() is invocation_span + parent_span = tracer.start_span("agent_run parent") + try: + assert trace.get_current_span() is parent_span + child_span = tracer.start_span("agent_run child") + try: + assert trace.get_current_span() is child_span + assert child_span.parent == parent_span.get_span_context() + assert parent_span.parent == invocation_span.get_span_context() + assert child_span.get_span_context().trace_id == ( + parent_span.get_span_context().trace_id + ) + assert parent_span.get_span_context().trace_id == ( + invocation_span.get_span_context().trace_id + ) + finally: + child_span.end() + assert trace.get_current_span() is parent_span + finally: + parent_span.end() + assert trace.get_current_span() is invocation_span + finally: + invocation_span.end() + + assert trace.get_current_span() is initial_span + + +def test_processor_isolates_concurrent_agent_contexts(): + exporter = _InMemoryExporter() + processor = _InMemorySpanProcessor(exporter) + provider = TracerProvider() + provider.add_span_processor(processor) + tracer = provider.get_tracer(__name__) + initial_span = trace.get_current_span() + + async def exercise(): + both_started = asyncio.Event() + started = 0 + + async def run_agent(name: str): + nonlocal started + span = tracer.start_span(f"agent_run {name}") + try: + assert trace.get_current_span() is span + started += 1 + if started == 2: + both_started.set() + await both_started.wait() + await asyncio.sleep(0) + assert trace.get_current_span() is span + trace_id = span.get_span_context().trace_id + finally: + span.end() + assert trace.get_current_span() is initial_span + return trace_id + + return await asyncio.gather(run_agent("one"), run_agent("two")) + + trace_ids = asyncio.run(exercise()) + + assert trace_ids[0] != trace_ids[1] + assert trace.get_current_span() is initial_span + assert processor._context_tokens == {} + + +def test_processor_restores_context_when_agent_run_raises(): + exporter = _InMemoryExporter() + processor = _InMemorySpanProcessor(exporter) + provider = TracerProvider() + provider.add_span_processor(processor) + tracer = provider.get_tracer(__name__) + initial_span = trace.get_current_span() + + with pytest.raises(RuntimeError, match="agent failed"): + span = tracer.start_span("agent_run failing") + try: + raise RuntimeError("agent failed") + finally: + span.end() + + assert trace.get_current_span() is initial_span + assert processor._context_tokens == {} + + +def test_processor_keeps_agent_current_for_non_agent_span(): + exporter = _InMemoryExporter() + processor = _InMemorySpanProcessor(exporter) + provider = TracerProvider() + provider.add_span_processor(processor) + tracer = provider.get_tracer(__name__) + initial_span = trace.get_current_span() + + agent_span = tracer.start_span("agent_run parent") + try: + non_agent_span = tracer.start_span("call_llm") + try: + assert trace.get_current_span() is agent_span + assert non_agent_span.parent == agent_span.get_span_context() + finally: + non_agent_span.end() + assert trace.get_current_span() is agent_span + finally: + agent_span.end() + + assert trace.get_current_span() is initial_span + assert processor._context_tokens == {} + + +def test_processor_cleans_up_record_only_span(): + exporter = _InMemoryExporter() + processor = _InMemorySpanProcessor(exporter) + provider = TracerProvider( + sampler=sampling.StaticSampler(sampling.Decision.RECORD_ONLY) + ) + provider.add_span_processor(processor) + tracer = provider.get_tracer(__name__) + initial_span = trace.get_current_span() + + span = tracer.start_span("agent_run unsampled") + assert trace.get_current_span() is span + span.end() + + assert trace.get_current_span() is initial_span + assert processor._context_tokens == {} + assert exporter._spans == [] + + +def test_processor_uses_readable_span_to_clean_up_context(): + exporter = _InMemoryExporter() + processor = _InMemorySpanProcessor(exporter) + provider = TracerProvider() + provider.add_span_processor(processor) + tracer = provider.get_tracer(__name__) + initial_span = trace.get_current_span() + + span = tracer.start_span("agent_run readable") + span.end() + + assert len(exporter._spans) == 1 + assert isinstance(exporter._spans[0], ReadableSpan) + assert exporter._spans[0] is not span + assert trace.get_current_span() is initial_span + assert processor._context_tokens == {} diff --git a/veadk/tracing/telemetry/exporters/inmemory_exporter.py b/veadk/tracing/telemetry/exporters/inmemory_exporter.py index e1f5e13cd..f530f0010 100644 --- a/veadk/tracing/telemetry/exporters/inmemory_exporter.py +++ b/veadk/tracing/telemetry/exporters/inmemory_exporter.py @@ -21,6 +21,7 @@ set_value, ) from opentelemetry.sdk.trace import ReadableSpan, export +from opentelemetry.trace import set_span_in_context from typing_extensions import override from veadk.tracing.telemetry.exporters.base_exporter import BaseExporter @@ -144,6 +145,7 @@ def __init__(self, exporter: _InMemoryExporter) -> None: exporter: _InMemoryExporter instance for storing processed spans """ super().__init__(exporter) + self._context_tokens = {} def on_start(self, span, parent_context) -> None: """Handle span start events with type-specific attribute setting. @@ -165,9 +167,10 @@ def on_start(self, span, parent_context) -> None: ctx = set_value( "suppress_language_model_instrumentation", True, context=ctx ) + ctx = set_span_in_context(span, ctx) token = attach(ctx) # mount context on `invocation` root span in Google ADK - setattr(span, "_invocation_token", token) # for later detach + self._context_tokens[span.context.span_id] = token if span.name.startswith("agent_run") or span.name.startswith("invoke_agent"): span.set_attribute("gen_ai.operation.name", "agent") @@ -180,8 +183,9 @@ def on_start(self, span, parent_context) -> None: ctx = set_value( "suppress_language_model_instrumentation", True, context=ctx ) + ctx = set_span_in_context(span, ctx) token = attach(ctx) - setattr(span, "_agent_run_token", token) # for later detach + self._context_tokens[span.context.span_id] = token def on_end(self, span: ReadableSpan) -> None: """Handle span end events with proper context cleanup. @@ -193,23 +197,22 @@ def on_end(self, span: ReadableSpan) -> None: Args: span: The span that has finished execution """ - if span.context: - if not span.context.trace_flags.sampled: - return - token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True)) - try: - self.span_exporter.export((span,)) - # pylint: disable=broad-exception-caught - except Exception: - logger.exception("Exception while exporting Span.") - detach(token) - - token = getattr(span, "_invocation_token", None) - if token: - detach(token) - - token = getattr(span, "_agent_run_token", None) - if token: + if not span.context: + return + + try: + if span.context.trace_flags.sampled: + token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True)) + try: + self.span_exporter.export((span,)) + # pylint: disable=broad-exception-caught + except Exception: + logger.exception("Exception while exporting Span.") + finally: + detach(token) + finally: + token = self._context_tokens.pop(span.context.span_id, None) + if token is not None: detach(token)