bitfab-py 0.76.0__tar.gz → 0.76.2__tar.gz

{bitfab_py-0.76.0 → bitfab_py-0.76.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bitfab-py
-Version: 0.76.0
+Version: 0.76.2
 Summary: Bitfab client for provider-based API calls with local BAML execution
 Author: Harvest Team
 Requires-Python: >=3.10

{bitfab_py-0.76.0 → bitfab_py-0.76.2}/bitfab/client.py

@@ -45,6 +45,7 @@ from bitfab.constants import (
 )
 from bitfab.datasets import DatasetsClient
 from bitfab.db_snapshot import DbSnapshotRef, build_snapshot_ref
+from bitfab.error_chain import describe_error
 from bitfab.experiment_id import resolve_experiment_id
 from bitfab.experiments import ExperimentsClient
 from bitfab.graders import GradersClient
@@ -477,6 +478,7 @@ class SpanContext(TypedDict, total=False):
     enclosingTrace: TraceLink
     variantChain: VariantChain
     spanName: str
+    traceFunctionKey: str
     rootTraceFunctionKey: str
     instrumentation: SpanInstrumentation
     declaredNode: bool
@@ -547,6 +549,10 @@ _span_stack: ContextVar[list[SpanContext]] = ContextVar(
     "bitfab_span_stack", default=None
 )
 
+_enclosing_trace_function_keys: ContextVar[tuple[str, ...]] = ContextVar(
+    "bitfab_enclosing_trace_function_keys", default=()
+)
+
 # Global dict to track active trace states (trace_id -> TraceState)
 _active_trace_states: dict[str, TraceState] = {}
 
@@ -680,11 +686,28 @@ def _trace_root_span_context() -> SpanContext | None:
     )
 
 
+def _running_trace_function_keys(stack: list[SpanContext]) -> list[str]:
+    keys: list[str] = []
+    candidates = list(_enclosing_trace_function_keys.get())
+    candidates.extend(context.get("traceFunctionKey", "") for context in stack)
+    for key in candidates:
+        if key and key not in keys:
+            keys.append(key)
+    return keys
+
+
 class CurrentSpan:
     """Handle to the current active span, allowing context to be added."""
 
-    def __init__(self, context: SpanContext) -> None:
+    def __init__(
+        self, context: SpanContext, trace_function_keys: list[str] | None = None
+    ) -> None:
         self._context = context
+        self._trace_function_keys = (
+            trace_function_keys
+            if trace_function_keys is not None
+            else _running_trace_function_keys([context])
+        )
 
     @property
     def id(self) -> str:
@@ -696,6 +719,15 @@ class CurrentSpan:
         """The trace ID for the current span."""
         return self._context["traceId"]
 
+    @property
+    def trace_function_keys(self) -> list[str]:
+        """The keys of every trace function running at this point, outermost
+        first, each listed once. A trace function called inside another one
+        comes after it, so a cost recorded here can be counted toward each of
+        them. Empty outside a trace.
+        """
+        return list(self._trace_function_keys)
+
     def add_context(self, context: dict[str, Any]) -> None:
         """Add a context entry to this span.
 
@@ -744,8 +776,26 @@ def _on_event_loop_thread() -> bool:
 class CurrentTrace:
     """Handle to the current active trace, allowing trace-level context to be set."""
 
-    def __init__(self, trace_id: str) -> None:
-        self._trace_id = trace_id
+    def __init__(self, context: SpanContext) -> None:
+        self._trace_id = context["traceId"]
+        self._root_trace_function_key = context.get("rootTraceFunctionKey")
+
+    @property
+    def trace_function_key(self) -> str:
+        """The key of the trace function the current trace is stored under,
+        which is the one that started it.
+
+        A ``@span`` or ``get_function`` call inside another trace function
+        joins that trace, so it reports the outer key. A ``trace()`` root
+        inside another root starts its own trace, so it reports its own key.
+        Empty string outside a trace.
+        """
+        return (
+            _resolve_root_trace_function_key(
+                self._root_trace_function_key, self._trace_id, None
+            )
+            or ""
+        )
 
     def _get_or_create_trace_state(self) -> TraceState:
         """Get existing trace state or create a new one."""
@@ -961,6 +1011,11 @@ class _NoOpCurrentSpan:
         """Return empty string when outside a span context."""
         return ""
 
+    @property
+    def trace_function_keys(self) -> list[str]:
+        """Return an empty list when outside a span context."""
+        return []
+
     def add_context(self, context: dict[str, Any]) -> None:
         pass
 
@@ -971,6 +1026,11 @@ class _NoOpCurrentSpan:
 class _NoOpCurrentTrace:
     """No-op trace handle returned when outside a span context."""
 
+    @property
+    def trace_function_key(self) -> str:
+        """Return empty string when outside a span context."""
+        return ""
+
     def set_session_id(self, session_id: str) -> None:
         pass
 
@@ -1001,7 +1061,7 @@ def get_current_span() -> CurrentSpan | _NoOpCurrentSpan:
     """
     stack = _get_span_stack()
     if stack:
