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..3b40e0d3 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,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)