GV

GeezerVison

The World Wide Web, restored to 800 pixels of glory.

Home|About|Gaming cabinet|Archives|How it works|Starter sites

On this site

  • Sign in
  • ActionsCI/CD and automation
  • CodespacesCloud development environments
  • All Features
  • Documentation
  • Training
  • Certifications
  • Build with Copilot
  • Customize Copilot
  • Integrations
  • Trending repos
  • GitHub Sponsors
  • Developer articles
  • Enterprises

This view

  • Original page
  • New address

GeezerVison is showing github.com — fix(tracing): fail open temporal span activities · scaleapi/scale-agentex-python@e90a03e

File tree

Expand file treeCollapse file tree
Open diff view settings
Filter options
  • src/agentex/lib
    • adk/_modules
      • tracing.py
    • sdk/state_machine
      • state_machine.py
  • tests/lib
    • adk
      • test_tracing_module.py
    • test_state_machine.py
Expand file treeCollapse file tree
Open diff view settings
Collapse file

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

Copy file name to clipboardExpand all lines: src/agentex/lib/adk/_modules/tracing.py
+54-16Lines changed: 54 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,9 @@
66
from datetime import timedelta
77
from typing import Any
88

9+
from temporalio import workflow
910
from temporalio.common import RetryPolicy
11+
from temporalio.exceptions import ActivityError, TimeoutError as TemporalTimeoutError, is_cancelled_exception
1012

1113
from agentex import AsyncAgentex # noqa: F401
1214
from agentex.lib.adk.utils._modules.client import create_async_agentex_client
@@ -26,6 +28,18 @@
2628
logger = make_logger(__name__)
2729

2830
DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1)
31+
TEMPORAL_SPAN_ACTIVITY_DROPPED_METRIC = "agentex.tracing.temporal_span_activity.dropped"
32+
33+
34+
def _record_temporal_span_activity_dropped(event_type: str) -> None:
35+
try:
36+
workflow.metric_meter().create_counter(
37+
TEMPORAL_SPAN_ACTIVITY_DROPPED_METRIC,
38+
description="Temporal tracing span activities dropped after fail-open",
39+
unit="1",
40+
).add(1, {"event_type": event_type})
41+
except Exception:
42+
pass
2943

3044