-        return CurrentSpan(stack[-1])
+        return CurrentSpan(stack[-1], _running_trace_function_keys(stack))
     return _NO_OP_SPAN
 
 
@@ -1045,10 +1105,22 @@ def get_current_trace() -> CurrentTrace | _NoOpCurrentTrace:
     """
     stack = _get_span_stack()
     if stack:
-        return CurrentTrace(stack[-1]["traceId"])
+        return CurrentTrace(stack[-1])
     return _NO_OP_TRACE
 
 
+@contextlib.contextmanager
+def _trace_root_scope(scope: tuple[list[SpanContext], tuple[str, ...]]) -> Any:
+    stack, enclosing_keys = scope
+    stack_token = _span_stack.set(stack)
+    keys_token = _enclosing_trace_function_keys.set(enclosing_keys)
+    try:
+        yield
+    finally:
+        _enclosing_trace_function_keys.reset(keys_token)
+        _span_stack.reset(stack_token)
+
+
 def _run_with_span_stack(stack: list[SpanContext], fn: Callable[[], T]) -> T:
     """Run a function with a new span stack context."""
     token = _span_stack.set(stack)
@@ -2660,6 +2732,7 @@ class Bitfab:
                     "spanId": span_id,
                     "threadId": executing_thread.ident or 0,
                     "spanName": span_name,
+                    "traceFunctionKey": trace_function_key,
                     "rootTraceFunctionKey": (
                         root_state_at_open.get("traceFunctionKey", trace_function_key)
                         if root_state_at_open
@@ -3258,7 +3331,7 @@ class Bitfab:
                             yielded.append(value)
                             yield value
                     except Exception as e:
-                        error_msg = str(e)
+                        error_msg = describe_error(e)
                         raise
                     finally:
                         token = _span_stack.set(new_stack)
@@ -3271,7 +3344,7 @@ class Bitfab:
                                 ):
                                     iterator.close()
                             except Exception as e:
-                                error_msg = str(e)
+                                error_msg = describe_error(e)
                                 raise
                             finally:
                                 if finalize is not None and error_msg is None:
@@ -3369,7 +3442,7 @@ class Bitfab:
                         error_msg = str(e) or e.__class__.__name__
                         raise
                     except Exception as e:
-                        error_msg = str(e)
+                        error_msg = describe_error(e)
                         raise
                     finally:
                         token = _span_stack.set(new_stack)
@@ -3385,7 +3458,7 @@ class Bitfab:
                                 error_msg = str(e) or e.__class__.__name__
                                 raise
                             except Exception as e:
-                                error_msg = str(e)
+                                error_msg = describe_error(e)
                                 raise
                             finally:
                                 if not sent:
@@ -3563,7 +3636,7 @@ class Bitfab:
                             span_ctx,
                             **base_params,
                             result=None,
-                            error=str(e),
+                            error=describe_error(e),
                             ended_at=_get_ended_at(),
                         )
                         raise
@@ -3649,7 +3722,7 @@ class Bitfab:
                             span_ctx,
                             **base_params,
                             result=None,
-                            error=str(e),
+                            error=describe_error(e),
                             ended_at=_get_ended_at(),
                         )
                         raise
@@ -4431,16 +4504,16 @@ class Bitfab:
                 surface="opt-out",
             )(inner)
 
-            def root_span_stack() -> list[SpanContext] | None:
+            def root_span_stack() -> tuple[list[SpanContext], tuple[str, ...]] | None:
                 stack = _get_span_stack()
                 enclosing = stack[-1].get("surface") if stack else None
                 if enclosing == "opt-in":
                     raise _mixed_tracing_error("trace", "span")
                 if enclosing != "opt-out":
-                    return stack
+                    return stack, _enclosing_trace_function_keys.get()
                 if _in_replay_item() or _in_seed_scope():
                     return None
-                return []
+                return [], tuple(_running_trace_function_keys(stack))
 
             if inspect.isasyncgenfunction(fn):
 
@@ -4458,21 +4531,15 @@ class Bitfab:
                     iterator = spanned(*args, **kwargs)
                     try:
                         while True:
-                            token = _span_stack.set(root_stack)
-                            try:
+                            with _trace_root_scope(root_stack):
                                 try:
                                     item = await iterator.__anext__()
                                 except StopAsyncIteration:
                                     return
-                            finally:
-                                _span_stack.reset(token)
                             yield item
                     finally:
-                        token = _span_stack.set(root_stack)
-                        try:
+                        with _trace_root_scope(root_stack):
                             await iterator.aclose()
-                        finally:
-                            _span_stack.reset(token)
 
             elif inspect.iscoroutinefunction(fn):
 
@@ -4481,11 +4548,8 @@ class Bitfab:
                     root_stack = root_span_stack()
                     if root_stack is None:
                         return await body(*args, **kwargs)
-                    token = _span_stack.set(root_stack)
-                    try:
+                    with _trace_root_scope(root_stack):
                         return await spanned(*args, **kwargs)
-                    finally:
-                        _span_stack.reset(token)
 
             else:
 
@@ -4494,9 +4558,8 @@ class Bitfab:
                     root_stack = root_span_stack()
                     if root_stack is None:
                         return body(*args, **kwargs)
-                    return _run_with_span_stack(
-                        root_stack, lambda: spanned(*args, **kwargs)
-                    )
+                    with _trace_root_scope(root_stack):
+                        return spanned(*args, **kwargs)
 
             _trace_root_wrappers.add(detached)
             return detached  # type: ignore[return-value]

bitfab_py-0.76.2/bitfab/error_chain.py

@@ -0,0 +1,101 @@
+"""Read the chain of causes below an exception.
+
+A database driver error usually reaches a traced function wrapped one or two
+levels down (``raise QueryError(...) from timeout``). Recording only the top
+message makes a connection timeout read as a broken query, so spans and replay
+results carry the deepest cause too.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+MAX_ERROR_CAUSE_DEPTH = 10
+"""How many causes below an exception are followed."""
+
+_CODE_ATTRIBUTES = ("code", "sqlstate", "pgcode", "errno")
+
+
+def _cause_of(error: BaseException) -> BaseException | None:
+    # The same rule Python's traceback uses: an explicit ``from`` cause wins,
+    # and an implicit context counts unless ``from None`` suppressed it. Bitfab's
+    # own replay errors also keep their cause on a ``cause`` attribute.
+    if error.__cause__ is not None:
+        return error.__cause__
+    attribute = getattr(error, "cause", None)
+    if isinstance(attribute, BaseException):
+        return attribute
+    if error.__suppress_context__:
+        return None
+    return error.__context__
+
+
+def _error_code(error: BaseException) -> str | int | None:
+    for attribute in _CODE_ATTRIBUTES:
+        value = getattr(error, attribute, None)
+        if isinstance(value, (str, int)) and not isinstance(value, bool):
+            return value
+    return None
+
+
+def error_causes(error: BaseException) -> list[BaseException]:
+    """The causes below ``error``, nearest first, stopping at the depth limit
+    or at a cause already seen."""
+    causes: list[BaseException] = []
+    seen = {id(error)}
+    current = error
+    while len(causes) < MAX_ERROR_CAUSE_DEPTH:
+        cause = _cause_of(current)
+        if cause is None or id(cause) in seen:
+            break
+        causes.append(cause)
+        seen.add(id(cause))
+        current = cause
+    return causes
+
+
+def describe_error(error: BaseException) -> str:
+    """The exception's message, followed by the deepest cause's message when
+    it wraps another exception, e.g. ``Failed query: ... (caused by:
+    connection timeout)``. The cause is left off when the message already
+    contains it."""
+    message = str(error)
+    causes = error_causes(error)
+    if not causes:
+        return message
+    deepest = str(causes[-1]).strip()
+    if not deepest or deepest in message:
+        return message
+    return f"{message} (caused by: {deepest})"
+
+
+def serialize_error(
+    error: BaseException,
+    extra: Callable[[BaseException], dict[str, Any]] | None = None,
+) -> dict[str, Any]:
+    """An exception as a JSON record: type, message, the driver's code when it
+    has one, and its causes nested under ``cause`` down to the depth limit.
+    ``extra`` adds fields a specific exception type carries."""
+    seen: set[int] = set()
+
+    def serialize(value: BaseException, depth: int) -> dict[str, Any]:
+        seen.add(id(value))
+        serialized: dict[str, Any] = {
+            "type": type(value).__name__,
+            "message": str(value),
+        }
+        code = _error_code(value)
+        if code is not None:
+            serialized["code"] = code
+        if extra is not None:
+            serialized.update(extra(value))
+        cause = _cause_of(value)
+        if (
+            cause is not None
+            and id(cause) not in seen
+            and depth < MAX_ERROR_CAUSE_DEPTH
+        ):
+            serialized["cause"] = serialize(cause, depth + 1)
+        return serialized
+
+    return serialize(error, 1)

