From 71dd6bd5f49df2b8ccf22fce7e38493b589d2f40 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 3 Nov 2026 11:29:41 -0500 Subject: [PATCH 1/2] fix(temporal): preserve trace context across continue-as-new MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > 🤖 **This text was AI-generated.** BraintrustPlugin did not attach span context to continue-as-new commands. Continued runs could lose the original trace or skip the span that requested the continuation. Propagate the current span through Temporal headers using the same helper as activities and child workflows. Each run keeps its own workflow span while remaining in the original trace: ```text before: run 1 run 2 run 3 after: run 1 └─ run 2 └─ nested span └─ run 3 ``` Cover both transitions in one real Temporal workflow chain, with and without a client parent and with cached and replayed workflow execution. Span end-time handling remains outside this change. Fixes https://linear.app/braintrustdata/issue/SDK-780 --- .../integrations/temporal/plugin.py | 7 +- .../integrations/temporal/test_temporal.py | 69 +++++++++++++++++++ 2 files changed, 75 insertions(+), 1 deletion(-) diff --git a/py/src/braintrust/integrations/temporal/plugin.py b/py/src/braintrust/integrations/temporal/plugin.py index 2211244e..2864cc20 100644 --- a/py/src/braintrust/integrations/temporal/plugin.py +++ b/py/src/braintrust/integrations/temporal/plugin.py @@ -76,6 +76,7 @@ - Trace local activities - Maintain parent-child relationships between client calls, workflows, and activities - Handle child workflows +- Preserve trace context across continue-as-new runs - Respect Temporal replay safety (no duplicate spans during replay) - Record cancellations as ``metadata["temporal.cancelled"] = True`` instead of errors """ @@ -85,7 +86,7 @@ import hashlib import uuid from collections.abc import Callable, Mapping -from typing import Any +from typing import Any, NoReturn import braintrust import temporalio.activity @@ -406,6 +407,10 @@ def start_child_workflow( input.headers = self._add_span_context_to_headers(input.headers) return super().start_child_workflow(input) + def continue_as_new(self, input: temporalio.worker.ContinueAsNewInput) -> NoReturn: + input.headers = self._add_span_context_to_headers(input.headers) + super().continue_as_new(input) + def _modify_workflow_runner(existing: WorkflowRunner | None) -> WorkflowRunner | None: if isinstance(existing, SandboxedWorkflowRunner): diff --git a/py/src/braintrust/integrations/temporal/test_temporal.py b/py/src/braintrust/integrations/temporal/test_temporal.py index f6b6c6dd..c053b19c 100644 --- a/py/src/braintrust/integrations/temporal/test_temporal.py +++ b/py/src/braintrust/integrations/temporal/test_temporal.py @@ -6,6 +6,7 @@ import threading import uuid from collections import Counter +from contextlib import nullcontext from dataclasses import dataclass from datetime import timedelta from typing import Any, cast @@ -328,6 +329,18 @@ async def run(self, mode: str) -> None: @temporalio.workflow.query def state(self) -> str: return "waiting" +@temporalio.workflow.defn +class ContinueAsNewWorkflow: + @temporalio.workflow.run + async def run(self, iteration: int) -> int: + await temporalio.workflow.execute_activity( + simple_activity, TaskInput(value=iteration), start_to_close_timeout=timedelta(seconds=10) + ) + if iteration < 2: + # Cover both workflow-level parenting and a nested span in the same chain. + with braintrust.start_span(name="continue") if iteration == 1 else nullcontext(): + temporalio.workflow.continue_as_new(iteration + 1) + return iteration class TestAutoInstrumentation: @@ -458,6 +471,62 @@ async def test_workflow_cancellation_metadata(self, temporal_env, memory_logger, assert span["metadata"]["temporal.cancelled"] is True assert not span.get("error") + @pytest.mark.parametrize("with_client_parent", [False, True]) + @pytest.mark.parametrize("max_cached_workflows", [1000, 0], ids=["cached", "replay"]) + @pytest.mark.asyncio + async def test_continue_as_new_preserves_trace_context( + self, temporal_env, memory_logger, with_client_parent, max_cached_workflows + ): + client = temporal_env.client + if with_client_parent: + client = await Client.connect( + client.service_client.config.target_host, + namespace=client.namespace, + plugins=[BraintrustPlugin(logger=memory_logger)], + ) + task_queue = f"test-queue-continue-{uuid.uuid4()}" + async with Worker( + temporal_env.client, + task_queue=task_queue, + workflows=[ContinueAsNewWorkflow], + activities=[simple_activity], + max_cached_workflows=max_cached_workflows, + plugins=[BraintrustPlugin(logger=memory_logger)], + ): + with braintrust.start_span(name="client-continue") if with_client_parent else nullcontext(): + assert ( + await client.execute_workflow( + ContinueAsNewWorkflow.run, + 0, + id=f"test-workflow-{uuid.uuid4()}", + task_queue=task_queue, + ) + == 2 + ) + + braintrust.flush() + spans = memory_logger.pop() + workflows = sorted( + _spans_named(spans, "temporal.workflow.ContinueAsNewWorkflow"), key=lambda span: span["metrics"]["start"] + ) + assert len(workflows) == 3 + assert len({span["metadata"]["temporal.run_id"] for span in workflows}) == 3 + assert len({span["span_id"] for span in workflows}) == 3 + assert len({span["root_span_id"] for span in spans}) == 1 + for span in workflows: + assert not span.get("error"), span + first, second, third = workflows + (continuation_span,) = _spans_named(spans, "continue") + assert second["span_parents"] == [first["span_id"]] + assert continuation_span["span_parents"] == [second["span_id"]] + assert third["span_parents"] == [continuation_span["span_id"]] + activities = _spans_named(spans, "temporal.activity.simple_activity") + assert len(activities) == 3 + assert {span["span_parents"][0] for span in activities} == {span["span_id"] for span in workflows} + if with_client_parent: + (client_span,) = _spans_named(spans, "client-continue") + assert workflows[0]["span_parents"] == [client_span["span_id"]] + @pytest.mark.parametrize("with_client_parent", [False, True]) @pytest.mark.asyncio async def test_plugin_activity_after_replay_stays_under_workflow_span( From 690ae661c5223909a6c10e8fa20475e01409e081 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 3 Nov 2026 13:26:25 -0500 Subject: [PATCH 2/2] style(temporal): format continue-as-new test --- py/src/braintrust/integrations/temporal/test_temporal.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/py/src/braintrust/integrations/temporal/test_temporal.py b/py/src/braintrust/integrations/temporal/test_temporal.py index c053b19c..b0e26872 100644 --- a/py/src/braintrust/integrations/temporal/test_temporal.py +++ b/py/src/braintrust/integrations/temporal/test_temporal.py @@ -329,6 +329,8 @@ async def run(self, mode: str) -> None: @temporalio.workflow.query def state(self) -> str: return "waiting" + + @temporalio.workflow.defn class ContinueAsNewWorkflow: @temporalio.workflow.run