From ffa01f952a41473b5a701fb48dcb4deb9b9c6747 Mon Sep 17 00:00:00 2001
From: Hurshal Patel <hurshal@braintrustdata.com>
Date: Wed, 1 Jul 2026 16:43:48 -0700
Subject: [PATCH] expose realtime flag for span fetcher

---
 py/src/braintrust/logger.py     |  4 +-
 py/src/braintrust/test_trace.py | 70 ++++++++++++++++++++++++++++++++-
 py/src/braintrust/trace.py      |  4 ++
 3 files changed, 75 insertions(+), 3 deletions(-)

diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py
index 34aa9b74..9da8d9c2 100644
--- a/py/src/braintrust/logger.py
+++ b/py/src/braintrust/logger.py
@@ -3200,6 +3200,7 @@ def __init__(
         pinned_version: None | int | str = None,
         mutate_record: Callable[[TMapping], TMapping] | None = None,
         _internal_btql: dict[str, Any] | None = None,
+        _internal_brainstore_realtime: bool = True,
     ):
         self.object_type = object_type
 
@@ -3215,6 +3216,7 @@ def __init__(
 
         self._fetched_data: list[TMapping] | None = None
         self._internal_btql = _internal_btql
+        self._internal_brainstore_realtime = _internal_brainstore_realtime
 
     def fetch(self, batch_size: int | None = None) -> Iterator[TMapping]:
         """
@@ -3282,7 +3284,7 @@ def _refetch(self, batch_size: int | None = None) -> list[TMapping]:
                             **(self._internal_btql or {}),
                         },
                         "use_columnstore": False,
-                        "brainstore_realtime": True,
+                        "brainstore_realtime": self._internal_brainstore_realtime,
                         "query_source": f"py_sdk_object_fetcher_{self.object_type}",
                         **({"version": self._pinned_version} if self._pinned_version is not None else {}),
                     },
diff --git a/py/src/braintrust/test_trace.py b/py/src/braintrust/test_trace.py
index ebfea43c..346b47a1 100644
--- a/py/src/braintrust/test_trace.py
+++ b/py/src/braintrust/test_trace.py
@@ -1,7 +1,7 @@
 """Tests for Trace functionality."""
 
 import pytest
-from braintrust.trace import CachedSpanFetcher, LocalTrace, SpanData
+from braintrust.trace import CachedSpanFetcher, LocalTrace, SpanData, SpanFetcher
 
 
 # Helper to create mock spans
@@ -328,6 +328,48 @@ async def fetch_fn(span_type):
         assert call_args[0] is None or call_args[0] == []
         assert len(result) == 1
 
+    @pytest.mark.parametrize(
+        ("brainstore_realtime", "expected"),
+        [
+            (None, True),
+            (False, False),
+        ],
+    )
+    def test_span_fetcher_threads_realtime_setting(self, brainstore_realtime, expected):
+        calls = []
+        state = _DummyState(calls)
+        kwargs = dict(
+            object_type="project_logs",
+            object_id="project-1",
+            root_span_id="root-1",
+            state=state,
+        )
+        if brainstore_realtime is not None:
+            kwargs["brainstore_realtime"] = brainstore_realtime
+        fetcher = SpanFetcher(**kwargs)
+
+        assert list(fetcher.fetch()) == []
+        assert calls[0]["json"]["brainstore_realtime"] is expected
+
+    @pytest.mark.asyncio
+    async def test_cached_span_fetcher_threads_realtime_setting(self):
+        calls = []
+        state = _DummyState(calls)
+
+        async def get_state():
+            return state
+
+        fetcher = CachedSpanFetcher(
+            object_type="project_logs",
+            object_id="project-1",
+            root_span_id="root-1",
+            get_state=get_state,
+            brainstore_realtime=False,
+        )
+
+        assert await fetcher.get_spans() == []
+        assert calls[0]["json"]["brainstore_realtime"] is False
+
 
 class _DummySpanCache:
     def get_by_root_span_id(self, root_span_id: str):
@@ -335,12 +377,36 @@ def get_by_root_span_id(self, root_span_id: str):
 
 
 class _DummyState:
-    def __init__(self):
+    def __init__(self, api_calls=None):
         self.span_cache = _DummySpanCache()
+        self.api_calls = api_calls
 
     def login(self):
         return None
 
+    def api_conn(self):
+        return _DummyApiConn(self.api_calls)
+
+
+class _DummyResponse:
+    text = ""
+
+    def raise_for_status(self):
+        return None
+
+    def json(self):
+        return {"data": []}
+
+
+class _DummyApiConn:
+    def __init__(self, calls):
+        self.calls = calls
+
+    def post(self, path, *args, **kwargs):
+        if self.calls is not None:
+            self.calls.append({"path": path, "args": args, **kwargs})
+        return _DummyResponse()
+
 
 class TestLocalTraceGetThread:
     @pytest.mark.asyncio
diff --git a/py/src/braintrust/trace.py b/py/src/braintrust/trace.py
index c254fbda..ee62d671 100644
--- a/py/src/braintrust/trace.py
+++ b/py/src/braintrust/trace.py
@@ -75,6 +75,7 @@ def __init__(
         state: BraintrustState,
         span_type_filter: list[str] | None = None,
         include_scorers: bool = False,
+        brainstore_realtime: bool = True,
     ):
         # Build the filter expression for root_span_id and optionally span_attributes.type
         filter_expr = self._build_filter(root_span_id, span_type_filter, include_scorers)
@@ -82,6 +83,7 @@ def __init__(
         super().__init__(
             object_type=object_type,
             _internal_btql={"filter": filter_expr},
+            _internal_brainstore_realtime=brainstore_realtime,
         )
         self._object_id = object_id
         self._state = state
@@ -171,6 +173,7 @@ def __init__(
         root_span_id: str | None = None,
         get_state: Callable[[], Awaitable[BraintrustState]] | None = None,
         fetch_fn: SpanFetchFn | None = None,
+        brainstore_realtime: bool = True,
     ):
         self._span_cache: dict[str, list[SpanData]] = {}
         self._all_fetched = False
@@ -204,6 +207,7 @@ async def _fetch_fn(
                     state=state,
                     span_type_filter=span_type,
                     include_scorers=include_scorers,
+                    brainstore_realtime=brainstore_realtime,
                 )
                 rows = list(fetcher.fetch())
                 return [