fix(tracing): fail open temporal span activities

Author: danielmillerp

src/agentex/lib/adk/_modules/tracing.py

@@ -6,7 +6,9 @@
 from datetime import timedelta
 from typing import Any
 
+from temporalio import workflow
 from temporalio.common import RetryPolicy
+from temporalio.exceptions import ActivityError, TimeoutError as TemporalTimeoutError, is_cancelled_exception
 
 from agentex import AsyncAgentex  # noqa: F401
 from agentex.lib.adk.utils._modules.client import create_async_agentex_client
@@ -180,14 +182,25 @@ async def start_span(
             task_id=task_id,
         )
         if in_temporal_workflow():
-            return await ActivityHelpers.execute_activity(
-                activity_name=TracingActivityName.START_SPAN,
-                request=params,
-                response_type=Span,
-                start_to_close_timeout=start_to_close_timeout,
-                retry_policy=retry_policy,
-                heartbeat_timeout=heartbeat_timeout,
-            )
+            try:
+                return await ActivityHelpers.execute_activity(
+                    activity_name=TracingActivityName.START_SPAN,
+                    request=params,
+                    response_type=Span,
+                    start_to_close_timeout=start_to_close_timeout,
+                    retry_policy=retry_policy,
+                    heartbeat_timeout=heartbeat_timeout,
+                )
+            except (ActivityError, TemporalTimeoutError) as err:
+                if is_cancelled_exception(err):
+                    raise
+                workflow.logger.warning(
+                    "Failed to start tracing span %r for trace_id=%r; continuing without tracing",
+                    name,
+                    trace_id,
+                    exc_info=True,
+                )
+                return None
         else:
             return await self._tracing_service.start_span(
                 trace_id=trace_id,
@@ -224,14 +237,25 @@ async def end_span(
             span=span,
         )
         if in_temporal_workflow():
-            return await ActivityHelpers.execute_activity(
-                activity_name=TracingActivityName.END_SPAN,
-                request=params,
-                response_type=Span,
-                start_to_close_timeout=start_to_close_timeout,
-                retry_policy=retry_policy,
-                heartbeat_timeout=heartbeat_timeout,
-            )
+            try:
+                return await ActivityHelpers.execute_activity(
+                    activity_name=TracingActivityName.END_SPAN,
+                    request=params,
+                    response_type=Span,
+                    start_to_close_timeout=start_to_close_timeout,
+                    retry_policy=retry_policy,
+                    heartbeat_timeout=heartbeat_timeout,
+                )
+            except (ActivityError, TemporalTimeoutError) as err:
+                if is_cancelled_exception(err):
+                    raise
+                workflow.logger.warning(
+                    "Failed to end tracing span %r for trace_id=%r; continuing without closing trace",
+                    span.id,
+                    trace_id,
+                    exc_info=True,
+                )
+                return span
         else:
             return await self._tracing_service.end_span(
                 trace_id=trace_id,

tests/lib/adk/test_tracing_module.py

@@ -3,6 +3,9 @@
 from datetime import datetime, timezone
 from unittest.mock import AsyncMock, patch
 
+import pytest
+from temporalio.exceptions import ActivityError
+
 import agentex.lib.adk._modules.tracing as _tracing_mod
 from agentex.types.span import Span
 from agentex.lib.adk._modules.tracing import TracingModule
@@ -26,6 +29,18 @@ def _make_module() -> tuple[AsyncMock, TracingModule]:
     return mock_service, module
 
 
+def _make_activity_error() -> ActivityError:
+    return ActivityError(
+        "activity timed out",
+        scheduled_event_id=1,
+        started_event_id=2,
+        identity="worker-1",
+        activity_type="start-span",
+        activity_id="activity-1",
+        retry_state=None,
+    )
+
+
 class TestStartSpan:
     async def test_start_span_with_task_id(self):
         mock_service, module = _make_module()
@@ -87,6 +102,102 @@ async def test_end_span_preserves_task_id(self):
         mock_service.end_span.assert_called_once_with(trace_id="trace-123", span=span)
 
 
+class TestTracingModuleTemporalPath:
+    async def test_start_span_in_workflow_returns_none_when_activity_fails(self):
+        mock_service, module = _make_module()
+
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
+                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
+                patch.object(_tracing_mod.workflow, "logger") as mock_logger:
+            mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error())
+            result = await module.start_span(trace_id="trace-123", name="test-span")
+
+        assert result is None
+        mock_logger.warning.assert_called_once()
+        mock_helpers.execute_activity.assert_called_once()
+        mock_service.start_span.assert_not_called()
+
+    async def test_end_span_in_workflow_returns_span_when_activity_fails(self):
+        mock_service, module = _make_module()
+        span = _make_span()
+
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
+                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
+                patch.object(_tracing_mod.workflow, "logger") as mock_logger:
+            mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error())
+            result = await module.end_span(trace_id="trace-123", span=span)
+
+        assert result == span
+        mock_logger.warning.assert_called_once()
+        mock_helpers.execute_activity.assert_called_once()
+        mock_service.end_span.assert_not_called()
+
+    async def test_context_manager_skips_end_when_temporal_start_fails(self):
+        mock_service, module = _make_module()
+
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
+                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
+                patch.object(_tracing_mod.workflow, "logger"):
+            mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error())
+            async with module.span(trace_id="trace-123", name="test-span") as span:
+                assert span is None
+
+        mock_helpers.execute_activity.assert_called_once()
+        mock_service.start_span.assert_not_called()
+        mock_service.end_span.assert_not_called()
+
+    async def test_start_span_in_workflow_propagates_unexpected_errors(self):
+        mock_service, module = _make_module()
+
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
+                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers:
+            mock_helpers.execute_activity = AsyncMock(side_effect=RuntimeError("bad response shape"))
+            try:
+                await module.start_span(trace_id="trace-123", name="test-span")
+            except RuntimeError as exc:
+                assert str(exc) == "bad response shape"
+            else:
+                raise AssertionError("Expected unexpected errors to propagate")
+
+        mock_helpers.execute_activity.assert_called_once()
+        mock_service.start_span.assert_not_called()
+
+    async def test_start_span_in_workflow_propagates_cancellation(self):
+        mock_service, module = _make_module()
+        activity_error = _make_activity_error()
+
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
+                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
+                patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), \
+                patch.object(_tracing_mod.workflow, "logger") as mock_logger:
+            mock_helpers.execute_activity = AsyncMock(side_effect=activity_error)
+
+            with pytest.raises(ActivityError):
+                await module.start_span(trace_id="trace-123", name="test-span")
+
+        mock_logger.warning.assert_not_called()
+        mock_helpers.execute_activity.assert_called_once()
+        mock_service.start_span.assert_not_called()
+
+    async def test_end_span_in_workflow_propagates_cancellation(self):
+        mock_service, module = _make_module()
+        span = _make_span()
+        activity_error = _make_activity_error()
+
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
+                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
+                patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), \
+                patch.object(_tracing_mod.workflow, "logger") as mock_logger:
+            mock_helpers.execute_activity = AsyncMock(side_effect=activity_error)
+
+            with pytest.raises(ActivityError):
+                await module.end_span(trace_id="trace-123", span=span)
+
+        mock_logger.warning.assert_not_called()
+        mock_helpers.execute_activity.assert_called_once()
+        mock_service.end_span.assert_not_called()
+
+
 class TestSpanContextManager:
     async def test_span_context_manager_forwards_task_id(self):
         mock_service, module = _make_module()