From 1e6c6a8b35a4b8948304fea84410929aca12a5f3 Mon Sep 17 00:00:00 2001 From: hylin Date: Sat, 15 Aug 2026 16:03:24 +0800 Subject: [PATCH 1/3] fix(tracing): preserve subagent trace context --- tests/tracing/test_inmemory_exporter.py | 57 +++++++++++++++++++ .../telemetry/exporters/inmemory_exporter.py | 16 +++--- 2 files changed, 65 insertions(+), 8 deletions(-) create mode 100644 tests/tracing/test_inmemory_exporter.py diff --git a/tests/tracing/test_inmemory_exporter.py b/tests/tracing/test_inmemory_exporter.py new file mode 100644 index 000000000..31c6cbd05 --- /dev/null +++ b/tests/tracing/test_inmemory_exporter.py @@ -0,0 +1,57 @@ +# 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. + +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider + +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 diff --git a/veadk/tracing/telemetry/exporters/inmemory_exporter.py b/veadk/tracing/telemetry/exporters/inmemory_exporter.py index e1f5e13cd..630c2d128 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. @@ -204,12 +208,8 @@ def on_end(self, span: ReadableSpan) -> None: 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: + token = self._context_tokens.pop(span.context.span_id, None) + if token is not None: detach(token) From 73b9eec77e8fdee9bb51522c1d92a0e94fbb5ee0 Mon Sep 17 00:00:00 2001 From: hylin Date: Mon, 17 Aug 2026 23:31:50 +0800 Subject: [PATCH 2/3] fix(tracing): clean up unsampled span context --- tests/tracing/test_inmemory_exporter.py | 123 +++++++++++++++++- .../telemetry/exporters/inmemory_exporter.py | 25 ++-- 2 files changed, 136 insertions(+), 12 deletions(-) diff --git a/tests/tracing/test_inmemory_exporter.py b/tests/tracing/test_inmemory_exporter.py index 31c6cbd05..90a3aeaad 100644 --- a/tests/tracing/test_inmemory_exporter.py +++ b/tests/tracing/test_inmemory_exporter.py @@ -12,8 +12,11 @@ # 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 TracerProvider +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider, sampling from veadk.tracing.telemetry.exporters.inmemory_exporter import ( _InMemoryExporter, @@ -55,3 +58,121 @@ def test_processor_preserves_nested_agent_trace_context(): 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 630c2d128..f530f0010 100644 --- a/veadk/tracing/telemetry/exporters/inmemory_exporter.py +++ b/veadk/tracing/telemetry/exporters/inmemory_exporter.py @@ -197,17 +197,20 @@ 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) - + 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) From a6c823a9c43d9726a964a2fba37d9641cd216ec7 Mon Sep 17 00:00:00 2001 From: hylin Date: Wed, 14 Oct 2026 22:29:45 +0800 Subject: [PATCH 3/3] fix(tracing): resolve main merge conflict in inmemory exporter Reconcile #904 with upstream main's bounded deque retention refactor: keep main's evolved _InMemoryExporter/_InMemoryExporter.__init__ and for_provider, while retaining the subagent trace-context fix (set_span_in_context + ContextVar token cleanup). Update the RECORD_ONLY assertion to be type-agnostic to the deque retention. Tests: 6 passed. --- tests/tracing/test_inmemory_exporter.py | 2 +- .../telemetry/exporters/inmemory_exporter.py | 214 ++++++++++++------ 2 files changed, 140 insertions(+), 76 deletions(-) diff --git a/tests/tracing/test_inmemory_exporter.py b/tests/tracing/test_inmemory_exporter.py index 90a3aeaad..dbfa2902b 100644 --- a/tests/tracing/test_inmemory_exporter.py +++ b/tests/tracing/test_inmemory_exporter.py @@ -157,7 +157,7 @@ def test_processor_cleans_up_record_only_span(): assert trace.get_current_span() is initial_span assert processor._context_tokens == {} - assert exporter._spans == [] + assert len(exporter._spans) == 0 def test_processor_uses_readable_span_to_clean_up_context(): diff --git a/veadk/tracing/telemetry/exporters/inmemory_exporter.py b/veadk/tracing/telemetry/exporters/inmemory_exporter.py index f530f0010..849dcce5a 100644 --- a/veadk/tracing/telemetry/exporters/inmemory_exporter.py +++ b/veadk/tracing/telemetry/exporters/inmemory_exporter.py @@ -12,6 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import sys +from collections import deque +from collections.abc import Mapping +from threading import Lock from typing import Sequence from opentelemetry.context import ( @@ -20,7 +24,7 @@ detach, set_value, ) -from opentelemetry.sdk.trace import ReadableSpan, export +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider, export from opentelemetry.trace import set_span_in_context from typing_extensions import override @@ -28,67 +32,85 @@ from veadk.utils.logger import get_logger logger = get_logger(__name__) +_registration_lock = Lock() + + +def _value_size(value) -> int: + """Estimate OTel attribute data without copying or serializing large strings""" + size = sys.getsizeof(value) + if isinstance(value, Mapping): + size += sum(_value_size(k) + _value_size(v) for k, v in value.items()) + elif isinstance(value, (tuple, list)): + size += sum(_value_size(item) for item in value) + return size + + +def _span_size(span: ReadableSpan) -> int: + """Estimate retained span data; this is a cache budget, not an RSS measurement""" + # Include fixed span/container bookkeeping and all variable-size OTel fields + size = 1024 + _value_size(span.name) + _value_size(span.attributes) + for context in (span.context, span.parent): + if context: + size += _value_size(dict(context.trace_state.items())) + if span.status: + size += _value_size(span.status.description) + if span.resource: + size += _value_size(span.resource.attributes) + size += _value_size(span.resource.schema_url) + scope = span.instrumentation_scope + if scope: + size += _value_size((scope.name, scope.version, scope.schema_url)) + size += _value_size(scope.attributes) + for event in span.events or (): + size += 128 + _value_size(event.name) + _value_size(event.attributes) + for link in span.links or (): + size += 128 + _value_size(link.attributes) + size += _value_size(dict(link.context.trace_state.items())) + return size # ======== Adapted from Google ADK ======== class _InMemoryExporter(export.SpanExporter): - """Internal span exporter that stores spans in memory for local analysis and debugging. - - This exporter collects and stores OpenTelemetry spans in memory rather than - sending them to external services. It's particularly useful for development, - testing, and local trace analysis scenarios where immediate access to span - data is needed. - - Key Features: - - In-memory span storage with session-based organization - - Trace ID tracking for span correlation - - Session-to-trace mapping for efficient filtering - - Support for span retrieval by session ID - - Attributes: - _spans: List of all collected ReadableSpan objects - trace_id: Current trace identifier from the most recent span - session_trace_dict: Mapping of session IDs to their associated trace IDs - """ + """Bounded local span history, independent of cloud exporter delivery""" - def __init__(self) -> None: - """Initialize the in-memory exporter with empty storage containers.""" + def __init__( + self, max_spans: int = 1024, max_bytes: int = 16 * 1024 * 1024 + ) -> None: super().__init__() - self._spans = [] + if max_spans < 0 or max_bytes < 0: + raise ValueError("In-memory tracing limits must be nonnegative") + self._max_spans = max_spans + self._max_bytes = max_bytes + self._spans: deque[ReadableSpan] = deque() + self._span_sizes: deque[int] = deque() + self._retained_bytes = 0 + self._lock = Lock() + self._shutdown = False self.trace_id = "" - self.session_trace_dict = {} @override def export(self, spans: Sequence[ReadableSpan]) -> export.SpanExportResult: - """Export spans to in-memory storage with session tracking. - - Processes and stores spans while maintaining session-to-trace mapping - for efficient retrieval. Extracts session information from LLM call spans - to enable session-based filtering. - - Args: - spans: Sequence of ReadableSpan objects to store - - Returns: - SpanExportResult.SUCCESS: Always returns success for in-memory storage - """ - for span in spans: - if span.context: - self.trace_id = span.context.trace_id - else: - logger.warning( - f"Span context is missing, failed to get `trace_id`. span: {span}" - ) - - if span.name == "call_llm": - attributes = dict(span.attributes or {}) - session_id = attributes.get("gen_ai.session.id", None) - if session_id: - if session_id not in self.session_trace_dict: - self.session_trace_dict[session_id] = [self.trace_id] - else: - self.session_trace_dict[session_id] += [self.trace_id] - self._spans.extend(spans) + """Retain recent spans within both limits without changing their contents""" + with self._lock: + if self._shutdown: + return export.SpanExportResult.SUCCESS + for span in spans: + if span.context: + self.trace_id = span.context.trace_id + if not self._max_spans or not self._max_bytes: + continue + size = _span_size(span) + if size > self._max_bytes: + continue + while self._spans and ( + len(self._spans) >= self._max_spans + or self._retained_bytes + size > self._max_bytes + ): + self._spans.popleft() + self._retained_bytes -= self._span_sizes.popleft() + self._spans.append(span) + self._span_sizes.append(size) + self._retained_bytes += size return export.SpanExportResult.SUCCESS @override @@ -103,31 +125,49 @@ def force_flush(self, timeout_millis: int = 30000) -> bool: """ return True - def get_finished_spans(self, session_id: str): - """Retrieve all spans associated with a specific session ID. - - Filters stored spans to return only those belonging to the specified - session, enabling session-scoped trace analysis and debugging. - - Args: - session_id: Session identifier to filter spans by - - Returns: - list[ReadableSpan]: List of spans associated with the session, - empty list if session not found or no spans available - """ - trace_ids = self.session_trace_dict.get(session_id, None) - if trace_ids is None or not trace_ids: - return [] - return [x for x in self._spans if x.context.trace_id in trace_ids] + def _session_traces(self) -> dict[str, dict[int, None]]: + """Build an index only from retained anchors while holding the cache lock""" + sessions: dict[str, dict[int, None]] = {} + for span in self._spans: + if span.name == "call_llm" and span.context: + session_id = (span.attributes or {}).get("gen_ai.session.id") + if isinstance(session_id, str) and session_id: + sessions.setdefault(session_id, {})[span.context.trace_id] = None + return sessions + + @property + def session_trace_dict(self) -> dict[str, list[int]]: + """Return a snapshot with no independently retained session metadata""" + with self._lock: + return {key: list(traces) for key, traces in self._session_traces().items()} + + def get_finished_spans(self, session_id: str) -> list[ReadableSpan]: + """Return retained session spans; evicted history is unavailable locally""" + with self._lock: + trace_ids = self._session_traces().get(session_id, {}) + return [ + span + for span in self._spans + if span.context and span.context.trace_id in trace_ids + ] def clear(self): - """Clear all stored spans and session mappings. + """Reset local history, accounting, session mappings, and the trace ID""" + with self._lock: + self._clear() - Removes all collected span data from memory, useful for cleanup - between test runs or to free memory in long-running processes. - """ + def _clear(self) -> None: self._spans.clear() + self._span_sizes.clear() + self._retained_bytes = 0 + self.trace_id = "" + + @override + def shutdown(self) -> None: + """Release history and stop retaining spans after provider shutdown""" + with self._lock: + self._shutdown = True + self._clear() class _InMemorySpanProcessor(export.SimpleSpanProcessor): @@ -239,7 +279,8 @@ class InMemoryExporter(BaseExporter): Note: - Cannot be added to exporter lists (validation prevents this) - Automatically managed by OpentelemetryTracer - - Memory usage grows with span count - consider periodic clearing + - Recent history is bounded by configured span and estimated byte limits + - Disabling retention preserves annotation and cloud export behavior - Session tracking requires properly configured session IDs """ @@ -253,5 +294,28 @@ def __init__(self, name: str = "inmemory_exporter") -> None: self.name = name - self._exporter = _InMemoryExporter() + from veadk.config import settings + + config = settings.opentelemetry_config + self._exporter = _InMemoryExporter( + max_spans=config.inmemory_max_spans if config.inmemory_enabled else 0, + max_bytes=config.inmemory_max_bytes, + ) self.processor = _InMemorySpanProcessor(self._exporter) + + @classmethod + def for_provider(cls, provider: TracerProvider) -> "InMemoryExporter": + """Reuse one annotation processor and bounded cache for each provider""" + with _registration_lock: + existing = getattr(provider, "_veadk_inmemory_exporter", None) + if existing is not None: + return existing + exporter = cls() + # Preserve first-processor annotation order without racing registration + processors = provider._active_span_processor + with processors._lock: + processors._span_processors = ( + exporter.processor, + ) + processors._span_processors + provider._veadk_inmemory_exporter = exporter + return exporter