{bitfab_py-0.76.0 → bitfab_py-0.76.2}/bitfab/replay.py

@@ -24,6 +24,7 @@ from bitfab.constants import (
     _replay_context,
     _seed_context,
 )
+from bitfab.error_chain import describe_error, serialize_error
 from bitfab.git_command import run_git
 from bitfab.integration_fixture import extract_integration_fixtures
 from bitfab.mock_matching import MATCHER
@@ -580,19 +581,16 @@ def report_replay_progress(
 
 def _replay_json_default(value: Any) -> Any:
     if isinstance(value, BaseException):
-        serialized = {
-            "type": type(value).__name__,
-            "message": str(value),
-        }
-        if isinstance(value, DbBranchReplayError):
-            serialized["code"] = value.code
-            serialized["original_trace_id"] = value.original_trace_id
-            if value.cause is not None:
-                serialized["cause"] = value.cause
-        return serialized
+        return serialize_error(value, _replay_error_fields)
     return str(value)
 
 
+def _replay_error_fields(error: BaseException) -> dict[str, Any]:
+    if isinstance(error, DbBranchReplayError):
+        return {"original_trace_id": error.original_trace_id}
+    return {}
+
+
 def serialize_replay_result(result: ReplayResult) -> str:
     """Serialize a replay result while retaining structured exception fields."""
     return json.dumps(result, indent=2, default=_replay_json_default)
@@ -653,7 +651,7 @@ def _replay_item_error_message(error: BaseException) -> str:
             "replaying it against the live database would produce a result that "
             "looks valid but did not use the historical data you asked for."
         )
