From 74bfd89ce54741825e2012ded254f58c49354fa4 Mon Sep 17 00:00:00 2001 From: redshift <213178690+1ftredsh@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:44:25 +0200 Subject: [PATCH] feat(ehbp): report cache splits in FINALIZE payment events The FINALIZE event for EHBP requests carried only input/output token counts, so cache reads billed at the discounted cached rate stayed invisible to usage analytics even when the enclave reported cached_prompt_tokens. Add cache_read/cache_creation token counts and their msats to the ehbp_usage event. Providers that report no cache split emit a 0 for each field rather than dropping the keys, so consumers see a stable event schema. --- routstr/upstream/ehbp.py | 14 ++- tests/unit/test_ehbp_finalize_payment.py | 108 ++++++++++++++++++++++- 2 files changed, 118 insertions(+), 4 deletions(-) diff --git a/routstr/upstream/ehbp.py b/routstr/upstream/ehbp.py index 74051b00..54e41c10 100644 --- a/routstr/upstream/ehbp.py +++ b/routstr/upstream/ehbp.py @@ -395,9 +395,7 @@ def _inject_cost_response_headers(headers: dict[str, str], cost_info: dict) -> N headers["X-Routstr-Computed-Cost-Msats"] = str(cost_info["computed_msats"]) headers["X-Routstr-Input-Cost-Msats"] = str(cost_info["input_msats"]) headers["X-Routstr-Output-Cost-Msats"] = str(cost_info["output_msats"]) - headers["X-Routstr-Cache-Read-Msats"] = str( - cost_info.get("cache_read_msats", 0) - ) + headers["X-Routstr-Cache-Read-Msats"] = str(cost_info.get("cache_read_msats", 0)) headers["X-Routstr-Cache-Creation-Msats"] = str( cost_info.get("cache_creation_msats", 0) ) @@ -675,6 +673,16 @@ async def finalize_ehbp_actual_cost_payment( "cost_charged": total_cost_msats, "input_tokens": cost_info.get("input_tokens", 0), "output_tokens": cost_info.get("output_tokens", 0), + # Cache splits are only knowable when the enclave reports + # ``cached_prompt_tokens``; absent that they are a measured zero on + # the token counts the provider did report (not an unknown), so the + # event key set stays stable for usage-analytics consumers. + "cache_read_input_tokens": cost_info.get("cache_read_input_tokens", 0), + "cache_creation_input_tokens": cost_info.get( + "cache_creation_input_tokens", 0 + ), + "cache_read_msats": cost_info.get("cache_read_msats", 0), + "cache_creation_msats": cost_info.get("cache_creation_msats", 0), "balance": key.balance, "reserved_balance": key.reserved_balance, "total_spent": key.total_spent, diff --git a/tests/unit/test_ehbp_finalize_payment.py b/tests/unit/test_ehbp_finalize_payment.py index 2f905e63..6e428a7f 100644 --- a/tests/unit/test_ehbp_finalize_payment.py +++ b/tests/unit/test_ehbp_finalize_payment.py @@ -1,6 +1,8 @@ from __future__ import annotations -from typing import Any, AsyncGenerator +import logging +from contextlib import contextmanager +from typing import Any, AsyncGenerator, Iterator from unittest.mock import AsyncMock, MagicMock import pytest @@ -106,6 +108,110 @@ async def test_finalize_actual_cost_payment_updates_balance_and_releases_reserve assert updated.total_spent == 1_200 +@contextmanager +def _capture_payments_logs() -> Iterator[list[logging.LogRecord]]: + """Collect ``routstr.payments`` records for the duration of the block. + + ``setup_logging()`` sets ``propagate=False`` on the ``routstr`` logger, so + pytest's ``caplog`` (attached at the root) never sees these records; a + handler on the payments logger itself does. + """ + records: list[logging.LogRecord] = [] + + class _RecordingHandler(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + records.append(record) + + payments_logger = logging.getLogger("routstr.payments") + handler = _RecordingHandler(level=logging.INFO) + previous_level = payments_logger.level + payments_logger.addHandler(handler) + payments_logger.setLevel(logging.INFO) + try: + yield records + finally: + payments_logger.removeHandler(handler) + payments_logger.setLevel(previous_level) + + +@pytest.mark.asyncio +async def test_finalize_actual_cost_payment_logs_cache_tokens( + session: AsyncSession, +) -> None: + """The FINALIZE event carries the cache splits, not just input/output.""" + key = ApiKey(hashed_key="ehbp-cache-logging", balance=10_000) + session.add(key) + await session.commit() + await pay_for_request(key, 3_000, session) + reservation = await get_reservation_snapshot(key, session) + + with _capture_payments_logs() as records: + charged = await finalize_ehbp_actual_cost_payment( + key, + session, + reserved_cost_for_model=3_000, + model_id="tinfoil/glm-5-2", + cost_info={ + "total_msats": 1_200, + "input_tokens": 5, + "output_tokens": 20, + "input_msats": 500, + "output_msats": 700, + "cache_read_input_tokens": 64, + "cache_creation_input_tokens": 0, + "cache_read_msats": 12, + "cache_creation_msats": 0, + }, + reservation_snapshot=reservation, + ) + + assert charged == 1_200 + finalize_records = [ + record for record in records if record.getMessage() == "FINALIZE" + ] + assert len(finalize_records) == 1 + record = finalize_records[0] + assert record.finalize_type == "ehbp_usage" + assert record.input_tokens == 5 + assert record.output_tokens == 20 + assert record.cache_read_input_tokens == 64 + assert record.cache_creation_input_tokens == 0 + assert record.cache_read_msats == 12 + assert record.cache_creation_msats == 0 + + +@pytest.mark.asyncio +async def test_finalize_actual_cost_payment_logs_zero_cache_when_absent( + session: AsyncSession, +) -> None: + """Providers that report no cache split still emit a stable key set.""" + key = ApiKey(hashed_key="ehbp-no-cache-logging", balance=10_000) + session.add(key) + await session.commit() + await pay_for_request(key, 3_000, session) + reservation = await get_reservation_snapshot(key, session) + + with _capture_payments_logs() as records: + await finalize_ehbp_actual_cost_payment( + key, + session, + reserved_cost_for_model=3_000, + model_id="tinfoil/glm-5-2", + cost_info={ + "total_msats": 1_200, + "input_tokens": 10, + "output_tokens": 20, + }, + reservation_snapshot=reservation, + ) + + record = next(record for record in records if record.getMessage() == "FINALIZE") + assert record.cache_read_input_tokens == 0 + assert record.cache_creation_input_tokens == 0 + assert record.cache_read_msats == 0 + assert record.cache_creation_msats == 0 + + @pytest.mark.asyncio async def test_unmeasured_ehbp_releases_reservation( session: AsyncSession,