Skip to content

Commit 25d8560

Browse files
committed
fix(adk): warn when record_usage replaces non-dict span data
1 parent 233e3b9 commit 25d8560

2 files changed

Lines changed: 33 additions & 0 deletions

File tree

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,11 @@ def record_usage(
9393
)
9494
)
9595

96+
if self.span.data is not None and not isinstance(self.span.data, dict):
97+
logger.warning(
98+
f"TurnSpan.record_usage: span.data is {type(self.span.data).__name__} "
99+
"(expected dict or None); existing data will be replaced."
100+
)
96101
data = self.span.data if isinstance(self.span.data, dict) else {}
97102
if blob:
98103
data["usage"] = blob

tests/lib/adk/test_tracing_module.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -321,6 +321,34 @@ async def test_turn_span_preserves_existing_data(self):
321321
assert ended_span.data["custom"] == "value"
322322
assert ended_span.data["usage"] == {"prompt_tokens": 3, "completion_tokens": 4}
323323

324+
async def test_turn_span_warns_and_replaces_non_dict_data(self, caplog):
325+
mock_service, module = _make_module()
326+
started = _make_span(data=[{"item": 1}])
327+
mock_service.start_span.return_value = started
328+
mock_service.end_span.return_value = started
329+
330+
with patch.object(_tracing_mod, "in_temporal_workflow", return_value=False):
331+
async with module.turn_span(trace_id="trace-123", name="turn") as turn:
332+
with caplog.at_level("WARNING"):
333+
turn.record_usage(usage={"input_tokens": 1, "output_tokens": 2})
334+
335+
assert any("existing data will be replaced" in message for message in caplog.messages)
336+
ended_span = mock_service.end_span.call_args.kwargs["span"]
337+
assert ended_span.data == {"usage": {"input_tokens": 1, "output_tokens": 2}}
338+
339+
async def test_turn_span_dict_data_does_not_warn(self, caplog):
340+
mock_service, module = _make_module()
341+
started = _make_span(data={"custom": "value"})
342+
mock_service.start_span.return_value = started
343+
mock_service.end_span.return_value = started
344+
345+
with patch.object(_tracing_mod, "in_temporal_workflow", return_value=False):
346+
async with module.turn_span(trace_id="trace-123", name="turn") as turn:
347+
with caplog.at_level("WARNING"):
348+
turn.record_usage(usage={"input_tokens": 1})
349+
350+
assert not any("existing data will be replaced" in message for message in caplog.messages)
351+
324352
async def test_turn_span_cost_only(self):
325353
mock_service, module = _make_module()
326354
started = _make_span()

0 commit comments

Comments
 (0)