-    return str(error)
+    return describe_error(error)
 
 
 class AdaptContext(TypedDict):
@@ -3041,7 +3039,7 @@ async def _execute_item_async(
         input=reported_input,
         result=fn_result,
         original_output=item.get("output"),
-        error=str(item_error) if item_error is not None else None,
+        error=describe_error(item_error) if item_error is not None else None,
         trace_error=fn_error,
         replay_error=replay_error,
         duration_ms=replay_duration_ms,

{bitfab_py-0.76.0 → bitfab_py-0.76.2}/bitfab/subtree.py

@@ -23,6 +23,7 @@ from dataclasses import dataclass, field
 from typing import Any, Callable
 
 from bitfab.constants import _replay_context
+from bitfab.error_chain import describe_error
 from bitfab.trace_metadata import record_caller_trace_metadata
 from bitfab.warn_once import warn_once
 
@@ -795,7 +796,9 @@ def _on_unwind(code: Any, instruction_offset: int, exception: BaseException) ->
     # the first time an exception passed through library code. Filter instead.
     if is_library_code(code.co_filename):
         return None
-    return _close(code, None, f"{type(exception).__name__}: {exception}")
+    return _close(
+        code, None, f"{type(exception).__name__}: {describe_error(exception)}"
+    )
 
 
 def _is_free_threaded() -> bool:
@@ -1033,7 +1036,11 @@ class BorrowedSpan:
             updated[session] = _without(stack, index)
         _open.set(updated)
         if exc_type is not None:
-            self.emit(None, f"{exc_type.__name__}: {exc}")
+            self.emit(
+                None,
+                f"{exc_type.__name__}: "
+                f"{describe_error(exc) if isinstance(exc, BaseException) else exc}",
+            )
 
 
 def borrowed_span(

{bitfab_py-0.76.0 → bitfab_py-0.76.2}/bitfab/thread_propagation.py

@@ -17,8 +17,8 @@ the SDK's own span-transport threads are excluded so a long-lived dispatch
 loop can never pin the context of whichever trace happened to spawn it.
 Callables that already carry a context (``asyncio.to_thread`` submits
 ``Context.run``) pass through untouched. Only Bitfab's own variables are
-carried (the span stack, the replay context, the seed context, and the submit
-origin: the identity of the thread that called submit()/start(), recorded onto
+carried (the span stack, the keys of the trace functions around a nested
+trace root, the replay context, the seed context, and the submit origin: the identity of the thread that called submit()/start(), recorded onto
 worker spans as ``runtime.submit_thread_id``), so no other library's
 contextvars change behavior.
 
@@ -54,7 +54,7 @@ _installed = False
 _original_submit: Callable[..., Any] | None = None
 _original_thread_start: Callable[..., Any] | None = None
 
-_CapturedState = tuple[Any, Any, Any, dict[str, Any]]
+_CapturedState = tuple[Any, tuple[str, ...], Any, Any, dict[str, Any]]
 
 
 def _span_stack_var() -> contextvars.ContextVar[Any]:
@@ -63,9 +63,16 @@ def _span_stack_var() -> contextvars.ContextVar[Any]:
     return _span_stack
 
 
+def _enclosing_keys_var() -> contextvars.ContextVar[tuple[str, ...]]:
+    from bitfab.client import _enclosing_trace_function_keys
+
+    return _enclosing_trace_function_keys
+
+
 def _capture() -> _CapturedState | None:
     try:
         stack = _span_stack_var().get()
+        enclosing_keys = _enclosing_keys_var().get()
         replay_ctx = _replay_context.get()
         seed_ctx = _seed_context.get()
     except Exception:
@@ -74,14 +81,16 @@ def _capture() -> _CapturedState | None:
         return None
     submitter = threading.current_thread()
     origin = {"thread_id": submitter.ident, "thread_name": submitter.name}
-    return (stack, replay_ctx, seed_ctx, origin)
+    return (stack, enclosing_keys, replay_ctx, seed_ctx, origin)
 
 
 @contextlib.contextmanager
 def _entered(state: _CapturedState) -> Any:
-    stack, replay_ctx, seed_ctx, origin = state
+    stack, enclosing_keys, replay_ctx, seed_ctx, origin = state
     span_var = _span_stack_var()
+    keys_var = _enclosing_keys_var()
     span_token = span_var.set(stack)
+    keys_token = keys_var.set(enclosing_keys)
     replay_token = _replay_context.set(replay_ctx)
     seed_token = _seed_context.set(seed_ctx)
     origin_token = _submit_origin.set(origin)
@@ -89,6 +98,7 @@ def _entered(state: _CapturedState) -> Any:
         yield
     finally:
         span_var.reset(span_token)
+        keys_var.reset(keys_token)
         _replay_context.reset(replay_token)
         _seed_context.reset(seed_token)
         _submit_origin.reset(origin_token)

{bitfab_py-0.76.0 → bitfab_py-0.76.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "bitfab-py"
-version = "0.76.0"
+version = "0.76.2"
 description = "Bitfab client for provider-based API calls with local BAML execution"
 authors = ["Harvest Team"]
 readme = "README.md"