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@fa5d4a9

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
+40-16Lines changed: 40 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
@@ -180,14 +182,25 @@ async def start_span(
180182
task_id=task_id,
181183
)
182184
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-
)
185+
try:
186+
return await ActivityHelpers.execute_activity(
187+
activity_name=TracingActivityName.START_SPAN,
188+
request=params,
189+
response_type=Span,
190+
start_to_close_timeout=start_to_close_timeout,
191+
retry_policy=retry_policy,
192+
heartbeat_timeout=heartbeat_timeout,
193+
)
194+
except (ActivityError, TemporalTimeoutError) as err:
195+
if is_cancelled_exception(err):
196+
raise
197+
workflow.logger.warning(
198+
"Failed to start tracing span %r for trace_id=%r; continuing without tracing",
199+
name,
200+
trace_id,
201+
exc_info=True,
202+
)
203+
return None
191204
else:
192205
return await self._tracing_service.start_span(
193206
trace_id=trace_id,
@@ -224,14 +237,25 @@ async def end_span(
224237
span=span,
225238
)
226239
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-
)
240+
try:
241+
return await ActivityHelpers.execute_activity(
242+
activity_name=TracingActivityName.END_SPAN,
243+
request=params,
244+
response_type=Span,
245+
start_to_close_timeout=start_to_close_timeout,
246+
retry_policy=retry_policy,
247+
heartbeat_timeout=heartbeat_timeout,
248+
)
249+
except (ActivityError, TemporalTimeoutError) as err:
250+
if is_cancelled_exception(err):
251+
raise
252+
workflow.logger.warning(
253+
"Failed to end tracing span %r for trace_id=%r; continuing without closing trace",
254+
span.id,
255+
trace_id,
256+
exc_info=True,
257+
)
258+
return span
235259
else:
236260
return await self._tracing_service.end_span(
237261
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
+111Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,9 @@
33
from datetime import datetime, timezone
44
from unittest.mock import AsyncMock, patch
55

6+
import pytest
7+
from temporalio.exceptions import ActivityError
8+
69
import agentex.lib.adk._modules.tracing as _tracing_mod
710
from agentex.types.span import Span
811
from agentex.lib.adk._modules.tracing import TracingModule
@@ -26,6 +29,18 @@ 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+
2944
class TestStartSpan:
3045
async def test_start_span_with_task_id(self):
3146
mock_service, module = _make_module()
@@ -87,6 +102,102 @@ async def test_end_span_preserves_task_id(self):
87102
mock_service.end_span.assert_called_once_with(trace_id="trace-123", span=span)
88103

89104

105+
class TestTracingModuleTemporalPath:
106+
async def test_start_span_in_workflow_returns_none_when_activity_fails(self):
107+
mock_service, module = _make_module()
108+
109+
with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
110+
patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
111+
patch.object(_tracing_mod.workflow, "logger") as mock_logger:
112+
mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error())
113+
result = await module.start_span(trace_id="trace-123", name="test-span")
114+
115+
assert result is None
116+
mock_logger.warning.assert_called_once()
117+
mock_helpers.execute_activity.assert_called_once()
118+
mock_service.start_span.assert_not_called()
119+
120+
async def test_end_span_in_workflow_returns_span_when_activity_fails(self):
121+
mock_service, module = _make_module()
122+
span = _make_span()
123+
124+
with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
125+
patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
126+
patch.object(_tracing_mod.workflow, "logger") as mock_logger:
127+
mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error())
128+
result = await module.end_span(trace_id="trace-123", span=span)
129+
130+
assert result == span
131+
mock_logger.warning.assert_called_once()
132+
mock_helpers.execute_activity.assert_called_once()
133+
mock_service.end_span.assert_not_called()
134+
135+
async def test_context_manager_skips_end_when_temporal_start_fails(self):
136+
mock_service, module = _make_module()
137+
138+
with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
139+
patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
140+
patch.object(_tracing_mod.workflow, "logger"):
141+
mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error())
142+
async with module.span(trace_id="trace-123", name="test-span") as span:
143+
assert span is None
144+
145+
mock_helpers.execute_activity.assert_called_once()
146+
mock_service.start_span.assert_not_called()
147+
mock_service.end_span.assert_not_called()
148+
149+
async def test_start_span_in_workflow_propagates_unexpected_errors(self):
150+
mock_service, module = _make_module()
151+
152+
with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
153+
patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers:
154+
mock_helpers.execute_activity = AsyncMock(side_effect=RuntimeError("bad response shape"))
155+
try:
156+
await module.start_span(trace_id="trace-123", name="test-span")
157+
except RuntimeError as exc:
158+
assert str(exc) == "bad response shape"
159+
else:
160+
raise AssertionError("Expected unexpected errors to propagate")
161+
162+
mock_helpers.execute_activity.assert_called_once()
163+
mock_service.start_span.assert_not_called()
164+
165+
async def test_start_span_in_workflow_propagates_cancellation(self):
166+
mock_service, module = _make_module()
167+
activity_error = _make_activity_error()
168+
169+
with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
170+
patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
171+
patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), \
172+
patch.object(_tracing_mod.workflow, "logger") as mock_logger:
173+
mock_helpers.execute_activity = AsyncMock(side_effect=activity_error)
174+
175+
with pytest.raises(ActivityError):
176+
await module.start_span(trace_id="trace-123", name="test-span")
177+
178+
mock_logger.warning.assert_not_called()
179+
mock_helpers.execute_activity.assert_called_once()
180+
mock_service.start_span.assert_not_called()
181+
182+
async def test_end_span_in_workflow_propagates_cancellation(self):
183+
mock_service, module = _make_module()
184+
span = _make_span()
185+
activity_error = _make_activity_error()
186+
187+
with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
188+
patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
189+
patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), \
190+
patch.object(_tracing_mod.workflow, "logger") as mock_logger:
191+
mock_helpers.execute_activity = AsyncMock(side_effect=activity_error)
192+
193+
with pytest.raises(ActivityError):
194+
await module.end_span(trace_id="trace-123", span=span)
195+
196+
mock_logger.warning.assert_not_called()
197+
mock_helpers.execute_activity.assert_called_once()
198+
mock_service.end_span.assert_not_called()
199+
200+
90201
class TestSpanContextManager:
91202
async def test_span_context_manager_forwards_task_id(self):
92203
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.