fix(tracing): preserve subagent trace context

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 == {}

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)