From ecbf1e44437eb9eb183a9262cdbf548a772e9006 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 2 Nov 2026 22:53:45 -0500 Subject: [PATCH 1/3] fix(temporal): keep pydantic-ai subagent spans nested and replay-safe MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > 🤖 **This text was AI-generated.** With `BraintrustPlugin` and Pydantic AI's Temporal durability integration, a tool such as `delegate_task` runs in workflow code and starts a sub-agent child workflow. Two things went wrong: 1. The child workflow was parented to the workflow span, so it showed up as a sibling of the `delegate_task` span instead of under it. The workflow outbound interceptor always sent the workflow span as the parent of activities, local activities, and child workflows, ignoring the current in-workflow span. 2. Temporal re-runs workflow code from the start whenever a workflow replays, for example after a worker restart or when the workflow is evicted from the worker cache. Spans opened in workflow code, such as Pydantic AI's `agent_run`, `delegate_task`, and `chat` spans or users' own spans, were logged again on every replay with new random IDs. The plugin also skipped the workflow span during replay, so the copies became separate root traces. ``` before after temporal.workflow.Parent temporal.workflow.Parent ├─ agent_run [parent_agent] └─ agent_run [parent_agent] │ └─ delegate_task ├─ chat test ├─ temporal.workflow.SubAgent <- sibling ├─ delegate_task agent_run [parent_agent] <- replay copy, │ └─ temporal.workflow.SubAgent └─ ... new root │ └─ agent_run [sub_agent] └─ chat test ``` The fix: - Parenting: the workflow outbound interceptor now uses the current span as the parent when it is a real span other than the workflow span, and falls back to the workflow span otherwise. - Core: `logger.py` gets an internal, context-local span hook. When a hook is set, `SpanImpl` takes default row, span, and trace IDs from it, and `log_internal` drops writes the hook rejects. Spans only look up the hook after one has been installed, so code that never uses Temporal pays for one boolean check per span. - Replay safety: each workflow run installs a hook that numbers spans per run and derives their IDs from the run ID and that number. Workflow code is deterministic, so a replayed span gets the ID it was first logged with. The hook drops writes while replaying. The plugin now always creates the workflow span, so it is the current span during replay too. This also fixes the workflow span's end time. Evicting a workflow unwinds its coroutine, which ended the span at the first eviction. The span now ends when the workflow really finishes, because the later end row merges into the same span. `test_temporal(latest)` now also installs `pydantic-ai-slim[temporal]`, pinned to the `pydantic-ai-integration` latest version. Pydantic AI's temporal extra needs temporalio>=1.27 and Python 3.11+, so older temporalio entries skip it. The new test covers both the deprecated `TemporalAgent` wrapper and the newer `TemporalDurability` capability (`metadata={"temporal": False}` on the delegating tool), with and without the workflow cache. On the capability path, Pydantic AI routes the model request to the activity before our wrapper runs, so the `chat` span sits under the model activity instead of around it. Each request is still traced once. Limitations: replay safety only applies to workflows that take the existing `braintrust-temporal-workflow-span-context-v1` patch branch. Spans opened in signal or update handlers do not get deterministic IDs yet, because those handlers run in their own context. Co-Authored-By: Claude Opus 5.5 --- py/noxfile.py | 5 + .../integrations/temporal/plugin.py | 108 +++++++++--- .../integrations/temporal/test_temporal.py | 130 +++++++++++++- .../temporal/test_temporal_pydantic_ai.py | 160 ++++++++++++++++++ py/src/braintrust/logger.py | 55 +++++- py/src/braintrust/test_logger.py | 52 ++++++ 6 files changed, 480 insertions(+), 30 deletions(-) create mode 100644 py/src/braintrust/integrations/temporal/test_temporal_pydantic_ai.py diff --git a/py/noxfile.py b/py/noxfile.py index 5a3003533..cac9052e2 100644 --- a/py/noxfile.py +++ b/py/noxfile.py @@ -827,6 +827,11 @@ def test_transformers(session, version): def test_temporal(session, version): _install_test_deps(session) _install_matrix_dep(session, "temporalio", version) + # Pydantic AI's Temporal durability extra needs temporalio>=1.27 and Python 3.11+, + # so the Temporal + Pydantic AI tests only run against the latest temporalio pin. + if version == LATEST and sys.version_info >= (3, 11): + pydantic_ai_version = _MATRIX["pydantic-ai-integration"]["latest"].split("==", 1)[1] + _session_install(session, f"pydantic-ai-slim[temporal]=={pydantic_ai_version}", silent=SILENT_INSTALLS) _run_tests(session, f"{INTEGRATION_DIR}/temporal", version=version) diff --git a/py/src/braintrust/integrations/temporal/plugin.py b/py/src/braintrust/integrations/temporal/plugin.py index 60881c9f9..a0c1baa5c 100644 --- a/py/src/braintrust/integrations/temporal/plugin.py +++ b/py/src/braintrust/integrations/temporal/plugin.py @@ -256,6 +256,8 @@ def __init__(self, next: temporalio.worker.WorkflowInboundInterceptor) -> None: self.payload_converter = temporalio.converter.PayloadConverter.default self._parent_span_context: Any | None = None self._workflow_span: Any | None = None + self._span_hooks: _WorkflowSpanHooks | None = None + self._span_hooks_token: Any | None = None def init(self, outbound: temporalio.worker.WorkflowOutboundInterceptor) -> None: super().init(_BraintrustWorkflowOutboundInterceptor(outbound, self)) @@ -282,25 +284,31 @@ async def execute_workflow(self, input: temporalio.worker.ExecuteWorkflowInput) parent = parent_span_context or logger.export() if temporalio.workflow.patched(_SPAN_CONTEXT_PATCH_ID): ids = _workflow_span_ids(info, parent) - if temporalio.workflow.unsafe.is_replaying(): - self._parent_span_context = _workflow_span_context(parent, ids) - else: - span = logger.start_span( - internal={"instrumentation": "temporal-auto"}, - name=f"temporal.workflow.{info.workflow_type}", - type="task", - parent=parent, - id=ids["row_id"], - span_id=ids["span_id"], - root_span_id=ids["root_span_id"], - metadata={ - "temporal.workflow_type": info.workflow_type, - "temporal.workflow_id": info.workflow_id, - "temporal.run_id": info.run_id, - }, - ) - span.set_current() - self._parent_span_context = _workflow_span_context(parent, ids) + # Spans opened in workflow code (the workflow span, user spans, Pydantic AI + # agent/tool spans) get deterministic IDs and are not written while replaying, + # so a replay rebuilds the same span tree instead of logging a second copy. + self._span_hooks = _WorkflowSpanHooks( + _workflow_span_id_base(info), legacy_ids=SpanComponentsV4.get_version(parent) < 4 + ) + self._span_hooks_token = braintrust.logger._internal_set_span_hooks(self._span_hooks) + # Always create the workflow span, even on replay, so it is the current span for + # in-workflow spans and its end is logged if the workflow finishes after a replay. + span = logger.start_span( + internal={"instrumentation": "temporal-auto"}, + name=f"temporal.workflow.{info.workflow_type}", + type="task", + parent=parent, + id=ids["row_id"], + span_id=ids["span_id"], + root_span_id=ids["root_span_id"], + metadata={ + "temporal.workflow_type": info.workflow_type, + "temporal.workflow_id": info.workflow_id, + "temporal.run_id": info.run_id, + }, + ) + span.set_current() + self._parent_span_context = _workflow_span_context(parent, ids) elif not temporalio.workflow.unsafe.is_replaying(): span = logger.start_span( internal={"instrumentation": "temporal-auto"}, @@ -326,10 +334,13 @@ async def execute_workflow(self, input: temporalio.worker.ExecuteWorkflowInput) self._workflow_span.log(error=e) raise finally: - if self._workflow_span: - with temporalio.workflow.unsafe.sandbox_unrestricted(): + with temporalio.workflow.unsafe.sandbox_unrestricted(): + if self._workflow_span: self._workflow_span.unset_current() self._workflow_span.end() + if self._span_hooks_token is not None: + braintrust.logger._internal_reset_span_hooks(self._span_hooks_token) + self._span_hooks_token = None class _BraintrustWorkflowOutboundInterceptor(temporalio.worker.WorkflowOutboundInterceptor): @@ -348,12 +359,25 @@ def _add_span_context_to_headers( ) -> Mapping[str, temporalio.api.common.v1.Payload]: # 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]) + parent = self._current_parent_span_context() + if parent: + payloads = self.root.payload_converter.to_payloads([parent]) if payloads: return {**headers, _HEADER_KEY: payloads[0]} return headers + def _current_parent_span_context(self) -> Any | None: + # Prefer a span opened inside workflow code (e.g. a Pydantic AI tool span) so activities + # and child workflows nest under it. Without span hooks (workflows started before the + # span-context patch), spans created during replay have fresh IDs, so only trust the + # current span on real execution there. + if self.root._span_hooks is not None or not temporalio.workflow.unsafe.is_replaying(): + with temporalio.workflow.unsafe.sandbox_unrestricted(): + span = _current_span() + if span is not braintrust.logger.NOOP_SPAN and span is not self.root._workflow_span: + return span.export() + return self.root._parent_span_context + def start_activity(self, input: temporalio.worker.StartActivityInput) -> temporalio.workflow.ActivityHandle: input.headers = self._add_span_context_to_headers(input.headers) return super().start_activity(input) @@ -378,10 +402,42 @@ def _modify_workflow_runner(existing: WorkflowRunner | None) -> WorkflowRunner | return existing -def _workflow_span_ids(info: temporalio.workflow.Info, parent: str) -> dict[str, str]: +class _WorkflowSpanHooks: + """Span hooks for one workflow run: deterministic IDs and no writes while replaying. + + Workflow code is deterministic, so the Nth span opened in a run is the same span on every + replay. Numbering spans per run therefore gives a replayed span the ID it was first logged with. + """ + + def __init__(self, base: str, legacy_ids: bool) -> None: + self._base = base + self._legacy_ids = legacy_ids + self._counter = 0 + + def next_ids(self) -> tuple[str, str, str]: + self._counter += 1 + ids = _deterministic_span_ids(f"{self._base}:span:{self._counter}", self._legacy_ids) + return ids["row_id"], ids["span_id"], ids["root_span_id"] + + def should_write(self) -> bool: + try: + return not temporalio.workflow.unsafe.is_replaying() + except Exception: + # Not in workflow code (e.g. the span is ended from outside the workflow event loop). + return True + + +def _workflow_span_id_base(info: temporalio.workflow.Info) -> str: namespace = getattr(info, "namespace", "") - base = f"{_SPAN_ID_NAMESPACE}:{namespace}:{info.workflow_type}:{info.workflow_id}:{info.run_id}" - if SpanComponentsV4.get_version(parent) >= 4: + return f"{_SPAN_ID_NAMESPACE}:{namespace}:{info.workflow_type}:{info.workflow_id}:{info.run_id}" + + +def _workflow_span_ids(info: temporalio.workflow.Info, parent: str) -> dict[str, str]: + return _deterministic_span_ids(_workflow_span_id_base(info), SpanComponentsV4.get_version(parent) < 4) + + +def _deterministic_span_ids(base: str, legacy_ids: bool) -> dict[str, str]: + if not legacy_ids: digest = hashlib.sha256(base.encode()).hexdigest() return { "row_id": str(uuid.uuid5(uuid.NAMESPACE_URL, base)), diff --git a/py/src/braintrust/integrations/temporal/test_temporal.py b/py/src/braintrust/integrations/temporal/test_temporal.py index 2260fdc40..6ec3dbaee 100644 --- a/py/src/braintrust/integrations/temporal/test_temporal.py +++ b/py/src/braintrust/integrations/temporal/test_temporal.py @@ -219,6 +219,29 @@ async def run(self, input: TaskInput) -> int: return child_result +@temporalio.workflow.defn +class DelegatingParentWorkflow: + """Parent workflow that starts a child workflow and an activity inside an in-workflow span. + + Mirrors Pydantic AI's TemporalAgent, where a tool such as ``delegate_task`` runs in + workflow code and starts a subagent child workflow. + """ + + @temporalio.workflow.run + async def run(self, input: TaskInput) -> int: + with braintrust.start_span(name="delegate_task", type="tool"): + child_result = await temporalio.workflow.execute_child_workflow( + ChildWorkflow.run, + input, + id=f"child-{temporalio.workflow.info().workflow_id}", + ) + return await temporalio.workflow.execute_activity( + simple_activity, + TaskInput(value=child_result), + start_to_close_timeout=timedelta(seconds=10), + ) + + @temporalio.workflow.defn class ReplayAfterSignalWorkflow: """Workflow that schedules an activity after a later workflow task.""" @@ -252,6 +275,37 @@ def state(self) -> str: return self._state +@temporalio.workflow.defn +class InWorkflowSpanReplayWorkflow: + """Opens in-workflow spans across workflow tasks, like a Pydantic AI agent_run span.""" + + def __init__(self) -> None: + self._continue = False + self._state = "starting" + + @temporalio.workflow.run + async def run(self, input: TaskInput) -> int: + with braintrust.start_span(name="agent_run", type="task"): + with braintrust.start_span(name="first_tool", type="tool"): + first = await temporalio.workflow.execute_activity( + simple_activity, input, start_to_close_timeout=timedelta(seconds=10) + ) + self._state = "waiting" + await temporalio.workflow.wait_condition(lambda: self._continue) + with braintrust.start_span(name="second_tool", type="tool"): + return await temporalio.workflow.execute_activity( + simple_activity, TaskInput(value=first), start_to_close_timeout=timedelta(seconds=10) + ) + + @temporalio.workflow.signal + def continue_workflow(self) -> None: + self._continue = True + + @temporalio.workflow.query + def state(self) -> str: + return self._state + + class TestAutoInstrumentation: """Tests for Temporal auto-instrumentation helpers.""" @@ -307,7 +361,7 @@ def _spans_named(spans: list[dict[str, Any]], name: str) -> list[dict[str, Any]] async def _wait_for_workflow_state(handle: Any, expected: str) -> None: for _ in range(50): - if await handle.query(ReplayAfterSignalWorkflow.state) == expected: + if await handle.query("state") == expected: return await asyncio.sleep(0.1) raise AssertionError(f"Workflow did not reach state {expected!r}") @@ -366,6 +420,9 @@ async def test_plugin_activity_after_replay_stays_under_workflow_span( assert len(workflow_spans) == 1 assert len(activity_spans) == 2 workflow_span = workflow_spans[0] + # Evicting the workflow between tasks unwinds its coroutine; the span must still end when the + # workflow really finishes, not when it was first evicted. + assert workflow_span["metrics"]["end"] >= max(span["metrics"]["start"] for span in activity_spans) assert workflow_span["context"]["span_origin"]["instrumentation"]["name"] == "temporal-auto" assert all( span["context"]["span_origin"]["instrumentation"]["name"] == "temporal-auto" for span in activity_spans @@ -449,3 +506,74 @@ async def test_plugin_workflow_tracing( assert [bool(span.get("error")) for span in activity_spans] == expected_activity_errors if any(expected_activity_errors): assert "Simulated failure on first attempt" in activity_spans[0]["error"] + + @pytest.mark.asyncio + async def test_child_workflow_and_activity_nest_under_in_workflow_span(self, temporal_env, memory_logger): + task_queue = f"test-queue-delegate-{uuid.uuid4()}" + async with Worker( + temporal_env.client, + task_queue=task_queue, + workflows=[DelegatingParentWorkflow, ChildWorkflow], + activities=[simple_activity], + plugins=[BraintrustPlugin(logger=memory_logger)], + ): + await temporal_env.client.execute_workflow( + DelegatingParentWorkflow.run, + TaskInput(value=1), + id=f"test-workflow-{uuid.uuid4()}", + task_queue=task_queue, + ) + + braintrust.flush() + spans = memory_logger.pop() + (parent_wf,) = _spans_named(spans, "temporal.workflow.DelegatingParentWorkflow") + (tool,) = _spans_named(spans, "delegate_task") + (child_wf,) = _spans_named(spans, "temporal.workflow.ChildWorkflow") + activities = _spans_named(spans, "temporal.activity.simple_activity") + + assert tool["span_parents"] == [parent_wf["span_id"]] + assert child_wf["span_parents"] == [tool["span_id"]] + parent_activity = next( + a + for a in activities + if a["metadata"]["temporal.workflow_id"] != child_wf["metadata"]["temporal.workflow_id"] + ) + assert parent_activity["span_parents"] == [tool["span_id"]] + + @pytest.mark.asyncio + async def test_in_workflow_spans_are_not_duplicated_on_replay(self, temporal_env, memory_logger): + task_queue = f"test-queue-span-replay-{uuid.uuid4()}" + async with Worker( + temporal_env.client, + task_queue=task_queue, + workflows=[InWorkflowSpanReplayWorkflow], + activities=[simple_activity], + max_cached_workflows=0, + plugins=[BraintrustPlugin(logger=memory_logger)], + ): + handle = await temporal_env.client.start_workflow( + InWorkflowSpanReplayWorkflow.run, + TaskInput(value=10), + id=f"test-workflow-{uuid.uuid4()}", + task_queue=task_queue, + ) + await _wait_for_workflow_state(handle, "waiting") + await handle.signal(InWorkflowSpanReplayWorkflow.continue_workflow) + assert await handle.result() == 30 + + braintrust.flush() + spans = memory_logger.pop() + (workflow_span,) = _spans_named(spans, "temporal.workflow.InWorkflowSpanReplayWorkflow") + (agent_run,) = _spans_named(spans, "agent_run") + (first_tool,) = _spans_named(spans, "first_tool") + (second_tool,) = _spans_named(spans, "second_tool") + # Spans still open across a replay must end when the workflow really finishes. + assert workflow_span["metrics"]["end"] >= second_tool["metrics"]["start"] + assert agent_run["metrics"]["end"] >= second_tool["metrics"]["start"] + assert agent_run["span_parents"] == [workflow_span["span_id"]] + assert first_tool["span_parents"] == [agent_run["span_id"]] + assert second_tool["span_parents"] == [agent_run["span_id"]] + activities = _spans_named(spans, "temporal.activity.simple_activity") + assert sorted(a["span_parents"][0] for a in activities) == sorted( + [first_tool["span_id"], second_tool["span_id"]] + ) diff --git a/py/src/braintrust/integrations/temporal/test_temporal_pydantic_ai.py b/py/src/braintrust/integrations/temporal/test_temporal_pydantic_ai.py new file mode 100644 index 000000000..66064342f --- /dev/null +++ b/py/src/braintrust/integrations/temporal/test_temporal_pydantic_ai.py @@ -0,0 +1,160 @@ +"""Tests for BraintrustPlugin combined with Pydantic AI's Temporal durability integration.""" + +import uuid +import warnings + +import pytest + + +pytest.importorskip("temporalio") +pytest.importorskip("pydantic_ai.durable_exec.temporal") + +import braintrust +import pydantic_ai.durable_exec.temporal as pydantic_ai_temporal +import temporalio.workflow +from braintrust.integrations.pydantic_ai import PydanticAIIntegration +from braintrust.integrations.temporal import BraintrustPlugin +from braintrust.integrations.temporal.test_temporal import _spans_named, memory_logger, temporal_env # noqa: F401 +from pydantic_ai import Agent +from pydantic_ai.durable_exec.temporal import AgentPlugin, PydanticAIPlugin +from pydantic_ai.models.test import TestModel +from temporalio.client import Client +from temporalio.worker import Worker + + +_HAS_DURABILITY_CAPABILITY = hasattr(pydantic_ai_temporal, "TemporalDurability") + + +def _api_for_current_workflow() -> str: + # Workflow IDs start with the API under test; see the test below. + return temporalio.workflow.info().workflow_id.split(":", 1)[0] + + +async def _delegate(task: str) -> str: + # Runs in workflow code and starts the sub-agent as a child workflow, like a `delegate_task` tool. + return await temporalio.workflow.execute_child_workflow( + SubAgentWorkflow.run, + task, + id=f"{temporalio.workflow.info().workflow_id}:sub", + ) + + +# Deprecated `TemporalAgent` wrapper: the only API on older pydantic-ai releases. +with warnings.catch_warnings(): + warnings.simplefilter("ignore") + _wrapper_sub = pydantic_ai_temporal.TemporalAgent( + Agent(TestModel(custom_output_text="sub done"), name="wrapper_sub_agent") + ) + _wrapper_parent_agent = Agent(TestModel(call_tools=["delegate_task"]), name="wrapper_parent_agent") + _wrapper_parent_agent.tool_plain(name="delegate_task")(_delegate) + # delegate_task starts a child workflow, so it has to run in workflow code, not an activity. + _wrapper_parent = pydantic_ai_temporal.TemporalAgent( + _wrapper_parent_agent, tool_activity_config={"": {"delegate_task": False}} + ) + +if _HAS_DURABILITY_CAPABILITY: + _capability_sub = Agent( + TestModel(custom_output_text="sub done"), + name="capability_sub_agent", + capabilities=[pydantic_ai_temporal.TemporalDurability()], + ) + _capability_parent = Agent( + TestModel(call_tools=["delegate_task"]), + name="capability_parent_agent", + capabilities=[pydantic_ai_temporal.TemporalDurability()], + ) + # `metadata={"temporal": False}` replaces `tool_activity_config` on the capability path. + _capability_parent.tool_plain(name="delegate_task", metadata={"temporal": False})(_delegate) + +_AGENTS = { + "wrapper": (_wrapper_parent, _wrapper_sub), + **({"capability": (_capability_parent, _capability_sub)} if _HAS_DURABILITY_CAPABILITY else {}), +} + + +@temporalio.workflow.defn +class SubAgentWorkflow: + @temporalio.workflow.run + async def run(self, prompt: str) -> str: + _, sub_agent = _AGENTS[_api_for_current_workflow()] + return (await sub_agent.run(prompt)).output + + +@temporalio.workflow.defn +class ParentAgentWorkflow: + @temporalio.workflow.run + async def run(self, prompt: str) -> str: + parent_agent, _ = _AGENTS[_api_for_current_workflow()] + return (await parent_agent.run(prompt)).output + + +# With no workflow cache, every workflow task replays the run from the start, as after a worker +# restart or cache eviction. Replays must not log the in-workflow Pydantic AI spans again. +@pytest.mark.parametrize("max_cached_workflows", [1000, 0], ids=["cached", "replay"]) +@pytest.mark.parametrize( + "api", + [ + "wrapper", + pytest.param( + "capability", + marks=pytest.mark.skipif( + not _HAS_DURABILITY_CAPABILITY, reason="TemporalDurability needs a newer pydantic-ai" + ), + ), + ], +) +@pytest.mark.asyncio +async def test_subagent_child_workflow_nests_under_delegate_task_span( + temporal_env, # noqa: F811 + memory_logger, # noqa: F811 + api, + max_cached_workflows, +): + assert PydanticAIIntegration.setup() + parent_agent, sub_agent = _AGENTS[api] + client = await Client.connect( + temporal_env.client.service_client.config.target_host, + namespace=temporal_env.client.namespace, + plugins=[PydanticAIPlugin()], + ) + task_queue = f"test-queue-pydantic-ai-{api}-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[ParentAgentWorkflow, SubAgentWorkflow], + max_cached_workflows=max_cached_workflows, + plugins=[AgentPlugin(parent_agent), AgentPlugin(sub_agent), BraintrustPlugin(logger=memory_logger)], + ): + await client.execute_workflow( + ParentAgentWorkflow.run, + "delegate something", + id=f"{api}:{uuid.uuid4()}", + task_queue=task_queue, + ) + + braintrust.flush() + spans = memory_logger.pop() + (parent_workflow_span,) = _spans_named(spans, "temporal.workflow.ParentAgentWorkflow") + (parent_agent_span,) = _spans_named(spans, f"agent_run [{parent_agent.name}]") + (tool_span,) = _spans_named(spans, "delegate_task") + (child_workflow_span,) = _spans_named(spans, "temporal.workflow.SubAgentWorkflow") + (sub_agent_span,) = _spans_named(spans, f"agent_run [{sub_agent.name}]") + + assert parent_agent_span["span_parents"] == [parent_workflow_span["span_id"]] + assert tool_span["span_parents"] == [parent_agent_span["span_id"]] + assert child_workflow_span["span_parents"] == [tool_span["span_id"]] + assert sub_agent_span["span_parents"] == [child_workflow_span["span_id"]] + assert {span["root_span_id"] for span in spans} == {parent_workflow_span["root_span_id"]} + + # Each model request is traced once: one `chat` span paired with its model activity span. The + # wrapper API traces the request in workflow code (activity under `chat`); the capability + # API hands it to the activity first (`chat` under activity). + chat_spans = [span for span in spans if span["span_attributes"]["name"].startswith("chat ")] + model_activity_spans = [span for span in spans if span["span_attributes"]["name"].endswith("__model_request")] + assert len(chat_spans) == len(model_activity_spans) == 3 + pairs = { + frozenset((span["span_id"], span["span_parents"][0])) + for span in chat_spans + model_activity_spans + if span["span_parents"][0] in {other["span_id"] for other in chat_spans + model_activity_spans} + } + assert len(pairs) == 3 diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 41dbb0277..f89c48523 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -35,6 +35,7 @@ Any, Generic, Literal, + Protocol, TypedDict, TypeVar, cast, @@ -4987,6 +4988,40 @@ def as_dataset(self, batch_size: int | None = None) -> Iterator[ExperimentDatase _EXEC_COUNTER = 0 +class _SpanHooks(Protocol): + """Internal hooks that let a runtime control span IDs and writes in its own context. + + Used by the Temporal integration: spans opened in workflow code need deterministic IDs and must + not be written while the workflow replays its history. Not a public API. + """ + + def next_ids(self) -> tuple[str, str, str]: + """Return ``(row_id, span_id, trace_id)`` for a new span that has no explicit IDs.""" + ... + + def should_write(self) -> bool: + """Return False to drop a span write.""" + ... + + +_span_hooks: contextvars.ContextVar["_SpanHooks | None"] = contextvars.ContextVar( + "braintrust_span_hooks", default=None +) +# Flipped the first time hooks are installed, so spans skip the context lookup until a runtime uses them. +_span_hooks_enabled = False + + +def _internal_set_span_hooks(hooks: "_SpanHooks | None") -> contextvars.Token: + """Install span hooks for the current context. Returns a token for `_internal_reset_span_hooks`.""" + global _span_hooks_enabled + _span_hooks_enabled = True + return _span_hooks.set(hooks) + + +def _internal_reset_span_hooks(token: contextvars.Token) -> None: + _span_hooks.reset(token) + + class SpanImpl(Span): """Primary implementation of the `Span` interface. See the `Span` interface for full details on each method. @@ -5091,14 +5126,26 @@ def __init__( # TODO: can be simplified after `event` is typed. id = event.pop("id", None) - if id is None or not isinstance(id, str): + if not isinstance(id, str): + id = None + span_id = internal.get("span_id", span_id) + root_span_id = internal.get("root_span_id", root_span_id) + + self._hooks = _span_hooks.get() if _span_hooks_enabled else None + if self._hooks is not None and (id is None or span_id is None): + hook_row_id, hook_span_id, hook_trace_id = self._hooks.next_ids() + id = id or hook_row_id + span_id = span_id or hook_span_id + root_span_id = root_span_id or hook_trace_id + + if id is None: id = str(uuid.uuid4()) self._id = id # Resolve all span IDs (span_id, root_span_id, span_parents) span_ids = _resolve_span_ids( - span_id=internal.get("span_id", span_id), - root_span_id=internal.get("root_span_id", root_span_id), + span_id=span_id, + root_span_id=root_span_id, parent_span_ids=parent_span_ids, lookup_span_parent=lookup_span_parent, id_generator=self.state.id_generator, @@ -5148,6 +5195,8 @@ def log(self, **event: Any) -> None: return self.log_internal(event=event, internal_data=None) def log_internal(self, event: dict[str, Any] | None = None, internal_data: dict[str, Any] | None = None) -> None: + if self._hooks is not None and not self._hooks.should_write(): + return serializable_partial_record, lazy_partial_record = split_logging_data(event, internal_data) # We both check for serializability and round-trip `partial_record` diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index 9bb585fac..514e4918f 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -3829,3 +3829,55 @@ def test_proxy_conn_base_url(proxy_url, expected_base_url): state = BraintrustState() state.proxy_url = proxy_url assert state.proxy_conn().base_url == expected_base_url + + +class _RecordingSpanHooks: + def __init__(self): + self.counter = 0 + self.write = True + + def next_ids(self): + self.counter += 1 + return f"row-{self.counter}", f"span-{self.counter}", f"trace-{self.counter}" + + def should_write(self): + return self.write + + +def test_span_hooks_supply_default_ids_and_gate_writes(with_memory_logger): + hooks = _RecordingSpanHooks() + test_logger = init_test_logger(__name__) + token = braintrust.logger._internal_set_span_hooks(hooks) + try: + with test_logger.start_span(name="root") as root: + with root.start_span(name="child"): + pass + # Explicit IDs win over the hooks and do not consume a hook ID. + with test_logger.start_span(name="explicit", id="explicit-row", span_id="explicit-span"): + pass + hooks.write = False + with root.start_span(name="dropped"): + pass + hooks.write = True + finally: + braintrust.logger._internal_reset_span_hooks(token) + + with test_logger.start_span(name="after_reset") as after_reset: + pass + + rows = {row["span_attributes"]["name"]: row for row in with_memory_logger.pop()} + assert set(rows) == {"root", "child", "explicit", "after_reset"} + assert (rows["root"]["id"], rows["root"]["span_id"], rows["root"]["root_span_id"]) == ( + "row-1", + "span-1", + "trace-1", + ) + assert (rows["child"]["id"], rows["child"]["span_id"], rows["child"]["root_span_id"]) == ( + "row-2", + "span-2", + "trace-1", + ) + assert rows["child"]["span_parents"] == ["span-1"] + assert (rows["explicit"]["id"], rows["explicit"]["span_id"]) == ("explicit-row", "explicit-span") + assert hooks.counter == 3 + assert not after_reset.span_id.startswith("span-") From 56ccf6edb08c5009c24f6d3833c805216502b7f4 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 3 Nov 2026 11:19:54 -0500 Subject: [PATCH 2/3] ci(pylint): skip temporal pydantic-ai test when pydantic-ai is unlintable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > 🤖 **This text was AI-generated.** On Python 3.10 the `pylint` session does not install pydantic-ai, because its latest release needs Python 3.11+. The session already drops the `pydantic_ai` integration's own files from the lint set in that case. The new `temporal/test_temporal_pydantic_ai.py` imports pydantic-ai from outside that directory, so pylint still linted it and failed `static_checks (3.10)`: ``` src/braintrust/integrations/temporal/test_temporal_pydantic_ai.py:13:0: E0401: Unable to import 'pydantic_ai.durable_exec.temporal' (import-error) ``` The session now has a small `cross_integration_files` map that ties files outside an integration's directory to that integration. They are skipped together whenever the integration cannot be linted: ```python cross_integration_files = { "pydantic_ai": {f"src/{INTEGRATION_DIR}/temporal/test_temporal_pydantic_ai.py"}, } ``` Co-Authored-By: Claude Opus 5.5 --- py/noxfile.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/py/noxfile.py b/py/noxfile.py index cac9052e2..5ecb41e73 100644 --- a/py/noxfile.py +++ b/py/noxfile.py @@ -1041,11 +1041,16 @@ def pylint(session): # The lint group skips crewai to avoid vulnerable transitive chromadb # versions, so skip the matching example too. files = [f for f in files if not f.startswith("../examples/crewai/")] + # Tests outside an integration's own directory that import it. + cross_integration_files = { + "pydantic_ai": {f"src/{INTEGRATION_DIR}/temporal/test_temporal_pydantic_ai.py"}, + } for integration in sorted(unlintable_integrations): files = [ f for f in files if not f.startswith(f"src/{INTEGRATION_DIR}/{integration}/") + and f not in cross_integration_files.get(integration, set()) and f != f"src/{INTEGRATION_DIR}/auto_test_scripts/test_auto_{integration}.py" and not f.startswith(f"src/{INTEGRATION_DIR}/auto_test_scripts/test_manual_wrap_{integration}_") and not f.startswith(f"../examples/{integration}/") From 7acf6e5032388ccfe86006de41c1643f7672cab3 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 3 Nov 2026 11:46:16 -0500 Subject: [PATCH 3/3] fix(temporal): route workflow span writes to the plugin logger on every thread MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > 🤖 **This text was AI-generated.** `BraintrustPlugin(logger=...)` sends writes to its logger through `_override_bg_logger`, which is thread-local. The plugin set it only on the thread that started the workflow. Temporal can run later activations of the same workflow on other worker threads, so writes made there skipped the plugin's logger and went to the global logger instead: ``` activation 1 (thread A, override set) -> span start row -> plugin logger activation 2 (thread B, no override) -> span end row -> global logger ``` This split spans across loggers and left rows with no start, which showed up in CI as `KeyError: 'span_attributes'` in the Pydantic AI + Temporal test. It only affects users who pass `logger=`. With the default, every thread uses the global logger anyway. The workflow span hook now re-applies the plugin's logger on the current thread before every write it allows, using the same `__braintrust_get_logger` call the workflow makes at startup. The activity path already did this per activity. A new test case runs workflow activations on a fresh thread every time, so the bug reproduces every time instead of by chance. Co-Authored-By: Claude Opus 5.5 --- .../integrations/temporal/plugin.py | 16 ++++++--- .../integrations/temporal/test_temporal.py | 34 ++++++++++++++++++- 2 files changed, 45 insertions(+), 5 deletions(-) diff --git a/py/src/braintrust/integrations/temporal/plugin.py b/py/src/braintrust/integrations/temporal/plugin.py index a0c1baa5c..3841ad094 100644 --- a/py/src/braintrust/integrations/temporal/plugin.py +++ b/py/src/braintrust/integrations/temporal/plugin.py @@ -82,7 +82,7 @@ import dataclasses import hashlib import uuid -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import Any import braintrust @@ -288,7 +288,9 @@ async def execute_workflow(self, input: temporalio.worker.ExecuteWorkflowInput) # agent/tool spans) get deterministic IDs and are not written while replaying, # so a replay rebuilds the same span tree instead of logging a second copy. self._span_hooks = _WorkflowSpanHooks( - _workflow_span_id_base(info), legacy_ids=SpanComponentsV4.get_version(parent) < 4 + _workflow_span_id_base(info), + legacy_ids=SpanComponentsV4.get_version(parent) < 4, + bind_logger=get_logger, ) self._span_hooks_token = braintrust.logger._internal_set_span_hooks(self._span_hooks) # Always create the workflow span, even on replay, so it is the current span for @@ -409,9 +411,12 @@ class _WorkflowSpanHooks: replay. Numbering spans per run therefore gives a replayed span the ID it was first logged with. """ - def __init__(self, base: str, legacy_ids: bool) -> None: + def __init__(self, base: str, legacy_ids: bool, bind_logger: Callable[[], Any]) -> None: self._base = base self._legacy_ids = legacy_ids + # Points this thread at the plugin's logger. The logger override is thread-local, and + # Temporal can run each activation of a workflow on a different worker thread. + self._bind_logger = bind_logger self._counter = 0 def next_ids(self) -> tuple[str, str, str]: @@ -421,10 +426,13 @@ def next_ids(self) -> tuple[str, str, str]: def should_write(self) -> bool: try: - return not temporalio.workflow.unsafe.is_replaying() + if temporalio.workflow.unsafe.is_replaying(): + return False except Exception: # Not in workflow code (e.g. the span is ended from outside the workflow event loop). return True + self._bind_logger() + return True def _workflow_span_id_base(info: temporalio.workflow.Info) -> str: diff --git a/py/src/braintrust/integrations/temporal/test_temporal.py b/py/src/braintrust/integrations/temporal/test_temporal.py index 6ec3dbaee..2d1e1554f 100644 --- a/py/src/braintrust/integrations/temporal/test_temporal.py +++ b/py/src/braintrust/integrations/temporal/test_temporal.py @@ -1,7 +1,9 @@ """Unit tests for Braintrust Temporal interceptor.""" import asyncio +import concurrent.futures import os +import threading import uuid from collections import Counter from dataclasses import dataclass @@ -326,6 +328,28 @@ def test_contrib_temporal_compat_import_deprecated(self): # Integration Tests +class _NewThreadPerTaskExecutor(concurrent.futures.ThreadPoolExecutor): + """Runs every workflow activation on a new thread. + + A shared worker pool can hand each activation of a workflow to a different thread; this makes + that happen every time, so thread-local state cannot be relied on between activations. + """ + + def submit(self, fn, /, *args, **kwargs): + future: concurrent.futures.Future = concurrent.futures.Future() + + def run() -> None: + if not future.set_running_or_notify_cancel(): + return + try: + future.set_result(fn(*args, **kwargs)) + except BaseException as e: + future.set_exception(e) + + threading.Thread(target=run, daemon=True).start() + return future + + @pytest_asyncio.fixture(scope="function") async def temporal_env(): """Create a Temporal test environment. @@ -508,13 +532,18 @@ async def test_plugin_workflow_tracing( assert "Simulated failure on first attempt" in activity_spans[0]["error"] @pytest.mark.asyncio - async def test_child_workflow_and_activity_nest_under_in_workflow_span(self, temporal_env, memory_logger): + @pytest.mark.parametrize("new_thread_per_task", [False, True], ids=["shared_pool", "new_thread_per_task"]) + async def test_child_workflow_and_activity_nest_under_in_workflow_span( + self, temporal_env, memory_logger, new_thread_per_task + ): task_queue = f"test-queue-delegate-{uuid.uuid4()}" + executor = _NewThreadPerTaskExecutor() if new_thread_per_task else None async with Worker( temporal_env.client, task_queue=task_queue, workflows=[DelegatingParentWorkflow, ChildWorkflow], activities=[simple_activity], + workflow_task_executor=executor, plugins=[BraintrustPlugin(logger=memory_logger)], ): await temporal_env.client.execute_workflow( @@ -526,6 +555,9 @@ async def test_child_workflow_and_activity_nest_under_in_workflow_span(self, tem braintrust.flush() spans = memory_logger.pop() + # Every row of every span reaches the plugin's logger, whichever thread ran the activation. + assert all("span_attributes" in span for span in spans) + assert all(span["metrics"].get("end") is not None for span in spans) (parent_wf,) = _spans_named(spans, "temporal.workflow.DelegatingParentWorkflow") (tool,) = _spans_named(spans, "delegate_task") (child_wf,) = _spans_named(spans, "temporal.workflow.ChildWorkflow")