core/agent: normalize provider usage structs in token_usage_from_api

src/plyngent/agent/messages_bridge.py

@@ -64,16 +64,6 @@ def _encode_tool_arguments(obj: dict[str, Any]) -> str:
     return json.dumps(obj, ensure_ascii=False, separators=(",", ":"))
 
 
-def _as_nonneg_int(value: object) -> int:
-    if isinstance(value, bool):
-        return 0
-    if isinstance(value, int):
-        return max(0, value)
-    if isinstance(value, float):
-        return max(0, int(value))
-    return 0
-
-
 def tool_items_to_anthropic_tools(
     tools: Sequence[AnyToolItem] | None,
 ) -> list[AnthropicToolDefinition]:
@@ -235,23 +225,17 @@ def anthropic_stop_to_finish_reason(
 
 def anthropic_usage_to_dict(usage: AnthropicUsage | dict[str, Any] | None) -> dict[str, Any] | None:
     """Normalize Anthropic usage Struct/dict to chat-compatible usage fields."""
-    if usage is None:
-        return None
-    if isinstance(usage, AnthropicUsage):
-        prompt = _as_nonneg_int(usage.input_tokens)
-        completion = _as_nonneg_int(usage.output_tokens)
-    else:
-        raw = cast("dict[str, object]", usage)
-        prompt = _as_nonneg_int(raw.get("input_tokens") or raw.get("prompt_tokens") or 0)
-        completion = _as_nonneg_int(raw.get("output_tokens") or raw.get("completion_tokens") or 0)
-    if prompt == 0 and completion == 0:
+    from .usage import token_usage_from_api
+
+    parsed = token_usage_from_api(usage)
+    if parsed is None:
         return None
     return {
-        "prompt_tokens": prompt,
-        "completion_tokens": completion,
-        "total_tokens": prompt + completion,
-        "input_tokens": prompt,
-        "output_tokens": completion,
+        "prompt_tokens": parsed.prompt_tokens,
+        "completion_tokens": parsed.completion_tokens,
+        "total_tokens": parsed.total_tokens,
+        "input_tokens": parsed.prompt_tokens,
+        "output_tokens": parsed.completion_tokens,
     }
 
 

src/plyngent/agent/usage.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, cast
 
+import msgspec
 from msgspec import UNSET, Struct
 
 from plyngent.agent.budget import estimate_message_chars, estimate_messages_chars
@@ -86,13 +87,16 @@ def _as_nonneg_int(value: object) -> int:
 
 
 def token_usage_from_api(usage: object) -> TokenUsage | None:
-    """Parse OpenAI-style usage dict; return None if missing/empty.
+    """Parse API usage into :class:`TokenUsage`; return None if missing/empty.
 
-    Accepts chat completions fields (``prompt_tokens`` / ``completion_tokens``)
-    and Responses fields (``input_tokens`` / ``output_tokens``).
+    Accepts chat completions fields (``prompt_tokens`` / ``completion_tokens``),
+    Responses fields (``input_tokens`` / ``output_tokens``), and msgspec Structs
+    carrying either pair (e.g. ``AnthropicUsage``).
     """
     if usage is None or usage is UNSET:
         return None
+    if isinstance(usage, Struct):
+        usage = msgspec.to_builtins(usage)
     if not isinstance(usage, dict):
         return None
     raw = cast("dict[str, object]", usage)

tests/test_agent/test_usage.py

@@ -58,6 +58,23 @@ def test_token_usage_from_api_responses_fields() -> None:
     assert u.total_tokens == 9
 
 
+def test_token_usage_from_api_struct() -> None:
+    from plyngent.lmproto.anthropic.model import AnthropicUsage
+
+    u = token_usage_from_api(AnthropicUsage(input_tokens=7, output_tokens=2))
+    assert u is not None
+    assert u.prompt_tokens == 7
+    assert u.completion_tokens == 2
+    assert u.total_tokens == 9
+    assert u.source == "api"
+
+
+def test_token_usage_from_api_empty_struct_returns_none() -> None:
+    from plyngent.lmproto.anthropic.model import AnthropicUsage
+
+    assert token_usage_from_api(AnthropicUsage()) is None
+
+
 def test_chars_to_tokens() -> None:
     assert chars_to_tokens(0) == 0
     assert chars_to_tokens(1) == 1