From 87fcf3c0aabface102ba9dc64768e000e07005b7 Mon Sep 17 00:00:00 2001 From: Evo Date: Sun, 28 Jun 2026 16:54:16 +0800 Subject: [PATCH 1/2] fix(llm-trace): stash litellm tool-call usage so token cost survives arg-parse failures (completes #2396) --- .../hindsight_api/engine/providers/litellm_llm.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/hindsight-api-slim/hindsight_api/engine/providers/litellm_llm.py b/hindsight-api-slim/hindsight_api/engine/providers/litellm_llm.py index 7ae414868f..1be253780f 100644 --- a/hindsight-api-slim/hindsight_api/engine/providers/litellm_llm.py +++ b/hindsight-api-slim/hindsight_api/engine/providers/litellm_llm.py @@ -408,6 +408,14 @@ async def call_with_tools( self._acompletion(**call_kwargs), timeout=self.timeout, ) + # Stash usage before the tool-call argument parse below, which + # can raise json.JSONDecodeError locally even though the provider + # already billed for these tokens; without this the error trace + # records 0/0 tokens (#2387). Mirrors call() and the anthropic/ + # gemini call_with_tools paths so the litellm tool path (and the + # LiteLLMRouterLLM subclass that inherits this method) completes + # the #2396 usage-on-error coverage. + stash_response_usage(_usage_from_litellm_response(response)) message = response.choices[0].message content = message.content From cf5bebb6748fe6c9a57a266291ac44b5246d437b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Tue, 30 Jun 2026 10:18:04 +0200 Subject: [PATCH 2/2] test(llm-trace): cover litellm tool-call arg-parse usage stash Add a real-provider regression test for the fix in this PR: the existing wrapper-level tools test uses a provider that already stashes, so it does not guard LiteLLMLLM.call_with_tools. This drives the real provider with a billed response whose tool arguments are malformed JSON and asserts the error trace keeps the provider-reported tokens (input/output/cached). The LiteLLMRouterLLM subclass inherits call_with_tools, so it is covered too. Verified it fails (input_tokens=None) when the stash line is removed. --- hindsight-api-slim/tests/test_llm_trace.py | 47 ++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/hindsight-api-slim/tests/test_llm_trace.py b/hindsight-api-slim/tests/test_llm_trace.py index 78bccd0192..3fc773901d 100644 --- a/hindsight-api-slim/tests/test_llm_trace.py +++ b/hindsight-api-slim/tests/test_llm_trace.py @@ -389,6 +389,53 @@ async def test_retain_extract_success_records_usage_once(registered_recorder): assert r.cached_tokens == 20 +# ── real provider: litellm tool-call arg-parse failure keeps usage (#2387) ──── + + +def _litellm_tool_response_with_usage(arguments: str): + """A successful LiteLLM (OpenAI-shaped) tool-call response carrying usage, + like ``call_with_tools`` sees right before it ``json.loads`` the tool + arguments.""" + function = SimpleNamespace(name="extract", arguments=arguments) + tool_call = SimpleNamespace(id="call_1", function=function) + message = SimpleNamespace(content=None, tool_calls=[tool_call]) + choice = SimpleNamespace(finish_reason="tool_calls", message=message) + usage = SimpleNamespace( + prompt_tokens=140, + completion_tokens=18, + total_tokens=158, + prompt_tokens_details=SimpleNamespace(cached_tokens=20), + ) + return SimpleNamespace(error=None, usage=usage, choices=[choice]) + + +@pytest.mark.asyncio +async def test_litellm_tool_call_arg_parse_failure_keeps_usage(registered_recorder): + """The litellm tool path bills the provider response, then ``json.loads`` the + tool-call arguments locally; malformed args raise after billing, so the error + trace must keep the provider-reported tokens. Exercises the real + ``LiteLLMLLM.call_with_tools`` stash that ``LiteLLMRouterLLM`` also inherits + (the wrapper-level tools test uses a provider that already stashes).""" + llm = LLMProvider(provider="litellm", api_key="test-key", base_url="https://example.test/v1", model="gpt-4o-mini") + # Valid response + usage, but the tool arguments are not valid JSON. + llm._provider_impl._acompletion = AsyncMock(return_value=_litellm_tool_response_with_usage("{not valid json")) + + with pytest.raises(json.JSONDecodeError): + await llm.call_with_tools( + messages=[{"role": "user", "content": "x"}], + tools=[], + scope="tools", + max_retries=0, + ) + + assert len(registered_recorder.records) == 1 + r = registered_recorder.records[0] + assert r.status == "error" + assert r.input_tokens == 140 + assert r.output_tokens == 18 + assert r.cached_tokens == 20 + + @pytest.mark.asyncio async def test_configured_provider_binds_bank_context(registered_recorder): llm = LLMProvider(provider="mock", api_key="", base_url="", model="mock")