Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion py/noxfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
10 changes: 1 addition & 9 deletions py/src/braintrust/integrations/langchain/callbacks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
24 changes: 22 additions & 2 deletions py/src/braintrust/otel/context.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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")
69 changes: 69 additions & 0 deletions py/src/braintrust/otel/test_otel_bt_integration.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
when created in mixed contexts.
"""

import contextvars
import logging
import os

import pytest
Expand Down Expand Up @@ -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"])
Loading