3145
class TracingModule:
@@ -180,14 +194,26 @@ async def start_span(
180194
task_id=task_id,
181195
)
182196
if in_temporal_workflow():
183-
return await ActivityHelpers.execute_activity(
184-
activity_name=TracingActivityName.START_SPAN,
185-
request=params,
186-
response_type=Span,
187-
start_to_close_timeout=start_to_close_timeout,
188-
retry_policy=retry_policy,
189-
heartbeat_timeout=heartbeat_timeout,
190-
)
197+
try:
198+
return await ActivityHelpers.execute_activity(
199+
activity_name=TracingActivityName.START_SPAN,
200+
request=params,
201+
response_type=Span,
202+
start_to_close_timeout=start_to_close_timeout,
203+
retry_policy=retry_policy,
204+
heartbeat_timeout=heartbeat_timeout,
205+
)
206+
except (ActivityError, TemporalTimeoutError) as err:
207+
if is_cancelled_exception(err):
208+
raise
209+
workflow.logger.warning(
210+
"Failed to start tracing span %r for trace_id=%r; continuing without tracing",
211+
name,
212+
trace_id,
213+
exc_info=True,
214+
)
215+
_record_temporal_span_activity_dropped("start")
216+
return None
191217
else:
192218
return await self._tracing_service.start_span(
193219
trace_id=trace_id,
@@ -224,14 +250,26 @@ async def end_span(
224250
span=span,
225251
)
226252
if in_temporal_workflow():
227-
return await ActivityHelpers.execute_activity(
228-
activity_name=TracingActivityName.END_SPAN,
229-
request=params,
230-
response_type=Span,
231-
start_to_close_timeout=start_to_close_timeout,
232-
retry_policy=retry_policy,
233-
heartbeat_timeout=heartbeat_timeout,
234-
)
253+
try:
254+
return await ActivityHelpers.execute_activity(
255+
activity_name=TracingActivityName.END_SPAN,
256+
request=params,
257+
response_type=Span,
258+
start_to_close_timeout=start_to_close_timeout,
259+
retry_policy=retry_policy,
260+
heartbeat_timeout=heartbeat_timeout,
261+
)
262+
except (ActivityError, TemporalTimeoutError) as err:
263+
if is_cancelled_exception(err):
264+
raise
265+
workflow.logger.warning(
266+
"Failed to end tracing span %r for trace_id=%r; continuing without closing trace",
267+
span.id,
268+
trace_id,
269+
exc_info=True,
270+
)
271+
_record_temporal_span_activity_dropped("end")
272+
return span
235273
else:
236274
return await self._tracing_service.end_span(
237275
trace_id=trace_id,
Collapse file

‎src/agentex/lib/sdk/state_machine/state_machine.py‎

Copy file name to clipboardExpand all lines: src/agentex/lib/sdk/state_machine/state_machine.py
+2-1Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,7 @@ async def reset_to_initial_state(self):
113113
"""
114114
Reset the state machine to its initial state.
115115
"""
116+
span = None
116117
if self._trace_transitions:
117118
if self._task_id is None:
118119
raise ValueError(
@@ -126,7 +127,7 @@ async def reset_to_initial_state(self):
126127

127128
await self.transition(self._initial_state)
128129

129-
if self._trace_transitions:
130+
if self._trace_transitions and span is not None:
130131
span.output = {"output_state": self._initial_state} # type: ignore[assignment,union-attr]
131132
await adk.tracing.end_span(trace_id=self._task_id, span=span)
132133

Collapse file

‎tests/lib/adk/test_tracing_module.py‎

Copy file name to clipboardExpand all lines: tests/lib/adk/test_tracing_module.py
+144-1Lines changed: 144 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
11
from __future__ import annotations
22

33
from datetime import datetime, timezone
4-
from unittest.mock import AsyncMock, patch
4+
from unittest.mock import AsyncMock, MagicMock, patch
5+
6+
import pytest
7+
from temporalio.exceptions import ActivityError
58

69
import agentex.lib.adk._modules.tracing as _tracing_mod
710
from agentex.types.span import Span
@@ -26,6 +29,24 @@ def _make_module() -> tuple[AsyncMock, TracingModule]:
2629
return mock_service, module
2730

2831

32+
def _make_activity_error() -> ActivityError:
33+
return ActivityError(
34+
"activity timed out",
35+
scheduled_event_id=1,
36+
started_event_id=2,
37+
identity="worker-1",
38+
activity_type="start-span",
39+
activity_id="activity-1",
40+
retry_state=None,
41+
)
42+
43+
44+
def _make_metric_meter() -> MagicMock:
45+
mock_meter = MagicMock()
46+
mock_meter.create_counter.return_value = MagicMock()
47+
return mock_meter
48+
49+
2950
class TestStartSpan:
3051
async def test_start_span_with_task_id(self):
3152
mock_service, module = _make_module()
@@ -87,6 +108,128 @@ async def test_end_span_preserves_task_id(self):
87108
mock_service.end_span.assert_called_once_with(trace_id="trace-123", span=span)
88109

89110

111+
class TestTracingModuleTemporalPath:
112+
async def test_start_span_in_workflow_returns_none_when_activity_fails(self):
113+
mock_service, module = _make_module()
114+
mock_meter = _make_metric_meter()
115+
116+
with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
117+
patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
118+
patch.object(_tracing_mod.workflow, "logger") as mock_logger, \
119+
patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter):
120+
mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error())
121+
result = await module.start_span(trace_id="trace-123", name="test-span")
122+
123+
assert result is None
124+
mock_logger.warning.assert_called_once()
125+
mock_meter.create_counter.assert_called_once_with(
126+
_tracing_mod.TEMPORAL_SPAN_ACTIVITY_DROPPED_METRIC,
127+
description="Temporal tracing span activities dropped after fail-open",
128+
unit="1",
129+
)
130+
mock_meter.create_counter.return_value.add.assert_called_once_with(
131+
1, {"event_type": "start"}
132+
)
133+
mock_helpers.execute_activity.assert_called_once()
134+
mock_service.start_span.assert_not_called()
135+
136+
async def test_end_span_in_workflow_returns_span_when_activity_fails(self):
137+
mock_service, module = _make_module()
138+
span = _make_span()
139+
mock_meter = _make_metric_meter()
140+
141+
with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
142+
patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
143+
patch.object(_tracing_mod.workflow, "logger") as mock_logger, \
144+
patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter):
145+
mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error())
146+
result = await module.end_span(trace_id="trace-123", span=span)
147+
148+
assert result == span
149+
mock_logger.warning.assert_called_once()
150+
mock_meter.create_counter.assert_called_once_with(
151+
_tracing_mod.TEMPORAL_SPAN_ACTIVITY_DROPPED_METRIC,
152+
description="Temporal tracing span activities dropped after fail-open",
153+
unit="1",
154+
)
155+
mock_meter.create_counter.return_value.add.assert_called_once_with(
156+
1, {"event_type": "end"}
157+
)
158+
mock_helpers.execute_activity.assert_called_once()
159+
mock_service.end_span.assert_not_called()
160+
161+
async def test_context_manager_skips_end_when_temporal_start_fails(self):
162+
mock_service, module = _make_module()
163+
164+
with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
165+
patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
166+
patch.object(_tracing_mod.workflow, "logger"):
167+
mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error())
168+
async with module.span(trace_id="trace-123", name="test-span") as span:
169+
assert span is None
170+
171+
mock_helpers.execute_activity.assert_called_once()
172+
mock_service.start_span.assert_not_called()
173+
mock_service.end_span.assert_not_called()
174+
175+
async def test_start_span_in_workflow_propagates_unexpected_errors(self):
176+
mock_service, module = _make_module()
177+
178+
with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
179+
patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers:
180+
mock_helpers.execute_activity = AsyncMock(side_effect=RuntimeError("bad response shape"))
181+
try:
182+
await module.start_span(trace_id="trace-123", name="test-span")
183+
except RuntimeError as exc:
184+
assert str(exc) == "bad response shape"
185+
else:
186+
raise AssertionError("Expected unexpected errors to propagate")
187+
188+
mock_helpers.execute_activity.assert_called_once()
189+
mock_service.start_span.assert_not_called()
190+
191+
async def test_start_span_in_workflow_propagates_cancellation(self):
192+
mock_service, module = _make_module()
193+
activity_error = _make_activity_error()
194+
mock_meter = _make_metric_meter()
195+
196+
with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
197+
patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
198+
patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), \
199+
patch.object(_tracing_mod.workflow, "logger") as mock_logger, \
200+
patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter):
201+
mock_helpers.execute_activity = AsyncMock(side_effect=activity_error)
202+
203+
with pytest.raises(ActivityError):
204+
await module.start_span(trace_id="trace-123", name="test-span")
205+
206+
mock_logger.warning.assert_not_called()
207+
mock_meter.create_counter.assert_not_called()
208+
mock_helpers.execute_activity.assert_called_once()
209+
mock_service.start_span.assert_not_called()
210+
211+
async def test_end_span_in_workflow_propagates_cancellation(self):
212+
mock_service, module = _make_module()
213+
span = _make_span()
214+
activity_error = _make_activity_error()
215+
mock_meter = _make_metric_meter()
216+
217+
with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
218+
patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
219+
patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), \
220+
patch.object(_tracing_mod.workflow, "logger") as mock_logger, \
221+
patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter):
222+
mock_helpers.execute_activity = AsyncMock(side_effect=activity_error)
223+
224+
with pytest.raises(ActivityError):
225+
await module.end_span(trace_id="trace-123", span=span)
226+
227+
mock_logger.warning.assert_not_called()
228+
mock_meter.create_counter.assert_not_called()
229+
mock_helpers.execute_activity.assert_called_once()
230+
mock_service.end_span.assert_not_called()
231+
232+
90233
class TestSpanContextManager:
91234
async def test_span_context_manager_forwards_task_id(self):
92235
mock_service, module = _make_module()
Collapse file

