From 1b66626bde35e036003975dd9f981d647e49c053 Mon Sep 17 00:00:00 2001 From: Dana Wensberg Date: Mon, 6 Apr 2026 23:38:07 -0400 Subject: [PATCH] fix(SUBCON-444): deserialize usage stats into proper dataclass instances _parse_run() was passing raw API dicts directly into Usage.models and Usage.platform_tools instead of converting them to ModelUsage and PlatformToolUsage instances. This meant attribute access (e.g. run.usage.models[0].input_tokens) failed at runtime. Also maps camelCase API keys (inputTokens, outputTokens, etc.) to snake_case and adds the missing duration_ms field to Usage. Co-Authored-By: Claude Opus 4.6 (1M context) --- subconscious/client.py | 24 +++++++- subconscious/types.py | 1 + tests/test_types.py | 130 ++++++++++++++++++++++++++++++++++++++++- 3 files changed, 152 insertions(+), 3 deletions(-) diff --git a/subconscious/client.py b/subconscious/client.py index 01f547b..ef6d74a 100644 --- a/subconscious/client.py +++ b/subconscious/client.py @@ -14,6 +14,8 @@ DoneEvent, Engine, ErrorEvent, + ModelUsage, + PlatformToolUsage, PollOptions, Run, RunInput, @@ -483,9 +485,27 @@ def _parse_run(self, data: Dict[str, Any]) -> Run: usage = None if "usage" in data and data["usage"]: + raw_usage = data["usage"] + models = [ + ModelUsage( + engine=m.get("engine", ""), + input_tokens=m.get("inputTokens", m.get("input_tokens", 0)), + output_tokens=m.get("outputTokens", m.get("output_tokens", 0)), + total_tokens=m.get("totalTokens", m.get("total_tokens", 0)), + ) + for m in raw_usage.get("models", []) + ] + platform_tools = [ + PlatformToolUsage( + tool_id=pt.get("toolId", pt.get("tool_id", "")), + calls=pt.get("calls", 0), + ) + for pt in raw_usage.get("platformTools", []) + ] usage = Usage( - models=data["usage"].get("models", []), - platform_tools=data["usage"].get("platformTools", []), + models=models, + platform_tools=platform_tools, + duration_ms=raw_usage.get("durationMs", raw_usage.get("duration_ms")), ) return Run( diff --git a/subconscious/types.py b/subconscious/types.py index 99ee906..b45cb15 100644 --- a/subconscious/types.py +++ b/subconscious/types.py @@ -102,6 +102,7 @@ class Usage: models: List[ModelUsage] = field(default_factory=list) platform_tools: List[PlatformToolUsage] = field(default_factory=list) + duration_ms: Optional[int] = None @dataclass diff --git a/tests/test_types.py b/tests/test_types.py index 85b9e00..7334b1e 100644 --- a/tests/test_types.py +++ b/tests/test_types.py @@ -6,10 +6,13 @@ FunctionTool, McpAuth, MCPTool, + ModelUsage, PlatformTool, + PlatformToolUsage, Tool, + Usage, ) -from subconscious.client import _normalize_tool +from subconscious.client import _normalize_tool, Subconscious # --------------------------------------------------------------------------- @@ -142,3 +145,128 @@ def test_platform_tool_unchanged(self): result = _normalize_tool(tool) assert result["id"] == "fast_search" assert result["options"] == {"limit": 10} + + +# --------------------------------------------------------------------------- +# Usage parsing (_parse_run) +# --------------------------------------------------------------------------- + +class TestParseRunUsage: + """Test that _parse_run correctly deserializes usage statistics.""" + + def _parse(self, data): + """Call _parse_run without needing a live client.""" + return Subconscious._parse_run(None, data) + + def test_usage_with_camel_case_api_response(self): + data = { + "runId": "run_abc", + "status": "succeeded", + "result": {"answer": "hello"}, + "usage": { + "models": [ + { + "engine": "tim-gpt", + "inputTokens": 150, + "outputTokens": 42, + "totalTokens": 192, + } + ], + "platformTools": [ + {"toolId": "fast_search", "calls": 3} + ], + "durationMs": 1234, + }, + } + run = self._parse(data) + assert run.run_id == "run_abc" + assert run.status == "succeeded" + assert run.result.answer == "hello" + + # Usage should be proper dataclass instances + assert isinstance(run.usage, Usage) + assert len(run.usage.models) == 1 + m = run.usage.models[0] + assert isinstance(m, ModelUsage) + assert m.engine == "tim-gpt" + assert m.input_tokens == 150 + assert m.output_tokens == 42 + assert m.total_tokens == 192 + + assert len(run.usage.platform_tools) == 1 + pt = run.usage.platform_tools[0] + assert isinstance(pt, PlatformToolUsage) + assert pt.tool_id == "fast_search" + assert pt.calls == 3 + + assert run.usage.duration_ms == 1234 + + def test_usage_with_snake_case_keys(self): + """Ensure snake_case keys also work (defensive).""" + data = { + "runId": "run_def", + "status": "succeeded", + "result": {"answer": "ok"}, + "usage": { + "models": [ + { + "engine": "tim-edge", + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + } + ], + "platformTools": [ + {"tool_id": "web_browse", "calls": 1} + ], + "duration_ms": 500, + }, + } + run = self._parse(data) + m = run.usage.models[0] + assert m.input_tokens == 10 + assert m.output_tokens == 5 + assert m.total_tokens == 15 + + pt = run.usage.platform_tools[0] + assert pt.tool_id == "web_browse" + assert run.usage.duration_ms == 500 + + def test_usage_with_multiple_models(self): + data = { + "runId": "run_multi", + "status": "succeeded", + "usage": { + "models": [ + {"engine": "tim-edge", "inputTokens": 10, "outputTokens": 5, "totalTokens": 15}, + {"engine": "tim-gpt", "inputTokens": 100, "outputTokens": 50, "totalTokens": 150}, + ], + "platformTools": [], + }, + } + run = self._parse(data) + assert len(run.usage.models) == 2 + assert run.usage.models[0].engine == "tim-edge" + assert run.usage.models[1].engine == "tim-gpt" + assert run.usage.models[1].input_tokens == 100 + + def test_no_usage_returns_none(self): + data = {"runId": "run_none", "status": "queued"} + run = self._parse(data) + assert run.usage is None + + def test_empty_usage_returns_none(self): + data = {"runId": "run_empty", "status": "queued", "usage": {}} + run = self._parse(data) + assert run.usage is None + + def test_usage_without_duration(self): + data = { + "runId": "run_nodur", + "status": "succeeded", + "usage": { + "models": [{"engine": "tim-gpt", "inputTokens": 1, "outputTokens": 1, "totalTokens": 2}], + }, + } + run = self._parse(data) + assert run.usage.duration_ms is None