mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-10-05 12:28:22 +00:00
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:
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user