‎tests/lib/test_state_machine.py‎

Copy file name to clipboard
+68Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
from __future__ import annotations
2+
3+
from typing import override
4+
from unittest.mock import AsyncMock, patch
5+
6+
from agentex.lib.sdk.state_machine import State, StateMachine, StateWorkflow
7+
from agentex.lib.utils.model_utils import BaseModel
8+
9+
10+
class ExampleData(BaseModel):
11+
value: int = 0
12+
13+
14+
class InitialWorkflow(StateWorkflow):
15+
transitions = ["next"]
16+
17+
@override
18+
async def execute(self, state_machine, state_machine_data=None):
19+
return "next"
20+
21+
22+
class NextWorkflow(StateWorkflow):
23+
transitions = ["initial"]
24+
25+
@override
26+
async def execute(self, state_machine, state_machine_data=None):
27+
return "initial"
28+
29+
30+
class ExampleStateMachine(StateMachine[ExampleData]):
31+
@override
32+
async def terminal_condition(self):
33+
return False
34+
35+
36+
def _make_state_machine() -> ExampleStateMachine:
37+
return ExampleStateMachine(
38+
initial_state="initial",
39+
states=[
40+
State(name="initial", workflow=InitialWorkflow()),
41+
State(name="next", workflow=NextWorkflow()),
42+
],
43+
task_id="task-123",
44+
state_machine_data=ExampleData(value=1),
45+
trace_transitions=True,
46+
)
47+
48+
49+
async def test_reset_to_initial_state_skips_end_span_when_start_span_fails_open():
50+
state_machine = _make_state_machine()
51+
await state_machine.transition("next")
52+
53+
with patch(
54+
"agentex.lib.sdk.state_machine.state_machine.adk.tracing.start_span",
55+
new=AsyncMock(return_value=None),
56+
) as start_span, patch(
57+
"agentex.lib.sdk.state_machine.state_machine.adk.tracing.end_span",
58+
new=AsyncMock(),
59+
) as end_span:
60+
await state_machine.reset_to_initial_state()
61+
62+
assert state_machine.get_current_state() == "initial"
63+
start_span.assert_awaited_once_with(
64+
trace_id="task-123",
65+
name="state_transition_reset",
66+
input={"input_state": "next"},
67+
)
68+
end_span.assert_not_awaited()

0 commit comments

Comments
0 (0)
You are visitor 01402520 — GeezerVison does not store the pages you look up.
Best viewed with Netscape Navigator 4.7 or Internet Explorer 6 at 800×600. No cookies required.