From 4bb63364eda22e7c9fcc03d0bff848acf6c17189 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 3 Nov 2026 11:29:45 -0500 Subject: [PATCH 1/2] fix(temporal): distinguish cancellations from execution errors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > 🤖 **This text was AI-generated.** Stopping a chat through Temporal cancellation previously logged cancellation exceptions as errors, while asyncio cancellations had no cancellation marker. Use Temporal's cancellation classifier for workflow and activity spans, including cancellations wrapped in activity or child workflow exceptions: ```python # Before: cancellation could populate error or leave no marker. # After: span.log(metadata={"temporal.cancelled": True}) ``` Preserve exception propagation and span cleanup. Ordinary failures continue to populate error. Cover direct, activity, and child workflow cancellation alongside an ordinary failure with real Temporal workflows. Fixes: https://linear.app/braintrustdata/issue/SDK-782 --- .../integrations/temporal/plugin.py | 17 +++- .../integrations/temporal/test_temporal.py | 81 ++++++++++++++++++- 2 files changed, 93 insertions(+), 5 deletions(-) diff --git a/py/src/braintrust/integrations/temporal/plugin.py b/py/src/braintrust/integrations/temporal/plugin.py index 3841ad09..0c03949d 100644 --- a/py/src/braintrust/integrations/temporal/plugin.py +++ b/py/src/braintrust/integrations/temporal/plugin.py @@ -77,8 +77,10 @@ - Maintain parent-child relationships between client calls, workflows, and activities - Handle child workflows - Respect Temporal replay safety (no duplicate spans during replay) +- Record cancellations as ``metadata["temporal.cancelled"] = True`` instead of errors """ +import asyncio import dataclasses import hashlib import uuid @@ -90,6 +92,7 @@ import temporalio.api.common.v1 import temporalio.client import temporalio.converter +import temporalio.exceptions import temporalio.worker import temporalio.workflow from braintrust.span_identifier_v3 import SpanComponentsV3 @@ -235,8 +238,11 @@ async def execute_activity(self, input: temporalio.worker.ExecuteActivityInput) try: return await super().execute_activity(input) - except Exception as e: - span.log(error=e) + except (Exception, asyncio.CancelledError) as e: + if temporalio.exceptions.is_cancelled_exception(e): + span.log(metadata={"temporal.cancelled": True}) + else: + span.log(error=e) raise finally: span.unset_current() @@ -330,10 +336,13 @@ async def execute_workflow(self, input: temporalio.worker.ExecuteWorkflowInput) try: return await super().execute_workflow(input) - except Exception as e: + except (Exception, asyncio.CancelledError) as e: if self._workflow_span: with temporalio.workflow.unsafe.sandbox_unrestricted(): - self._workflow_span.log(error=e) + if temporalio.exceptions.is_cancelled_exception(e): + self._workflow_span.log(metadata={"temporal.cancelled": True}) + else: + self._workflow_span.log(error=e) raise finally: with temporalio.workflow.unsafe.sandbox_unrestricted(): diff --git a/py/src/braintrust/integrations/temporal/test_temporal.py b/py/src/braintrust/integrations/temporal/test_temporal.py index 2d1e1554..54df618c 100644 --- a/py/src/braintrust/integrations/temporal/test_temporal.py +++ b/py/src/braintrust/integrations/temporal/test_temporal.py @@ -21,6 +21,7 @@ import temporalio.activity import temporalio.api.common.v1 import temporalio.converter +import temporalio.exceptions import temporalio.testing import temporalio.worker import temporalio.workflow @@ -29,7 +30,7 @@ from braintrust.span_identifier_v3 import SpanComponentsV3, SpanObjectTypeV3 from braintrust.span_identifier_v4 import SpanComponentsV4 from braintrust.test_helpers import init_test_logger, preserve_env_vars -from temporalio.client import Client +from temporalio.client import Client, WorkflowFailureError from temporalio.common import RetryPolicy from temporalio.worker import Worker @@ -308,6 +309,42 @@ def state(self) -> str: return self._state +@temporalio.activity.defn +async def waiting_activity() -> None: + while True: + temporalio.activity.heartbeat() + await asyncio.sleep(0.1) + + +@temporalio.workflow.defn +class CancellationWorkflow: + @temporalio.workflow.run + async def run(self, mode: str) -> None: + if mode == "temporal": + raise temporalio.exceptions.CancelledError("Chat stopped") + if mode == "failure": + raise temporalio.exceptions.ApplicationError("Cancellation failed", non_retryable=True) + if mode == "activity": + await temporalio.workflow.execute_activity( + waiting_activity, + start_to_close_timeout=timedelta(minutes=5), + heartbeat_timeout=timedelta(seconds=5), + cancellation_type=temporalio.workflow.ActivityCancellationType.WAIT_CANCELLATION_COMPLETED, + ) + elif mode == "child": + await temporalio.workflow.execute_child_workflow( + CancellationWorkflow.run, + "timer", + id=f"child-{temporalio.workflow.info().workflow_id}", + ) + else: + await temporalio.workflow.wait_condition(lambda: False) + + @temporalio.workflow.query + def state(self) -> str: + return "waiting" + + class TestAutoInstrumentation: """Tests for Temporal auto-instrumentation helpers.""" @@ -394,6 +431,48 @@ async def _wait_for_workflow_state(handle: Any, expected: str) -> None: class TestBraintrustPluginIntegration: """Integration tests for BraintrustPlugin with real Temporal workflows.""" + @pytest.mark.parametrize("mode", ["temporal", "activity", "child", "failure"]) + @pytest.mark.asyncio + async def test_workflow_cancellation_metadata(self, temporal_env, memory_logger, mode): + task_queue = f"test-queue-{uuid.uuid4()}" + async with Worker( + temporal_env.client, + task_queue=task_queue, + workflows=[CancellationWorkflow], + activities=[waiting_activity], + plugins=[BraintrustPlugin(logger=memory_logger)], + ): + handle = await temporal_env.client.start_workflow( + CancellationWorkflow.run, + mode, + id=f"test-workflow-{uuid.uuid4()}", + task_queue=task_queue, + ) + if mode not in {"temporal", "failure"}: + await _wait_for_workflow_state(handle, "waiting") + await handle.cancel() + with pytest.raises(WorkflowFailureError) as exc: + await handle.result() + if mode == "failure": + assert isinstance(exc.value.cause, temporalio.exceptions.ApplicationError) + else: + assert isinstance(exc.value.cause, temporalio.exceptions.CancelledError) + + spans = memory_logger.pop() + expected_span_counts = {"temporal.workflow.CancellationWorkflow": 2 if mode == "child" else 1} + if mode == "activity": + expected_span_counts["temporal.activity.waiting_activity"] = 1 + assert Counter(span["span_attributes"]["name"] for span in spans) == expected_span_counts + for span in spans: + assert span["context"]["span_origin"]["instrumentation"]["name"] == "temporal-auto" + assert "end" in span["metrics"] + if mode == "failure": + assert "Cancellation failed" in span["error"] + assert "temporal.cancelled" not in span["metadata"] + else: + assert span["metadata"]["temporal.cancelled"] is True + assert not span.get("error") + @pytest.mark.parametrize("with_client_parent", [False, True]) @pytest.mark.asyncio async def test_plugin_activity_after_replay_stays_under_workflow_span( From 440f84e0daf6480596e694ef87e25ddf633df2b0 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 3 Nov 2026 11:58:42 -0500 Subject: [PATCH 2/2] test(temporal): verify eviction does not mark workflows cancelled MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > 🤖 **This text was AI-generated.** Cache eviction cancels workflow tasks internally. Extend the existing replay test to assert that a successful workflow has neither temporal.cancelled metadata nor an error. The replay span hook already suppresses these writes. The child cancellation test could inspect spans before the child finished cancelling. Wait for completion before checking the exported spans: ```text before: request child cancellation -> inspect spans after: request child cancellation -> await completion -> inspect spans ``` Refs: https://github.com/braintrustdata/braintrust-sdk-python/pull/1059#discussion_r4457597665 Refs: https://linear.app/braintrustdata/issue/SDK-782 --- py/src/braintrust/integrations/temporal/test_temporal.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/py/src/braintrust/integrations/temporal/test_temporal.py b/py/src/braintrust/integrations/temporal/test_temporal.py index 54df618c..3b40e0d3 100644 --- a/py/src/braintrust/integrations/temporal/test_temporal.py +++ b/py/src/braintrust/integrations/temporal/test_temporal.py @@ -336,6 +336,7 @@ async def run(self, mode: str) -> None: CancellationWorkflow.run, "timer", id=f"child-{temporalio.workflow.info().workflow_id}", + cancellation_type=temporalio.workflow.ChildWorkflowCancellationType.WAIT_CANCELLATION_COMPLETED, ) else: await temporalio.workflow.wait_condition(lambda: False) @@ -523,6 +524,8 @@ 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] + assert "temporal.cancelled" not in workflow_span["metadata"] + assert not workflow_span.get("error") # 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)