From d8f615ec013330d261189ffe0bbe13c05a4858b7 Mon Sep 17 00:00:00 2001 From: kevin-lozada-santos Date: Sun, 13 Sep 2026 22:50:56 -0400 Subject: [PATCH] fix: preserve cached and reasoning tokens without double billing --- CHANGELOG.md | 9 + src/tokenops/control/boundary.py | 22 +-- src/tokenops/control/core.py | 13 +- src/tokenops/control/integration.py | 8 +- .../control/policies/context_compaction.py | 13 +- src/tokenops/control/usage.py | 43 +++++ src/tokenops/providers/anthropic.py | 6 +- src/tokenops/providers/openai.py | 6 + src/tokenops/providers/types.py | 4 + tests/test_provider_usage.py | 161 ++++++++++++++++++ 10 files changed, 252 insertions(+), 33 deletions(-) create mode 100644 src/tokenops/control/usage.py create mode 100644 tests/test_provider_usage.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 2458fc2..7708645 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed + +- Carry cache-read and reasoning tokens from native SDK responses, bundled adapters, + and flat agent steps into disjoint `Usage` buckets, avoiding double billing while + retaining total prompt size for context-compaction trends. Flat dispatch/step + input and output remain inclusive totals; direct `Usage` producers must exclude + cached/reasoning subsets from input/output (#137, thanks @kevin-lozada-santos). + Cache-write premiums and streaming usage remain outside this change. + ## [0.3.0] - 2026-09-12 ### Added diff --git a/src/tokenops/control/boundary.py b/src/tokenops/control/boundary.py index f6e8e44..2e85eb7 100644 --- a/src/tokenops/control/boundary.py +++ b/src/tokenops/control/boundary.py @@ -15,6 +15,7 @@ from tokenops.control.attribution import _build_attribution, require_registration from tokenops.control.context import current_governance, current_span from tokenops.control.core import NodeType, Observation, Usage +from tokenops.control.usage import usage_from_counts _KIND_MAP: dict[str, NodeType] = { "llm": "llm", @@ -73,24 +74,9 @@ def observation_from_crossing( if node_type == "llm": usage_obj = getattr(result, "usage", None) - if usage_obj is not None: - usage = Usage( - input=int( - getattr(usage_obj, "prompt_tokens", 0) - or getattr(usage_obj, "input_tokens", 0) - or 0 - ), - output=int( - getattr(usage_obj, "completion_tokens", 0) - or getattr(usage_obj, "output_tokens", 0) - or 0 - ), - ) - else: - usage = Usage( - input=int(getattr(result, "input_tokens", 0) or 0), - output=int(getattr(result, "output_tokens", 0) or 0), - ) + usage = usage_from_counts( + usage_obj if usage_obj is not None else result, native=usage_obj is not None + ) text = getattr(result, "content", None) if text is None and hasattr(result, "completion"): text = getattr(result, "completion", result) diff --git a/src/tokenops/control/core.py b/src/tokenops/control/core.py index 7bc8471..6f36c67 100644 --- a/src/tokenops/control/core.py +++ b/src/tokenops/control/core.py @@ -61,11 +61,14 @@ @dataclass(frozen=True, kw_only=True) class Usage: - """Provider-reported token *totals* for one call — never the streamed text. - - ``cached`` and ``reasoning`` mirror the hidden, costly categories providers report - (OpenAI ``*_tokens_details``; Anthropic ``cache_read_input_tokens``). Track them or - the spend most likely to surprise you stays invisible. + """Disjoint, independently priceable token buckets for one call. + + ``input`` excludes cache reads; ``output`` excludes reasoning. ``cached`` and + ``reasoning`` are billed separately, never added to inclusive provider totals. + Prompt/context size is ``input + cached``; completion size is ``output + reasoning``. + Dispatch adapters and agent steps report inclusive totals and are normalized at + the boundary. Native Anthropic input already excludes cache reads. Cache-write + premium accounting is not yet supported. """ input: int = 0 diff --git a/src/tokenops/control/integration.py b/src/tokenops/control/integration.py index 4a28e5d..be3ed9c 100644 --- a/src/tokenops/control/integration.py +++ b/src/tokenops/control/integration.py @@ -24,8 +24,9 @@ from chronicle import wrap_llm from tokenops.control.context import current_span -from tokenops.control.core import Attribution, CallRequest, Observation, Usage +from tokenops.control.core import Attribution, CallRequest, Observation from tokenops.control.crossing import install_crossing_hook +from tokenops.control.usage import usage_from_counts def tool_signature(name: str, args) -> str: @@ -62,10 +63,7 @@ def step_to_observation( span = _span_fields(service=service) if action == "model": tu = getattr(step, "tokens", None) - usage = Usage( - input=getattr(tu, "input_tokens", 0) if tu else 0, - output=getattr(tu, "output_tokens", 0) if tu else 0, - ) + usage = usage_from_counts(tu) boundary_tags = { "node_type": "llm", "provider": provider, diff --git a/src/tokenops/control/policies/context_compaction.py b/src/tokenops/control/policies/context_compaction.py index e9b2925..b4e155d 100644 --- a/src/tokenops/control/policies/context_compaction.py +++ b/src/tokenops/control/policies/context_compaction.py @@ -2,7 +2,7 @@ LLD row: Detect: est_input ≥ ctx_max OR est_input rising over recent(run, W) (estimate from last - llm step's usage.input, never tokenize on the hot path). + llm step's usage.input + usage.cached, never tokenize on the hot path). Fix: MUTATE the outgoing prompt: (1) move volatile values below the static prefix to restore the prompt-cache discount, (2) dedup tool outputs by hash, (3) summarize only filler, pinning system prompt, schema, constraints, state. No hook → degrade @@ -46,10 +46,15 @@ def pre_call(self, request: CallRequest, view: LedgerView) -> Signal | None: rising = False if len(recent_llm) >= 2: rising = all( - (a.usage.input if a.usage else 0) <= (b.usage.input if b.usage else 0) + (a.usage.input + a.usage.cached if a.usage else 0) + <= (b.usage.input + b.usage.cached if b.usage else 0) for a, b in zip(recent_llm, recent_llm[1:]) - ) and (recent_llm[-1].usage.input if recent_llm[-1].usage else 0) > ( - recent_llm[0].usage.input if recent_llm[0].usage else 0 + ) and ( + recent_llm[-1].usage.input + recent_llm[-1].usage.cached + if recent_llm[-1].usage + else 0 + ) > ( + recent_llm[0].usage.input + recent_llm[0].usage.cached if recent_llm[0].usage else 0 ) if est >= self.ctx_max or (rising and est >= self.ctx_max // 2): return Signal( diff --git a/src/tokenops/control/usage.py b/src/tokenops/control/usage.py new file mode 100644 index 0000000..c21d668 --- /dev/null +++ b/src/tokenops/control/usage.py @@ -0,0 +1,43 @@ +"""Normalize provider counts once into independently priceable Usage buckets. + +SDK totals include details, except native Anthropic input which excludes cache reads. +Flat dispatch/agent-step counts always use inclusive input/output totals. +Cache creation pricing remains a separate follow-up. +""" + +from __future__ import annotations + +from tokenops.control.core import Usage + + +def _count(obj: object, primary: str, fallback: str = "") -> int: + value = getattr(obj, primary, None) + if value is None and fallback: + value = getattr(obj, fallback, None) + return int(value or 0) + + +def usage_from_counts(counts: object, *, native: bool = False) -> Usage: + """Read SDK usage (native=True) or the inclusive flat dispatch/step contract.""" + input_tokens = _count(counts, "prompt_tokens", "input_tokens") + output_tokens = _count(counts, "completion_tokens", "output_tokens") + if native and hasattr(counts, "cache_read_input_tokens"): + # Anthropic input_tokens already excludes cached reads. + cached = _count(counts, "cache_read_input_tokens") + reasoning = 0 + else: + if native: + input_details = getattr(counts, "prompt_tokens_details", None) + if input_details is None: + input_details = getattr(counts, "input_tokens_details", None) + output_details = getattr(counts, "completion_tokens_details", None) + if output_details is None: + output_details = getattr(counts, "output_tokens_details", None) + cached = _count(input_details, "cached_tokens") + reasoning = _count(output_details, "reasoning_tokens") + else: + cached = _count(counts, "cached_tokens") + reasoning = _count(counts, "reasoning_tokens") + input_tokens -= cached + output_tokens -= reasoning + return Usage(input=input_tokens, output=output_tokens, cached=cached, reasoning=reasoning) diff --git a/src/tokenops/providers/anthropic.py b/src/tokenops/providers/anthropic.py index 4922054..49edf37 100644 --- a/src/tokenops/providers/anthropic.py +++ b/src/tokenops/providers/anthropic.py @@ -27,6 +27,10 @@ def messages( text_blocks = [b.text for b in response.content if b.type == "text"] return ModelResponse( content="".join(text_blocks), - input_tokens=usage.input_tokens if usage else 0, + # ModelResponse uses inclusive totals even though the native SDK does not. + input_tokens=(usage.input_tokens + (getattr(usage, "cache_read_input_tokens", 0) or 0)) + if usage + else 0, + cached_tokens=getattr(usage, "cache_read_input_tokens", 0) or 0, output_tokens=usage.output_tokens if usage else 0, ) diff --git a/src/tokenops/providers/openai.py b/src/tokenops/providers/openai.py index 7a2eb9d..e449b51 100644 --- a/src/tokenops/providers/openai.py +++ b/src/tokenops/providers/openai.py @@ -31,6 +31,12 @@ def chat( content=content, input_tokens=usage.prompt_tokens if usage else 0, output_tokens=usage.completion_tokens if usage else 0, + cached_tokens=getattr(getattr(usage, "prompt_tokens_details", None), "cached_tokens", 0) + or 0, + reasoning_tokens=getattr( + getattr(usage, "completion_tokens_details", None), "reasoning_tokens", 0 + ) + or 0, ) # Anthropic via OpenAI-compatible path not used; delegate to anthropic module diff --git a/src/tokenops/providers/types.py b/src/tokenops/providers/types.py index 8c4fa63..c6ec7ff 100644 --- a/src/tokenops/providers/types.py +++ b/src/tokenops/providers/types.py @@ -5,6 +5,10 @@ @dataclass class ModelResponse: + """Inclusive input/output totals; cached/reasoning are subsets, not extra tokens.""" + content: str input_tokens: int = 0 output_tokens: int = 0 + cached_tokens: int = 0 + reasoning_tokens: int = 0 diff --git a/tests/test_provider_usage.py b/tests/test_provider_usage.py new file mode 100644 index 0000000..95d602e --- /dev/null +++ b/tests/test_provider_usage.py @@ -0,0 +1,161 @@ +"""Provider usage reaches pricing exactly once, in disjoint token buckets.""" + +from types import SimpleNamespace as NS +from unittest.mock import patch + +from tokenops.control.boundary import observation_from_crossing +from tokenops.control.context import SpanContext, run_scope +from tokenops.control.core import Attribution, Usage +from tokenops.control.integration import step_to_observation +from tokenops.control.models import RunRegistration +from tokenops.control.pricing import Rate, build_price_book +from tokenops.providers import anthropic as ant +from tokenops.providers import openai as oa +from tokenops.providers.types import ModelResponse + +PROMPT, CACHED, OUTPUT = 814835, 677898, 17947 +PRICE = build_price_book({"fixture": Rate(500000, 3000000, cached=50000)}) + + +def crossing(result, provider="openai"): + with run_scope(RunRegistration(run_id="offline-137"), SpanContext(span_id="s", service="test")): + return observation_from_crossing( + boundary_id="test.chat", + kind="llm", + service="test", + input_state={}, + result=result, + provider=provider, + model="fixture", + ).usage + + +def oa_response(): + return NS( + usage=NS( + prompt_tokens=PROMPT, + completion_tokens=OUTPUT, + prompt_tokens_details=NS(cached_tokens=CACHED), + completion_tokens_details=NS(reasoning_tokens=100), + ), + choices=[NS(message=NS(content="offline"))], + ) + + +def ant_response(): + return NS( + usage=NS( + input_tokens=PROMPT - CACHED, + output_tokens=OUTPUT, + cache_read_input_tokens=CACHED, + cache_creation_input_tokens=0, + ), + content=[NS(type="text", text="offline")], + ) + + +def test_openai_boundary_preserves_cache_and_reasoning(): + u = crossing(oa_response()) + assert u == Usage(input=136937, output=17847, cached=677898, reasoning=100) + + +def test_responses_shape_preserves_cache_and_reasoning(): + u = crossing( + NS( + usage=NS( + input_tokens=1000, + output_tokens=200, + input_tokens_details=NS(cached_tokens=800), + output_tokens_details=NS(reasoning_tokens=150), + ) + ) + ) + assert u == Usage(input=200, output=50, cached=800, reasoning=150) + + +def test_anthropic_boundary_prices_native_disjoint_cache_reads(): + u = crossing(ant_response(), "anthropic") + assert u == Usage(input=136937, output=17947, cached=677898) + assert PRICE("anthropic", "fixture", u) == 156204 + + +def test_flat_fallback_normalizes_inclusive_fields(): + assert crossing( + NS(input_tokens=1000, output_tokens=200, cached_tokens=800, reasoning_tokens=150) + ) == Usage(input=200, output=50, cached=800, reasoning=150) + + +def test_agent_step_normalizes_inclusive_fields(): + step = NS( + action="model", + tokens=NS(input_tokens=1000, output_tokens=200, cached_tokens=800, reasoning_tokens=150), + ) + u = step_to_observation( + step, Attribution(user="test", agent="test", run_id="offline"), ts=0 + ).usage + assert u == Usage(input=200, output=50, cached=800, reasoning=150) + + +def test_bundled_openai_adapter_preserves_details_through_boundary(): + client = NS(chat=NS(completions=NS(create=lambda **kw: oa_response()))) + with patch.object(oa, "OpenAI", return_value=client): + result = oa.chat("fixture", []) + assert result.cached_tokens == 677898 + assert crossing(result) == Usage(input=136937, output=17847, cached=677898, reasoning=100) + + +def test_bundled_anthropic_adapter_preserves_details_through_boundary(): + client = NS(messages=NS(create=lambda **kw: ant_response())) + with patch.object(ant.anthropic, "Anthropic", return_value=client): + result = ant.messages("fixture", []) + assert result.input_tokens == 814835 + assert result.cached_tokens == 677898 + assert crossing(result, "anthropic") == Usage(input=136937, output=17947, cached=677898) + + +def test_legacy_two_count_response_control(): + assert crossing(ModelResponse("legacy", 12, 4)) == Usage(input=12, output=4) + + +def test_high_cache_fixture_is_priced_once(): + assert PRICE("openai", "fixture", crossing(oa_response())) == 156204 + + +def test_details_can_be_missing_or_none(): + for details in ({}, {"prompt_tokens_details": None, "completion_tokens_details": None}): + assert crossing(NS(usage=NS(prompt_tokens=12, completion_tokens=4, **details))) == Usage( + input=12, output=4 + ) + assert crossing(NS(usage=None)) == Usage() + + +def test_zero_totals_do_not_fall_through_to_another_alias(): + assert ( + crossing( + NS(usage=NS(prompt_tokens=0, input_tokens=99, completion_tokens=0, output_tokens=99)) + ) + == Usage() + ) + + +def test_reasoning_is_not_double_billed(): + price = build_price_book({"x": Rate(0, 1000000, reasoning=2000000)}) + usage = crossing(NS(input_tokens=0, output_tokens=200, reasoning_tokens=150)) + assert usage == Usage(output=50, reasoning=150) + assert price("", "x", usage) == 350 + + +def test_context_trend_includes_cached_prompt_tokens(): + from tokenops.control.core import CallRequest + from tokenops.control.policies.context_compaction import ContextCompactionDetector + + request = CallRequest( + attr=Attribution(user="t", agent="t", run_id="t"), + provider="fixture", + model="fixture", + estimated_input_tokens=6000, + ) + steps = [NS(node_type="llm", usage=Usage(input=1000, cached=n)) for n in (3000, 5000, 7000)] + signal = ContextCompactionDetector(10000).pre_call(request, NS(recent=lambda *args: steps)) + assert signal is not None + assert signal.evidence["rising"] is True