diff --git a/py/noxfile.py b/py/noxfile.py index 5ecb41e7..8eb5f7d1 100644 --- a/py/noxfile.py +++ b/py/noxfile.py @@ -929,7 +929,7 @@ def test_otel(session, version): _install_matrix_dep(session, "opentelemetry-api", version) _install_matrix_dep(session, "opentelemetry-sdk", version) _install_matrix_dep(session, "opentelemetry-exporter-otlp-proto-http", version) - _run_tests(session, "braintrust/test_otel.py", version=version) + _run_tests(session, ["braintrust/test_otel.py", "braintrust/otel/"], version=version) @nox.session() diff --git a/py/src/braintrust/integrations/langchain/callbacks.py b/py/src/braintrust/integrations/langchain/callbacks.py index bf93bc1f..dbd413b8 100644 --- a/py/src/braintrust/integrations/langchain/callbacks.py +++ b/py/src/braintrust/integrations/langchain/callbacks.py @@ -246,15 +246,7 @@ def _end_span( dataset_record_id=dataset_record_id, ) - # Async callbacks may unset from a different context; span state is - # tracked in self.spans, so this ValueError is benign. - try: - span.unset_current() - except ValueError as e: - if "was created in a different Context" in str(e): - pass - else: - raise + span.unset_current() span.end() diff --git a/py/src/braintrust/otel/context.py b/py/src/braintrust/otel/context.py index ea9703bc..35299bca 100644 --- a/py/src/braintrust/otel/context.py +++ b/py/src/braintrust/otel/context.py @@ -1,6 +1,7 @@ """Unified context management using OTEL's built-in context.""" import logging +from contextvars import Token from typing import Any from braintrust.context import ParentSpanIds, SpanInfo @@ -90,8 +91,12 @@ def unset_current_span(self, context_token: Any = None) -> None: """Unset the current active span from OTEL context.""" if context_token: - # Detaching the token restores the previous context - context.detach(context_token) + # OTel's default ContextVarsRuntimeContext hands out contextvars tokens. + if isinstance(context_token, Token): + _reset_contextvars_token(context_token) + else: + # Non-contextvars OTel runtimes keep their own detach semantics. + context.detach(context_token) else: # No token means we need to explicitly clear the span # This shouldn't normally happen, but handle it gracefully @@ -108,6 +113,21 @@ def get_parent_span_ids(self) -> ParentSpanIds | None: ) +def _reset_contextvars_token(token: Token) -> None: + """Reset like ``context.detach()``, but treat a cross-Context reset as benign. + + Async frameworks (e.g. LangGraph) often end a span in a different Context than + it started in; ``context.detach()`` would log that at ERROR. + """ + try: + token.var.reset(token) + except ValueError: + # Only raised for a token from another Context; leave both contexts as-is. + log.debug("Skipping detach of span context token created in a different Context") + except Exception: # pylint: disable=broad-except + log.exception("Failed to detach context") + + def _is_otel_span(span: Any) -> bool: """Check if the span object is an OTEL span.""" return hasattr(span, "get_span_context") diff --git a/py/src/braintrust/otel/test_otel_bt_integration.py b/py/src/braintrust/otel/test_otel_bt_integration.py index b6d9ebff..f4e05406 100644 --- a/py/src/braintrust/otel/test_otel_bt_integration.py +++ b/py/src/braintrust/otel/test_otel_bt_integration.py @@ -5,6 +5,8 @@ when created in mixed contexts. """ +import contextvars +import logging import os import pytest @@ -387,5 +389,72 @@ def test_bt_span_without_explicit_parent_inherits_from_otel(otel_fixture): assert bt_child["span_parents"] == [otel_span_id] +def test_unset_current_span_restores_previous_context(otel_fixture): + from opentelemetry import context + + logger = init_test_logger(__name__) + with logger.start_span(name="parent") as parent: + child = logger.start_span(name="child") + child.set_current() + assert context.get_value("braintrust_span") is child + + child.unset_current() + child.end() + + assert context.get_value("braintrust_span") is parent + + +def test_unset_current_span_from_another_context_does_not_log_error(otel_fixture, caplog): + """Regression for #1062: LangGraph-style async callbacks start and end a span + in different contextvars Contexts; that must not log an OTel ERROR.""" + from opentelemetry import context + + logger = init_test_logger(__name__) + owner = contextvars.copy_context() + other = contextvars.copy_context() + + span = logger.start_span(name="cross_context") + other_span = logger.start_span(name="other_active") + owner.run(span.set_current) + other.run(other_span.set_current) + + with caplog.at_level(logging.DEBUG): + other.run(span.unset_current) + span.end() + other_span.end() + + otel_errors = [r for r in caplog.records if r.levelno >= logging.ERROR and r.name.startswith("opentelemetry")] + assert not otel_errors, [r.getMessage() for r in otel_errors] + # The destination context's own active span must be left untouched. + assert other.run(context.get_value, "braintrust_span") is other_span + + +def test_unset_current_span_reused_token_still_logs_without_raising(otel_fixture, caplog): + from braintrust.otel.context import ContextManager + from opentelemetry import context + + otel_context_manager = ContextManager() + token = context.attach(context.set_value("braintrust_span", "child")) + otel_context_manager.unset_current_span(token) + + with caplog.at_level(logging.ERROR, logger="braintrust.otel.context"): + otel_context_manager.unset_current_span(token) + + assert any(r.getMessage() == "Failed to detach context" for r in caplog.records) + + +def test_unset_current_span_non_contextvars_token_uses_otel_detach(otel_fixture, monkeypatch): + from braintrust.otel.context import ContextManager + from opentelemetry import context + + token = object() + calls = [] + monkeypatch.setattr(context, "detach", calls.append) + + ContextManager().unset_current_span(token) + + assert calls == [token] + + if __name__ == "__main__": pytest.main([__file__, "-v"])