fix(temporal): distinguish cancellations from execution errors

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():

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,43 @@ 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}",
+                cancellation_type=temporalio.workflow.ChildWorkflowCancellationType.WAIT_CANCELLATION_COMPLETED,
+            )
+        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 +432,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(
@@ -444,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)