diff --git a/tests/test_tokens.py b/tests/test_tokens.py index 133cbd2..bf2cdc0 100644 --- a/tests/test_tokens.py +++ b/tests/test_tokens.py @@ -115,6 +115,34 @@ def test_missing_usage_defaults_to_zero(): assert c.prompt == 0 and c.completion == 0 and c.cache_read == 0 and c.cache_creation == 0 +def test_malformed_usage_values_default_to_zero(): + """Malformed usage values should not break trace emission.""" + msg = { + "type": "assistant", + "requestId": "r_bad_usage", + "timestamp": "2026-06-22T00:00:00Z", + "message": { + "id": "r_bad_usage", + "model": "claude-haiku", + "content": [{"type": "text", "text": "ok"}], + "usage": { + "input_tokens": "unknown", + "output_tokens": "NaN", + "cache_read_input_tokens": None, + "cache_creation_input_tokens": {"bad": "shape"}, + }, + }, + } + calls = tokens.llm_calls([msg]) + assert len(calls) == 1 + c = calls[0] + assert c.prompt == 0 + assert c.completion == 0 + assert c.cache_read == 0 + assert c.cache_creation == 0 + assert c.text == "ok" + + def test_start_and_end_timestamps(): """start_ts = first frame's timestamp, end_ts = last frame's timestamp.""" msgs = [ diff --git a/traceroot_observability/tokens.py b/traceroot_observability/tokens.py index a45136e..6233355 100644 --- a/traceroot_observability/tokens.py +++ b/traceroot_observability/tokens.py @@ -68,6 +68,14 @@ def _tool_calls(content) -> list: return out +def _usage_int(value) -> int: + """Return a token count as int, defaulting malformed values to zero.""" + try: + return int(value or 0) + except (TypeError, ValueError): + return 0 + + def llm_calls(assistant_msgs: list[dict]) -> list[LlmCall]: """ Collapse a flat list of assistant transcript messages into per-requestId LlmCall @@ -118,16 +126,16 @@ def llm_calls(assistant_msgs: list[dict]) -> list[LlmCall]: u = msg.get("usage") or {} if rid not in seen_input: # First sighting: count input and cache tokens exactly once. - inp = int(u.get("input_tokens") or 0) - cr = int(u.get("cache_read_input_tokens") or 0) - cc = int(u.get("cache_creation_input_tokens") or 0) + inp = _usage_int(u.get("input_tokens")) + cr = _usage_int(u.get("cache_read_input_tokens")) + cc = _usage_int(u.get("cache_creation_input_tokens")) call.prompt += inp + cr + cc call.cache_read += cr call.cache_creation += cc seen_input.add(rid) # Output: add only the delta over the running max for this requestId. - out = int(u.get("output_tokens") or 0) + out = _usage_int(u.get("output_tokens")) prior = out_max.get(rid, 0) if out > prior: call.completion += out - prior