fix: minor usage tracking improvements

application/api/v1/routes.py

@@ -121,10 +121,13 @@ def _conversation_belongs_to_agent(
 
 
 def _response_usage(agent: Any) -> Dict[str, Any]:
-    """Return the upstream provider's latest OpenAI-shaped usage object."""
-    usage = getattr(getattr(agent, "llm", None), "_last_usage", None)
-    if isinstance(usage, dict):
-        return usage
+    """Return the turn's cumulative usage in Chat Completions shape.
+
+    Reads the per-instance accumulator so multi-round tool turns report
+    the sum of every LLM call, matching what ``token_usage`` rows bill;
+    the accumulator carries provider-exact counts whenever the upstream
+    reported them (see ``_prefer_provider_usage``).
+    """
     tokens = getattr(getattr(agent, "llm", None), "token_usage", {}) or {}
     prompt = int(tokens.get("prompt_tokens", 0) or 0)
     completion = int(tokens.get("generated_tokens", 0) or 0)

application/llm/base.py

@@ -317,10 +317,12 @@ class BaseLLM(ABC):
     ):
         # Non-streaming counterpart to ``_emit_stream_finished_log``. Paired
         # with ``llm_gen_start`` so cost dashboards can join start/finish for
-        # non-streaming calls just as they do for streams. Token counts are
-        # client-side estimates from ``gen_token_usage``; ``status`` is
-        # ``"error"`` when the call raised. A distinct event name keeps
-        # non-stream calls out of stream dashboards.
+        # non-streaming calls just as they do for streams. Token counts come
+        # from ``gen_token_usage``: provider-exact when the vendor reported
+        # usage (OpenAI-family chat + Responses), tiktoken estimates
+        # otherwise; ``status`` is ``"error"`` when the call raised. A
+        # distinct event name keeps non-stream calls out of stream
+        # dashboards.
         extra = {
             "model": model,
             "provider": self.provider_name,
@@ -361,10 +363,10 @@ class BaseLLM(ABC):
         error=None,
     ):
         # Paired with ``llm_stream_start`` so cost dashboards can sum tokens
