Skip to content

fix(temporal): pass Exception directly to span.log(error=...) #611

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Abhijeet Prasad (AbhiPrasad) merged 1 commit into from
Jul 21, 2026
Merged

fix(temporal): pass Exception directly to span.log(error=...) #611

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
70 changes: 12 additions & 58 deletions py/src/braintrust/integrations/temporal/plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,22 +99,16 @@
from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner


# Braintrust dynamically chooses its context implementation at runtime based on
# BRAINTRUST_OTEL_COMPAT environment variable. When first accessed, it reads
# os.environ which is restricted in the sandbox. Therefore if the first use
# is inside the sandbox, it will fail. So we eagerly reference it here to
# force initialization at import time (before sandbox evaluation).
# Braintrust picks its context implementation on first access by reading
# BRAINTRUST_OTEL_COMPAT from os.environ, which the sandbox blocks. Force
# init at import time so the first access doesn't happen inside the sandbox.
try:
braintrust.current_span()
except Exception:
# It's okay if this fails (e.g., no logger initialized yet)
pass

# Store module-level reference to braintrust.current_span to avoid re-importing
# inside extern functions (which can trigger sandbox restrictions)
_current_span = braintrust.current_span

# Header key for passing span context between client, workflows, and activities
_HEADER_KEY = "_braintrust-span"
_SPAN_ID_NAMESPACE = "braintrust.temporal.workflow"
_SPAN_CONTEXT_PATCH_ID = "braintrust-temporal-workflow-span-context-v1"
Expand All @@ -138,42 +132,28 @@ class BraintrustInterceptor(temporalio.client.Interceptor, temporalio.worker.Int
"""

def __init__(self, logger: Any | None = None) -> None:
"""Initialize interceptor.

Args:
logger: Optional background logger for testing.
"""
self.payload_converter = temporalio.converter.PayloadConverter.default
self._bg_logger = logger
# Capture logger instance at init time for cross-thread use
if logger:
braintrust.logger._state._override_bg_logger.logger = logger
self._logger = braintrust.current_logger()

def _get_logger(self) -> Any | None:
"""Get logger for creating spans.

Sets thread-local override if background logger provided (for testing),
then returns captured logger instance.
"""
if self._bg_logger:
braintrust.logger._state._override_bg_logger.logger = self._bg_logger
return self._logger

def intercept_client(self, next: temporalio.client.OutboundInterceptor) -> temporalio.client.OutboundInterceptor:
"""Intercept client calls to propagate span context to workflows."""
return _BraintrustClientOutboundInterceptor(next, self)

def intercept_activity(
self, next: temporalio.worker.ActivityInboundInterceptor
) -> temporalio.worker.ActivityInboundInterceptor:
"""Intercept activity executions to create activity spans."""
return _BraintrustActivityInboundInterceptor(next, self)

def workflow_interceptor_class(
self, input: temporalio.worker.WorkflowInterceptorClassInput
) -> type["BraintrustWorkflowInboundInterceptor"] | None:
"""Return workflow interceptor class to propagate context to activities."""
input.unsafe_extern_functions["__braintrust_get_logger"] = self._get_logger
return BraintrustWorkflowInboundInterceptor

Expand All @@ -182,7 +162,6 @@ def _span_context_to_headers(
span_context: Any,
headers: Mapping[str, temporalio.api.common.v1.Payload],
) -> Mapping[str, temporalio.api.common.v1.Payload]:
"""Add span context to headers."""
if span_context:
payloads = self.payload_converter.to_payloads([span_context])
if payloads:
Expand All @@ -193,7 +172,6 @@ def _span_context_to_headers(
return headers

def _span_context_from_headers(self, headers: Mapping[str, temporalio.api.common.v1.Payload]) -> Any | None:
"""Extract span context from headers."""
if _HEADER_KEY not in headers:
return None
header_payload = headers.get(_HEADER_KEY)
Expand All @@ -215,12 +193,10 @@ def __init__(self, next: temporalio.client.OutboundInterceptor, root: Braintrust
async def start_workflow(
self, input: temporalio.client.StartWorkflowInput
) -> temporalio.client.WorkflowHandle[Any, Any]:
# Get current span context and add it to workflow headers
current_span = _current_span()
if current_span:
span_context = current_span.export()
input.headers = self.root._span_context_to_headers(span_context, input.headers)

return await super().start_workflow(input)


Expand All @@ -237,15 +213,12 @@ def __init__(

async def execute_activity(self, input: temporalio.worker.ExecuteActivityInput) -> Any:
info = temporalio.activity.info()

# Extract parent span context from headers
parent_span_context = self.root._span_context_from_headers(input.headers)

logger = self.root._get_logger()
if not logger:
return await super().execute_activity(input)

# Create Braintrust span for activity execution, linked to workflow span
span = logger.start_span(
internal={"instrumentation": "temporal-auto"},
name=f"temporal.activity.{info.activity_type}",
Expand All @@ -261,10 +234,9 @@ async def execute_activity(self, input: temporalio.worker.ExecuteActivityInput)
span.set_current()

try:
result = await super().execute_activity(input)
return result
return await super().execute_activity(input)
except Exception as e:
span.log(error=str(e))
span.log(error=e)
raise
finally:
span.unset_current()
Expand All @@ -289,7 +261,6 @@ def init(self, outbound: temporalio.worker.WorkflowOutboundInterceptor) -> None:
super().init(_BraintrustWorkflowOutboundInterceptor(outbound, self))

async def execute_workflow(self, input: temporalio.worker.ExecuteWorkflowInput) -> Any:
# Extract parent span context from workflow headers (set by client)
parent_span_context = None
if _HEADER_KEY in input.headers:
header_payload = input.headers.get(_HEADER_KEY)
Expand All @@ -298,14 +269,11 @@ async def execute_workflow(self, input: temporalio.worker.ExecuteWorkflowInput)
if payloads:
parent_span_context = payloads[0]

# Store parent span context for activities (will be overwritten if we create a workflow span)
self._parent_span_context = parent_span_context

# Create a span for the workflow execution using sandbox_unrestricted
# to bypass the sandbox restrictions on logger state access
# sandbox_unrestricted lets us touch logger state that the workflow sandbox would otherwise block.
span = None
with temporalio.workflow.unsafe.sandbox_unrestricted():
# Get logger via extern function (supports test logger parameter)
get_logger = temporalio.workflow.extern_functions()["__braintrust_get_logger"]
logger = get_logger()

Expand Down Expand Up @@ -351,12 +319,11 @@ async def execute_workflow(self, input: temporalio.worker.ExecuteWorkflowInput)
self._workflow_span = span

try:
result = await super().execute_workflow(input)
return result
return await super().execute_workflow(input)
except Exception as e:
if self._workflow_span:
with temporalio.workflow.unsafe.sandbox_unrestricted():
self._workflow_span.log(error=str(e))
self._workflow_span.log(error=e)
raise
finally:
if self._workflow_span:
Expand All @@ -379,13 +346,8 @@ def __init__(
def _add_span_context_to_headers(
self, headers: Mapping[str, temporalio.api.common.v1.Payload]
) -> Mapping[str, temporalio.api.common.v1.Payload]:
"""Add parent span context to headers if available.

Note: We always pass span context through headers, even during replay,
so activities can maintain proper parent-child relationships. The replay
safety is handled in the activity interceptor, which only creates spans
when the activity actually executes (not during replay).
"""
# Always propagate context — even during replay — so activities can find their parent.
# Replay safety is enforced in the activity interceptor, which only spans on real execution.
if self.root._parent_span_context:
payloads = self.root.payload_converter.to_payloads([self.root._parent_span_context])
if payloads:
Expand All @@ -410,7 +372,6 @@ def start_child_workflow(


def _modify_workflow_runner(existing: WorkflowRunner | None) -> WorkflowRunner | None:
"""Add braintrust to sandbox passthrough modules."""
if isinstance(existing, SandboxedWorkflowRunner):
new_restrictions = existing.restrictions.with_passthrough_modules("braintrust")
return dataclasses.replace(existing, restrictions=new_restrictions)
Expand Down Expand Up @@ -471,20 +432,13 @@ class BraintrustPlugin(SimplePlugin):
"""

def __init__(self, logger: Any | None = None) -> None:
"""Initialize the Braintrust plugin.

Args:
logger: Optional background logger for testing.
"""
interceptor = BraintrustInterceptor(logger=logger)
import inspect

params = inspect.signature(SimplePlugin.__init__).parameters

# temporalio plugin APIs differ by version. Newer releases may expose
# a merged `interceptors` kwarg, while older releases keep separate
# client/worker interceptor kwargs. Build kwargs dynamically so pylint
# does not validate unsupported keywords against the installed version.
# Newer temporalio exposes a merged `interceptors` kwarg; older releases keep separate
# client_/worker_interceptors. Build kwargs dynamically to stay compatible.
kwargs: dict[str, Any] = {
"name": "braintrust",
"workflow_runner": _modify_workflow_runner,
Expand Down