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.
This commit is contained in:
redshift
2026-09-17 20:10:44 +02:00
parent b5c2b1b10b
commit 74bfd89ce5
2 changed files with 118 additions and 4 deletions
+11 -3
View File
@@ -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,
+107 -1
View File
@@ -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,