-        # by user/agent/provider. Token counts are client-side estimates
-        # from ``stream_token_usage``; vendor-reported counts (incl.
-        # ``cached_tokens`` for prompt caching) require per-provider
-        # extraction in each ``_raw_gen_stream`` and aren't wired yet.
+        # by user/agent/provider. Token counts come from
+        # ``stream_token_usage``: provider-exact when the vendor reported
+        # usage (OpenAI-family via ``stream_options.include_usage`` and the
+        # Responses API), tiktoken estimates for providers that don't.
         extra = {
             "model": model,
             "provider": self.provider_name,

application/llm/openai.py

@@ -442,8 +442,10 @@ class OpenAILLM(BaseLLM):
             request_params["tools"] = tools
         if response_format:
             request_params["response_format"] = response_format
+        self._last_usage = None
         response = self.client.chat.completions.create(**request_params)
         logging.debug("OpenAI request completed")
+        self._record_chat_usage(getattr(response, "usage", None))
         if tools:
             return response.choices[0]
         else:
@@ -506,11 +508,19 @@ class OpenAILLM(BaseLLM):
             request_params["tools"] = tools
         if response_format:
             request_params["response_format"] = response_format
+        # Ask for the terminal usage-only chunk (choices=[]) so token rows
+        # get provider-exact counts; servers that ignore stream_options
+        # leave the tiktoken-estimate fallback in place.
+        stream_options = dict(request_params.get("stream_options") or {})
+        stream_options.setdefault("include_usage", True)
+        request_params["stream_options"] = stream_options
+        self._last_usage = None
         response = self.client.chat.completions.create(**request_params)
 
         try:
             for line in response:
                 logging.debug(f"OpenAI stream line: {line}")
+                self._record_chat_usage(getattr(line, "usage", None))
                 if not getattr(line, "choices", None):
                     continue
 
@@ -782,6 +792,42 @@ class OpenAILLM(BaseLLM):
         result["summary"] = serialized
         return result
 
+    def _record_chat_usage(self, usage) -> None:
+        """Capture provider-reported Chat Completions usage for this call.
+
+        Counterpart to ``_record_responses_metadata`` for the chat path;
+        provider totals are stored as-is (see ``_prefer_provider_usage``).
+        """
+        if usage is None:
+            return
+        try:
+            prompt = int(getattr(usage, "prompt_tokens", 0) or 0)
+            completion = int(getattr(usage, "completion_tokens", 0) or 0)
+            total = int(getattr(usage, "total_tokens", 0) or 0)
+        except (TypeError, ValueError):
+            return
+        if not prompt and not completion:
+            # A zeroed usage object (some proxies) must not clobber estimates.
+            return
+        result = {
+            "prompt_tokens": prompt,
+            "completion_tokens": completion,
+            "total_tokens": total or prompt + completion,
+        }
+        input_details = getattr(usage, "prompt_tokens_details", None)
+        output_details = getattr(usage, "completion_tokens_details", None)
+        try:
+            cached = int(getattr(input_details, "cached_tokens", 0) or 0)
+            reasoning = int(getattr(output_details, "reasoning_tokens", 0) or 0)
+        except (TypeError, ValueError):
+            cached = 0
+            reasoning = 0
+        if cached:
+            result["prompt_tokens_details"] = {"cached_tokens": cached}
+        if reasoning:
+            result["completion_tokens_details"] = {"reasoning_tokens": reasoning}
+        self._last_usage = result
+
     def _record_responses_metadata(self, response):
         rid = getattr(response, "id", None)
         if rid:

application/usage.py

@@ -1,5 +1,6 @@
 import logging
 import time
+from typing import Any, Dict
 
 from application.storage.db.repositories.token_usage import TokenUsageRepository
 from application.storage.db.session import db_session
@@ -140,8 +141,17 @@ def _persist_call_usage(llm, call_usage):
         logger.exception("token_usage persist failed")
 
 
-def _prefer_provider_usage(llm, call_usage):
-    """Replace estimates with upstream counts when a provider reported them."""
+def _prefer_provider_usage(llm: Any, call_usage: Dict[str, int]) -> Dict[str, int]:
+    """Replace estimates with upstream counts when a provider reported them.
+
+    Invariant: provider totals are billing-parity bins. Upstream
+    ``prompt_tokens`` already includes cached-read tokens and
+    ``completion_tokens`` already includes reasoning/refusal tokens, so
+    they map 1:1 onto our two columns. Never subtract the
+    ``*_tokens_details`` breakdowns (``cached_tokens``,
+    ``reasoning_tokens``) back out of these bins — that would break
+    parity with what providers bill.
+    """
     reported = getattr(llm, "_last_usage", None)
     if not isinstance(reported, dict):
         return call_usage

tests/api/v1/test_routes_extended.py

@@ -194,6 +194,31 @@ def _patch_v1_db(conn):
         yield
 
 
+@pytest.mark.unit
+def test_response_usage_reports_cumulative_turn_totals():
+    from types import SimpleNamespace
+
+    from application.api.v1.routes import _response_usage
+
+    # Multi-round tool turns must report the accumulator's turn total,
+    # not the final round's provider snapshot.
+    agent = SimpleNamespace(
+        llm=SimpleNamespace(
+            _last_usage={
+                "prompt_tokens": 5,
+                "completion_tokens": 1,
+                "total_tokens": 6,
+            },
+            token_usage={"prompt_tokens": 30, "generated_tokens": 12},
+        )
+    )
+    assert _response_usage(agent) == {
+        "prompt_tokens": 30,
+        "completion_tokens": 12,
+        "total_tokens": 42,
+    }
+
+
 @pytest.mark.unit
 class TestLookupAgentHappy:
     def test_returns_agent_for_valid_key(self, pg_conn):

tests/llm/test_openai_responses.py

@@ -415,6 +415,103 @@ def test_chat_path_drops_tool_controls_when_tools_are_unavailable(monkeypatch):
     assert "parallel_tool_calls" not in request
 
 
+@pytest.mark.unit
+def test_chat_nonstream_records_provider_usage(monkeypatch):
+    llm = _make_llm(monkeypatch, ModelCapabilities(api_flavor="chat_completions"))
+    llm._last_usage = {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}
+    llm.client.chat.completions.create = MagicMock(
+        return_value=_ns(
+            choices=[_ns(message=_ns(content="ok"))],
+            usage=_ns(prompt_tokens=11, completion_tokens=5, total_tokens=16),
+        )
+    )
+
+    result = llm._raw_gen(llm, "model", [{"role": "user", "content": "hi"}])
+
+    assert result == "ok"
+    assert llm._last_usage == {
+        "prompt_tokens": 11,
+        "completion_tokens": 5,
+        "total_tokens": 16,
+    }
+
+
+@pytest.mark.unit
+def test_chat_stream_requests_and_records_terminal_usage(monkeypatch):
+    llm = _make_llm(monkeypatch, ModelCapabilities(api_flavor="chat_completions"))
+    llm._last_usage = {"prompt_tokens": 9, "completion_tokens": 9, "total_tokens": 18}
+    llm.client.chat.completions.create = MagicMock(
+        return_value=[
+            _ns(choices=[_ns(delta=_ns(content="Hi"), finish_reason=None)]),
+            _ns(
+                choices=[],
+                usage=_ns(prompt_tokens=7, completion_tokens=2, total_tokens=9),
+            ),
+        ]
+    )
+
+    out = list(
+        llm._raw_gen_stream(llm, "model", [{"role": "user", "content": "hi"}])
+    )
+
+    assert out == ["Hi"]
+    request = llm.client.chat.completions.create.call_args.kwargs
+    assert request["stream_options"] == {"include_usage": True}
+    assert llm._last_usage == {
+        "prompt_tokens": 7,
+        "completion_tokens": 2,
+        "total_tokens": 9,
+    }
+
+
+@pytest.mark.unit
+def test_chat_stream_without_usage_clears_stale_counts(monkeypatch):
+    llm = _make_llm(monkeypatch, ModelCapabilities(api_flavor="chat_completions"))
+    llm._last_usage = {"prompt_tokens": 9, "completion_tokens": 9, "total_tokens": 18}
+    llm.client.chat.completions.create = MagicMock(
+        return_value=[_ns(choices=[_ns(delta=_ns(content="Hi"), finish_reason=None)])]
+    )
+
+    assert list(
+        llm._raw_gen_stream(llm, "model", [{"role": "user", "content": "hi"}])
+    ) == ["Hi"]
+    assert llm._last_usage is None
+
+
+@pytest.mark.unit
+def test_zeroed_chat_usage_does_not_clobber_estimates(monkeypatch):
+    llm = _make_llm(monkeypatch, ModelCapabilities(api_flavor="chat_completions"))
+    llm.client.chat.completions.create = MagicMock(
+        return_value=_ns(
+            choices=[_ns(message=_ns(content="ok"))],
+            usage=_ns(prompt_tokens=0, completion_tokens=0, total_tokens=0),
+        )
+    )
+
+    llm._raw_gen(llm, "model", [{"role": "user", "content": "hi"}])
+
+    assert llm._last_usage is None
+
+
+@pytest.mark.unit
+def test_record_chat_usage_captures_detail_bins(monkeypatch):
+    llm = _make_llm(monkeypatch, ModelCapabilities(api_flavor="chat_completions"))
+    llm._record_chat_usage(_ns(
+        prompt_tokens=10,
+        completion_tokens=7,
+        total_tokens=17,
+        prompt_tokens_details=_ns(cached_tokens=4),
+        completion_tokens_details=_ns(reasoning_tokens=3),
+    ))
+    assert llm._last_usage == {
+        "prompt_tokens": 10,
+        "completion_tokens": 7,
+        "total_tokens": 17,
+        "prompt_tokens_details": {"cached_tokens": 4},
+        "completion_tokens_details": {"reasoning_tokens": 3},
+    }
+
+
 @pytest.mark.unit
 def test_build_responses_params_store_with_previous_id(monkeypatch):
     llm = _make_llm(monkeypatch, _responses_caps(), store_responses=True)