fix(devserver): resolve scorer parents from the request state

py/src/braintrust/devserver/test_server.py

@@ -5,7 +5,7 @@
 from typing import Any
 
 import pytest
-from braintrust.logger import BraintrustState
+from braintrust.logger import BraintrustState, Experiment
 from braintrust.test_helpers import (
     has_devserver_installed,
     init_test_exp,
@@ -46,7 +46,10 @@ def do_POST(self) -> None:
         thread.join()
 
 
-def test_ui_scorer_forwards_org_name_and_active_parent_to_function_invoke(function_server, with_memory_logger):
+@pytest.mark.parametrize("global_span_active", [False, True])
+def test_ui_scorer_forwards_org_name_and_active_parent_to_function_invoke(
+    function_server, with_memory_logger, global_span_active
+):
     if not has_devserver_installed():
         pytest.skip("Devserver dependencies not installed (requires .[cli])")
 
@@ -56,13 +59,18 @@ def test_ui_scorer_forwards_org_name_and_active_parent_to_function_invoke(functi
     state = BraintrustState()
     state.org_name = "test-org"
     state.proxy_url = base_url
+    state._override_bg_logger.logger = with_memory_logger
+
+    global_experiment = init_test_exp("test-project")
+    request_experiment = Experiment(global_experiment._lazy_metadata, state=state)
 
     scorer = make_scorer(state, "ui-scorer", {"function_id": "function-id"}, "project-id")
 
-    with init_test_exp("test-project").start_span(name="eval") as row_span:
-        with row_span.start_span(name="ui-scorer", type="score") as scorer_span:
-            result = scorer("input", "output", "expected", {})
-            expected_parent = scorer_span.export()
+    with global_experiment.start_span(name="unrelated", set_current=global_span_active):
+        with request_experiment.start_span(name="eval") as row_span:
+            with row_span.start_span(name="ui-scorer", type="score") as scorer_span:
+                result = scorer("input", "output", "expected", {})
+                expected_parent = scorer_span.export()
 
     assert result == {"score": 1}
     [(org_name, request)] = requests

py/src/braintrust/logger.py

@@ -2797,13 +2797,17 @@ def current_logger() -> "Logger | None":
     return _state._cv_logger.get() or _state._local_logger
 
 
-def current_span() -> Span:
+def current_span(state: BraintrustState | None = None) -> Span:
     """Return the currently-active span for logging (set by running a span under a context manager). If there is no active span, returns a no-op span object, which supports the same interface as spans but does no logging.
 
+    Uses the supplied state, or the global state if none is provided.
     See `Span` for full details.
     """
 
-    span_info = _state.context_manager.get_current_span_info()
+    if state is None:
+        state = _state
+
+    span_info = state.context_manager.get_current_span_info()
     if span_info and hasattr(span_info.span_object, "span_id"):
         # This is a BT span
         return span_info.span_object
@@ -2849,7 +2853,7 @@ def _get_span_parent_object_and_propagated_state(
     if state is None:
         state = _state
 
-    span = current_span()
+    span = current_span(state=state)
     if span != NOOP_SPAN:
         return span, None
 

py/src/braintrust/test_logger.py

@@ -1964,6 +1964,31 @@ async def child_task():
     assert braintrust.current_logger() is parent_logger
 
 
+@pytest.mark.parametrize("global_span_active", [False, True])
+def test_current_span_and_parent_use_explicit_state(with_memory_logger, global_span_active):
+    global_experiment = init_test_exp("global-experiment")
+    state = BraintrustState()
+    state._override_bg_logger.logger = with_memory_logger
+    request_experiment = logger.Experiment(global_experiment._lazy_metadata, state=state)
+
+    with global_experiment.start_span(name="global", set_current=global_span_active) as global_span:
+        expected_global_span = global_span if global_span_active else logger.NOOP_SPAN
+        assert logger.current_span(state=state) is logger.NOOP_SPAN
+        with request_experiment.start_span(name="request") as request_span:
+            assert logger.current_span() is expected_global_span
+            assert logger.current_span(state=state) is request_span
+            assert logger.get_span_parent_object(state=state) is request_span
+            with logger.start_span(name="child", state=state) as child_span:
+                assert logger.current_span(state=state) is child_span
+            assert logger.current_span(state=state) is request_span
+        assert logger.current_span(state=state) is logger.NOOP_SPAN
+        assert logger.get_span_parent_object(state=state) is global_experiment
+
+    rows = {row["span_attributes"]["name"]: row for row in with_memory_logger.pop()}
+    assert rows["child"]["span_parents"] == [rows["request"]["span_id"]]
+    assert rows["child"]["root_span_id"] == rows["request"]["root_span_id"]
+
+
 def test_span_set_current(with_memory_logger):
     """Test that span.set_current() makes the span accessible via current_span()."""
     init_test_logger(__name__)