laantungir/routstr-core
1
0
Fork 0
mirror of https://github.com/Routstr/routstr-core.git synced 2026-09-14 02:35:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

better price estimation

Browse Source
This commit is contained in:
9qeklajc
2026-08-27 00:41:26 +02:00
parent dd2885ca5d
commit 5f31c49eef
15 changed files with 456 additions and 226 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/client/payments.md
+3 -2
View File
@@ -38,8 +38,9 @@ Every time you make a request to `/v1/chat/completions` (or others), the cost is
`Cost = (Input_Tokens * Price_Input) + (Output_Tokens * Price_Output) + Request_Fee`
- Prices are defined per model (see `/v1/models`).
- If you stream the response, the balance is deducted incrementally or finalized at the end of the stream.
- If your balance hits 0 mid-stream, the connection is closed.
- Routstr reserves an authorization ceiling before forwarding, then finalizes the request at measured token cost.
- If a successful upstream omits usage, Routstr estimates input tokens from the provider-bound request and output tokens from the returned body or streamed deltas, then applies normal model pricing.
- A reservation is only a temporary hold. Missing usage or unusable prices must never turn the full reservation into the charge; if no auditable estimate can be priced, the reservation is released without charge.
### Headers
docs/ehbp-proxy-support.md
+3 -3
View File
@@ -58,10 +58,10 @@ Contains the shared opaque EHBP transport and billing helpers:
- `EHBPForwardingTarget` — provider-specific target URL plus extra headers
- `forward_ehbp_request()` — forwards the encrypted body, captures Tinfoil
usage from a response header or streaming HTTP trailer, and finalizes bearer
billing at actual cost (falling back to max cost when usage is unavailable)
billing at actual cost (releasing the reservation when usage is unavailable)
- `forward_ehbp_x_cashu_request()` — redeems the Cashu token, refunds the full
token on upstream failure, and refunds the difference between the redeemed
amount and actual cost (or max cost when usage is unavailable)
amount and actual cost (or the full amount when usage is unavailable)
### Provider support
@@ -84,7 +84,7 @@ The proxy is a **blind relay** for EHBP requests. It cannot decrypt the body
Cost tracking happens at the proxy level. Routstr reserves or redeems up to
`max_cost_for_model`, then Tinfoil's out-of-band usage header/trailer allows it
to finalize at actual token cost. If trusted usage is missing or invalid, the
proxy safely falls back to max-cost billing.
proxy releases/refunds rather than treating the authorization ceiling as usage.
## End-to-end flow
docs/tinfoil-direct-integration.md
+8 -7
View File
@@ -197,9 +197,9 @@ Properties:
This is the only architecture that preserves end-to-end encryption from the user to the PPQ/Tinfoil enclave while still letting Routstr mediate payment. The key requirement is that usage/cost metadata must be returned outside the encrypted body, ideally as a response header available before body streaming begins.
## Current Routstr problem
## Original Routstr problem
The current EHBP implementation charges successful EHBP requests at `max_cost_for_model` because Routstr cannot decrypt the response body:
The original EHBP implementation charged successful EHBP requests at `max_cost_for_model` because Routstr could not decrypt the response body:
```text
successful EHBP request -> charge full reserved max cost
@@ -369,7 +369,7 @@ Possible approaches:
- PPQ private models are billed per actual input/output tokens.
- Private model rates are available from `GET /v1/models?type=all`.
- Current Routstr EHBP billing at max cost is wrong for PPQ private models.
- Max-cost EHBP fallback is wrong for PPQ private models; current code releases/refunds when trusted usage metadata is absent.
- Direct Tinfoil integration inside Routstr would enable exact usage billing but would make Routstr see plaintext.
- A blind EHBP relay preserves privacy but requires PPQ/Tinfoil to expose usage/cost in plaintext headers/trailers.
- The preferred solution is to keep Routstr blind and have PPQ return billing metadata outside the encrypted body.
@@ -410,8 +410,9 @@ and `routstr/upstream/ehbp.py`.
actual served model's pricing is used for cost calculation.
- `forward_ehbp_request()` (bearer auth): if `X-Tinfoil-Usage-Metrics` is
present in the response header, finalizes with `adjust_payment_for_tokens()`
for exact billing; otherwise falls back to max-cost. Billing uses the
actual served model when it differs from the requested one.
for exact billing; otherwise releases the reservation. The encrypted body
cannot be estimated locally, and the authorization ceiling is not billed.
Billing uses the actual served model when it differs from the requested one.
- `forward_ehbp_x_cashu_request()`: if usage is available, computes the
refund from actual cost instead of max cost, using the actual served
model's pricing when applicable.
@@ -425,10 +426,10 @@ and `routstr/upstream/ehbp.py`.
|---|---|---|
| Bearer, non-streaming | `X-Tinfoil-Usage-Metrics` response header | Exact token cost via `adjust_payment_for_tokens` |
| Bearer, streaming | `X-Tinfoil-Usage-Metrics` HTTP trailer | Exact token cost (h11 captures trailers) |
| Bearer, no usage header/trailer | N/A | Max-cost fallback |
| Bearer, no usage header/trailer | N/A | Release reservation; zero charge |
| X-Cashu, non-streaming | `X-Tinfoil-Usage-Metrics` response header | Refund = `redeemed - actual_cost` |
| X-Cashu, streaming | `X-Tinfoil-Usage-Metrics` HTTP trailer | Refund = `redeemed - actual_cost` (h11 captures trailers) |
| X-Cashu, no usage header/trailer | N/A | Refund = `redeemed - max_cost` |
| X-Cashu, no usage header/trailer | N/A | Full refund |
### Cost response headers
routstr/core/logging.py
+1 -1
View File
@@ -16,7 +16,7 @@ DO NOT modify or remove these messages without updating the usage tracking logic
- The 'token_cost', 'model', 'input_tokens', and 'output_tokens' fields are extracted for dashboard metrics
3. "Max cost payment finalized" (INFO) - routstr/auth.py
- Used as the successful completion fallback when token usage is unavailable
- Used for explicit flat-price/MaxCostData settlements; missing usage alone must not create this charge
- The 'charged_amount', 'model', 'input_tokens', and 'output_tokens' fields are extracted for dashboard metrics
4. "Payment processed successfully" (INFO) - routstr/auth.py
routstr/payment/cost_calculation.py
+12 -36
View File
@@ -70,30 +70,6 @@ def _empty_cost(cls: type[CostData] = CostData) -> CostData:
)
def _unmeasured_cost(max_cost: int) -> MaxCostData:
"""Build the bounded fallback for a response whose usage cannot be measured.
Missing usage must NOT settle at zero — that hands out free inference. The
request was authorized up to ``max_cost`` (the reservation), so the safe,
bounded settlement is to charge exactly that. Token components stay zero
because they are genuinely unknown; ``total_msats`` carries the authorized
max so max-cost finalization debits the reservation instead of nothing.
"""
return MaxCostData(
base_msats=0,
input_msats=0,
output_msats=0,
total_msats=max(0, max_cost),
total_usd=0.0,
input_tokens=0,
output_tokens=0,
cache_read_input_tokens=0,
cache_creation_input_tokens=0,
cache_read_msats=0,
cache_creation_msats=0,
)
async def calculate_cost(
response_data: dict,
max_cost: int,
@@ -134,11 +110,11 @@ async def calculate_cost(
if usage is None:
logger.warning(
"No usage data in response — settling at the reserved max cost "
"(bounded fallback), not zero. Dashboard will show this request "
"as `(0+0)` tokens. Most common cause: upstream stream did not "
"include a final usage chunk (OpenAI-compat backends require "
"`stream_options.include_usage=true`).",
"No usage data or local estimate in response — releasing the "
"reservation without charging it as usage. Dashboard will show "
"this request as `(0+0)` tokens. Most common cause: upstream "
"stream did not include a final usage chunk (OpenAI-compat "
"backends require `stream_options.include_usage=true`).",
extra={
"max_cost_msats": max_cost,
"model": response_data.get("model", "unknown"),
@@ -147,7 +123,7 @@ async def calculate_cost(
else None,
},
)
return _unmeasured_cost(max_cost)
return _empty_cost(MaxCostData)
usage_data = response_data.get("usage") or {}
if not isinstance(usage_data, dict):
@@ -276,10 +252,10 @@ async def calculate_cost(
rates = (input_rate, output_rate, cache_read_rate, cache_creation_rate)
if not all(is_usable_rate(rate) for rate in rates):
logger.warning(
"No usable token pricing — billing at flat MaxCostData. "
"Token counts %s in the upstream response but cannot be "
"priced; the request will appear in dashboards with the "
"raw counts and a fixed max-cost charge.",
"No usable token pricing — releasing the reservation instead of "
"treating its ceiling as the charge. Token counts %s in the "
"upstream response but cannot be converted to money; the request "
"will appear in dashboards with raw counts and a zero charge.",
"are present" if (input_tokens > 0 or output_tokens > 0) else "are zero",
extra={
"base_cost_msats": max_cost,
@@ -291,10 +267,10 @@ async def calculate_cost(
},
)
return MaxCostData(
base_msats=max_cost,
base_msats=0,
input_msats=0,
output_msats=0,
total_msats=max_cost,
total_msats=0,
input_tokens=input_tokens,
output_tokens=output_tokens,
cache_read_input_tokens=cache_read_tokens,
routstr/upstream/base.py
+77 -42
View File
@@ -61,7 +61,7 @@ from .cache_breakpoints import (
inject_anthropic_cache_breakpoints,
is_explicit_cache_model,
)
from .count_tokens import count_tokens_locally
from .count_tokens import MissingUsageEstimator, count_tokens_locally
from .litellm_routing import detect_litellm_prefix
from .rate_limit import UPSTREAM_RATE_LIMIT, classify_rate_limit
@@ -705,8 +705,8 @@ class BaseUpstreamProvider:
# OpenAI-compatible streaming responses omit ``usage`` unless the
# request sets ``stream_options.include_usage = true``. Without it
# we can't reconcile token counts at end of stream and the
# request gets billed at max-cost with zero tokens. Discriminate
# we can't reconcile token counts at end of stream and must use
# the local request/response estimator. Discriminate
# chat-completions-shaped requests by the ``messages`` field so we
# don't poke unrelated endpoints.
if (
@@ -1021,6 +1021,7 @@ class BaseUpstreamProvider:
model_obj: Model | None = None,
reservation_snapshot: ReservationSnapshot | None = None,
client: httpx.AsyncClient | None = None,
request_body: bytes | None = None,
) -> StreamingResponse:
"""Handle streaming chat completion responses with token usage tracking and cost adjustment.
@@ -1041,6 +1042,8 @@ class BaseUpstreamProvider:
snapshot_key, snapshot_session
)
usage_estimator = MissingUsageEstimator(request_body, model_obj)
logger.debug(
"Processing streaming chat completion",
extra={
@@ -1070,7 +1073,7 @@ class BaseUpstreamProvider:
try:
await adjust_payment_for_tokens(
fresh_key,
{"model": last_model_seen or "unknown", "usage": None},
usage_estimator.response_data(last_model_seen),
new_session,
max_cost_for_model,
model_obj,
@@ -1157,6 +1160,7 @@ class BaseUpstreamProvider:
obj = None
if isinstance(obj, dict):
usage_estimator.observe(obj)
self._apply_provider_field(obj)
if obj.get("model"):
last_model_seen = str(obj.get("model"))
@@ -1246,13 +1250,8 @@ class BaseUpstreamProvider:
if fresh_key:
cost_data: dict
try:
adjustment_input = (
usage_chunk_data
if usage_chunk_data is not None
else {
"model": last_model_seen or "unknown",
"usage": None,
}
adjustment_input = usage_estimator.billing_data(
usage_chunk_data, last_model_seen
)
cost_data = await adjust_payment_for_tokens(
fresh_key,
@@ -1361,6 +1360,7 @@ class BaseUpstreamProvider:
requested_model: str | None = None,
model_obj: Model | None = None,
reservation_snapshot: ReservationSnapshot | None = None,
request_body: bytes | None = None,
) -> Response:
"""Handle non-streaming chat completion responses with token usage tracking and cost adjustment.
@@ -1402,6 +1402,13 @@ class BaseUpstreamProvider:
if "id" not in response_json or not isinstance(response_json["id"], str):
response_json["id"] = f"chatcmpl-{uuid.uuid4()}"
if not isinstance(response_json.get("usage"), dict):
usage_estimator = MissingUsageEstimator(request_body, model_obj)
usage_estimator.observe(response_json)
response_json["usage"] = usage_estimator.openai_response_data(
response_json.get("model")
)["usage"]
cost_data = await adjust_payment_for_tokens(
key,
response_json,
@@ -1500,6 +1507,7 @@ class BaseUpstreamProvider:
model_obj: Model | None = None,
reservation_snapshot: ReservationSnapshot | None = None,
client: httpx.AsyncClient | None = None,
request_body: bytes | None = None,
) -> StreamingResponse:
"""Handle streaming Responses API responses with token usage tracking and cost adjustment.
@@ -1511,6 +1519,8 @@ class BaseUpstreamProvider:
Returns:
StreamingResponse with cost data injected at the end
"""
usage_estimator = MissingUsageEstimator(request_body, model_obj)
logger.debug(
"Processing streaming Responses API completion",
extra={
@@ -1541,7 +1551,7 @@ class BaseUpstreamProvider:
try:
await adjust_payment_for_tokens(
fresh_key,
{"model": last_model_seen or "unknown", "usage": None},
usage_estimator.response_data(last_model_seen),
new_session,
max_cost_for_model,
model_obj,
@@ -1633,8 +1643,11 @@ class BaseUpstreamProvider:
"response.incomplete",
):
usage_chunk_data = obj
if not usage_estimator.output_text:
usage_estimator.observe(obj)
return
usage_estimator.observe(obj)
yield prefix + b"data: " + json.dumps(obj).encode() + b"\n\n"
else:
if final:
@@ -1674,13 +1687,8 @@ class BaseUpstreamProvider:
if fresh_key:
cost_data: dict
try:
adjustment_input = (
usage_chunk_data
if usage_chunk_data is not None
else {
"model": last_model_seen or "unknown",
"usage": None,
}
adjustment_input = usage_estimator.billing_data(
usage_chunk_data, last_model_seen
)
cost_data = await adjust_payment_for_tokens(
fresh_key,
@@ -1792,6 +1800,7 @@ class BaseUpstreamProvider:
requested_model: str | None = None,
model_obj: Model | None = None,
reservation_snapshot: ReservationSnapshot | None = None,
request_body: bytes | None = None,
) -> Response:
"""Handle non-streaming Responses API responses with token usage tracking and cost adjustment.
@@ -1831,6 +1840,13 @@ class BaseUpstreamProvider:
},
)
if not isinstance(response_json.get("usage"), dict):
usage_estimator = MissingUsageEstimator(request_body, model_obj)
usage_estimator.observe(response_json)
response_json["usage"] = usage_estimator.response_data(
response_json.get("model")
)["usage"]
if requested_model:
response_json["model"] = requested_model
if "id" not in response_json or not isinstance(response_json["id"], str):
@@ -1945,9 +1961,9 @@ class BaseUpstreamProvider:
return
try:
# Finalize with "unknown" model and no usage to release reservation/charge max cost
# (no routed identity here by design: the None usage settles at
# MaxCostData before any pricing lookup can happen).
# Generic opaque streams have no request/response token seam.
# Missing usage therefore releases the reservation; the hold is
# never treated as evidence of consumption.
await adjust_payment_for_tokens(
key,
{"model": "unknown", "usage": None},
@@ -1982,7 +1998,10 @@ class BaseUpstreamProvider:
requested_model: str | None = None,
model_obj: Model | None = None,
reservation_snapshot: ReservationSnapshot | None = None,
request_body: bytes | None = None,
) -> StreamingResponse:
usage_estimator = MissingUsageEstimator(request_body, model_obj)
async def stream_with_cost(
max_cost_for_model: int,
) -> AsyncGenerator[bytes, None]:
@@ -2036,13 +2055,9 @@ class BaseUpstreamProvider:
usage_finalized = True
return None
try:
fallback: dict = {
"model": last_model_seen or "unknown",
"usage": None,
}
cost_data = await adjust_payment_for_tokens(
fresh_key,
fallback,
usage_estimator.response_data(last_model_seen),
new_session,
max_cost_for_model,
model_obj,
@@ -2081,6 +2096,7 @@ class BaseUpstreamProvider:
try:
data = json.loads(line[6:])
if isinstance(data, dict):
usage_estimator.observe(data)
msg = data.get("message", {})
if msg and msg.get("model"):
last_model_seen = str(msg.get("model"))
@@ -2278,6 +2294,7 @@ class BaseUpstreamProvider:
requested_model: str | None = None,
model_obj: Model | None = None,
reservation_snapshot: ReservationSnapshot | None = None,
request_body: bytes | None = None,
) -> Response:
try:
content = await response.aread()
@@ -2296,6 +2313,12 @@ class BaseUpstreamProvider:
if path.endswith("count_tokens") and "usage" not in response_json:
input_tokens = response_json.get("input_tokens", 0)
response_json["usage"] = {"input_tokens": input_tokens}
elif not isinstance(response_json.get("usage"), dict):
usage_estimator = MissingUsageEstimator(request_body, model_obj)
usage_estimator.observe(response_json)
response_json["usage"] = usage_estimator.response_data(
response_json.get("model")
)["usage"]
cost_data = await adjust_payment_for_tokens(
key,
@@ -2403,11 +2426,18 @@ class BaseUpstreamProvider:
requested_model,
model_obj,
reservation_snapshot,
request_body,
)
response_json = messages_dispatch.coerce_litellm_payload(result)
if requested_model and "model" in response_json:
response_json["model"] = requested_model
if not isinstance(response_json.get("usage"), dict):
usage_estimator = MissingUsageEstimator(request_body, model_obj)
usage_estimator.observe(response_json)
response_json["usage"] = usage_estimator.response_data(
response_json.get("model")
)["usage"]
cost_data = await adjust_payment_for_tokens(
key,
@@ -2521,10 +2551,13 @@ class BaseUpstreamProvider:
requested_model: str | None,
model_obj: Model | None = None,
reservation_snapshot: ReservationSnapshot | None = None,
request_body: bytes | None = None,
) -> StreamingResponse:
"""Re-emit a litellm Anthropic-event iterator as live SSE bytes
with cost reconciliation appended at end of stream."""
usage_estimator = MissingUsageEstimator(request_body, model_obj)
async def stream_with_cost() -> AsyncGenerator[bytes, None]:
usage_finalized = False
last_model_seen: str | None = None
@@ -2541,12 +2574,10 @@ class BaseUpstreamProvider:
if usage_finalized:
return None
logger.warning(
"Finalizing /v1/messages stream with no usage data — "
"client will be billed at max-cost with zero tokens. "
"Likely cause: upstream omitted `usage` from the SSE "
"stream (check that the request includes "
"`stream_options.include_usage=true` and that the "
"upstream actually emits a final usage chunk).",
"Finalizing /v1/messages stream with locally estimated "
"usage because the upstream omitted `usage` from SSE. "
"Check that the upstream emits a final usage chunk; the "
"reservation ceiling will not be used as the charge.",
extra={
"key_hash": key.hashed_key[:8] + "...",
"model": last_model_seen or "unknown",
@@ -2560,13 +2591,9 @@ class BaseUpstreamProvider:
usage_finalized = True
return None
try:
fallback: dict = {
"model": last_model_seen or "unknown",
"usage": None,
}
cost_data = await adjust_payment_for_tokens(
fresh_key,
fallback,
usage_estimator.response_data(last_model_seen),
new_session,
max_cost_for_model,
model_obj,
@@ -2599,6 +2626,7 @@ class BaseUpstreamProvider:
async for annotated in messages_dispatch.stream_annotated_events(
iterator, requested_model
):
usage_estimator.observe(annotated.event)
if annotated.model:
last_model_seen = annotated.model
# Anthropic SSE reports usage cumulatively across
@@ -3046,6 +3074,7 @@ class BaseUpstreamProvider:
requested_model=original_model_id,
model_obj=model_obj,
reservation_snapshot=reservation_snapshot,
request_body=request_body,
)
background_tasks = BackgroundTasks()
background_tasks.add_task(response.aclose)
@@ -3064,6 +3093,7 @@ class BaseUpstreamProvider:
requested_model=original_model_id,
model_obj=model_obj,
reservation_snapshot=reservation_snapshot,
request_body=request_body,
)
finally:
await response.aclose()
@@ -3081,6 +3111,7 @@ class BaseUpstreamProvider:
requested_model=original_model_id,
model_obj=model_obj,
reservation_snapshot=reservation_snapshot,
request_body=request_body,
)
finally:
await response.aclose()
@@ -3131,6 +3162,7 @@ class BaseUpstreamProvider:
model_obj=model_obj,
reservation_snapshot=reservation_snapshot,
client=client,
request_body=request_body,
)
# Handle both non-streaming chat completions and embeddings
@@ -3144,6 +3176,7 @@ class BaseUpstreamProvider:
requested_model=original_model_id,
model_obj=model_obj,
reservation_snapshot=reservation_snapshot,
request_body=request_body,
)
finally:
await response.aclose()
@@ -3408,6 +3441,7 @@ class BaseUpstreamProvider:
model_obj=model_obj,
reservation_snapshot=reservation_snapshot,
client=client,
request_body=transformed_body,
)
if response.status_code == 200:
@@ -3420,6 +3454,7 @@ class BaseUpstreamProvider:
requested_model=original_model_id,
model_obj=model_obj,
reservation_snapshot=reservation_snapshot,
request_body=transformed_body,
)
finally:
await response.aclose()
@@ -4799,11 +4834,11 @@ class BaseUpstreamProvider:
model = payload["model"]
if usage_data is None:
# Settlement invariant: a terminal request is never silently
# zero-billed and never silently keeps the whole token. Unmeasured
# usage settles at the authorization ceiling and refunds the rest.
# No request body is available at this X-Cashu settlement seam, so
# an auditable input/output estimate cannot be built. Refund the
# token rather than treating the authorization ceiling as usage.
logger.warning(
"No usage in streaming Responses API response — settling at authorized max",
"No usage in streaming Responses API response — refunding instead of charging the authorized max",
extra={
"model": model,
"amount": amount,
routstr/upstream/count_tokens.py
+160 -8
View File
@@ -23,7 +23,7 @@ import litellm
from fastapi.responses import Response
from ..core import get_logger
from ..payment.helpers import estimate_tokens
from ..payment.helpers import estimate_prompt_tokens, estimate_tokens
from ..payment.models import Model
logger = get_logger(__name__)
@@ -39,6 +39,13 @@ def _parse_request_body(request_body: bytes | None) -> dict[str, Any]:
return parsed if isinstance(parsed, dict) else {}
def _model_name(model_obj: Model | None, body: dict[str, Any]) -> str:
if model_obj is not None:
return model_obj.forwarded_model_id or model_obj.id or ""
body_model = body.get("model")
return body_model if isinstance(body_model, str) else ""
def _count_with_litellm(model: str, body: dict[str, Any]) -> int:
messages = body.get("messages")
if not isinstance(messages, list):
@@ -67,6 +74,157 @@ def _count_with_litellm(model: str, body: dict[str, Any]) -> int:
)
def _count_text_with_litellm(model: str, text: str) -> int:
return int(
litellm.token_counter(
model=model,
text=text,
count_response_tokens=True,
)
)
def _generated_text(value: object) -> list[str]:
"""Extract generated text/tool arguments without counting response metadata."""
generated_keys = {
"arguments",
"content",
"delta",
"output_text",
"partial_json",
"reasoning",
"reasoning_content",
"text",
"thinking",
}
parts: list[str] = []
def walk(item: object, key: str | None = None) -> None:
if isinstance(item, str):
if key in generated_keys:
parts.append(item)
return
if isinstance(item, list):
for child in item:
walk(child, key)
return
if isinstance(item, dict):
for child_key, child in item.items():
walk(child, child_key)
walk(value)
return parts
class MissingUsageEstimator:
"""Estimate billable usage when an upstream omits its usage trailer.
The reservation is deliberately absent from this class: it is an
authorization ceiling, not an input to usage measurement.
"""
def __init__(self, request_body: bytes | None, model_obj: Model | None) -> None:
self.body = _parse_request_body(request_body)
self.model_name = _model_name(model_obj, self.body)
self._output_parts: list[str] = []
self._input_tokens: int | None = None
def _estimate_input_tokens(self) -> int:
if self._input_tokens is not None:
return self._input_tokens
try:
self._input_tokens = _count_with_litellm(self.model_name, self.body)
except Exception as exc:
self._input_tokens = estimate_prompt_tokens(self.body)
logger.debug(
"litellm request token count failed; using local estimator",
extra={
"model": self.model_name,
"error": str(exc),
"error_type": type(exc).__name__,
"estimated_tokens": self._input_tokens,
},
)
return self._input_tokens
@property
def output_text(self) -> str:
return "".join(self._output_parts)
def observe(self, response_data: object) -> None:
if isinstance(response_data, dict):
event_type = response_data.get("type")
if isinstance(event_type, str) and event_type.endswith(".done"):
# Responses API ``*.done`` events repeat text already streamed
# via ``*.delta`` events; counting both would double-bill.
return
self._output_parts.extend(_generated_text(response_data))
def billing_data(
self,
response_data: dict[str, Any] | None,
model: str | None = None,
) -> dict[str, Any]:
"""Use measured usage when present, otherwise return a local estimate."""
if isinstance(response_data, dict):
usage = response_data.get("usage")
if not isinstance(usage, dict):
nested = response_data.get("response")
usage = nested.get("usage") if isinstance(nested, dict) else None
if isinstance(usage, dict) and usage:
return {
"model": model or response_data.get("model") or self.model_name,
"usage": usage,
}
if not self._output_parts:
self.observe(response_data)
return self.response_data(model)
def response_data(self, model: str | None = None) -> dict[str, Any]:
text = self.output_text
try:
output_tokens = (
_count_text_with_litellm(self.model_name, text) if text else 0
)
except Exception as exc:
output_tokens = len(text) // 3
logger.debug(
"litellm response token count failed; using local estimator",
extra={
"model": self.model_name,
"error": str(exc),
"error_type": type(exc).__name__,
"estimated_tokens": output_tokens,
},
)
input_tokens = max(0, int(self._estimate_input_tokens()))
output_tokens = max(0, int(output_tokens))
return {
"model": model or self.model_name or "unknown",
"usage": {
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"total_tokens": input_tokens + output_tokens,
"estimated": True,
},
}
def openai_response_data(self, model: str | None = None) -> dict[str, Any]:
"""Same estimate in the OpenAI chat-completions usage dialect."""
data = self.response_data(model)
usage = data["usage"]
return {
"model": data["model"],
"usage": {
"prompt_tokens": usage["input_tokens"],
"completion_tokens": usage["output_tokens"],
"total_tokens": usage["total_tokens"],
"estimated": True,
},
}
def count_tokens_locally(
request_body: bytes | None,
model_obj: Model | None,
@@ -75,13 +233,7 @@ def count_tokens_locally(
touching the upstream. Always returns 200; never raises."""
body = _parse_request_body(request_body)
model_name = ""
if model_obj is not None:
model_name = model_obj.forwarded_model_id or model_obj.id or ""
if not model_name:
body_model = body.get("model")
if isinstance(body_model, str):
model_name = body_model
model_name = _model_name(model_obj, body)
input_tokens: int
try:
routstr/upstream/ehbp.py
+27 -82
View File
@@ -319,10 +319,10 @@ async def _compute_ehbp_actual_cost(
) -> dict:
"""Compute the actual cost in msats from Tinfoil usage metrics.
Falls back to ``max_cost_for_model`` when usage is absent (streaming) or
cannot be priced. The result is clamped to ``[min_request_msat,
max_cost_for_model]`` so the refund never exceeds the reservation and is
never zero.
When usage is present, the result is clamped to ``[min_request_msat,
max_cost_for_model]``. Missing or unpriceable usage returns zero: encrypted
EHBP bodies cannot be estimated locally, and the authorization ceiling is
not evidence of consumption.
When the usage-metrics header includes ``model=<name>`` and it differs
from ``model_obj.id``, the actual served model's pricing is used for the
@@ -335,7 +335,7 @@ async def _compute_ehbp_actual_cost(
"""
usage_dict = parse_tinfoil_usage_metrics(usage_header)
if usage_dict is None:
return _build_cost_info(max_cost_for_model)
return _build_cost_info(0)
# The enclave may serve a different model than the one requested (e.g.
# due to failover). The usage-metrics header's ``model=<name>`` carries
@@ -406,19 +406,19 @@ async def _compute_ehbp_actual_cost(
)
except Exception as e:
logger.warning(
"EHBP usage cost calculation failed, falling back to max cost",
"EHBP usage cost calculation failed; releasing instead of charging max cost",
extra={
"model": pricing_model_id,
"error": str(e),
"usage": usage_dict,
},
)
return _build_cost_info(max_cost_for_model, actual_model=actual_model)
return _build_cost_info(0, actual_model=actual_model)
if isinstance(cost, MaxCostData):
logger.warning(
"EHBP calculate_cost returned MaxCostData (no model pricing), "
"falling back to max cost",
"EHBP calculate_cost returned MaxCostData (no usable pricing); "
"releasing instead of charging max cost",
extra={
"model": pricing_model_id,
"max_cost_for_model": max_cost_for_model,
@@ -426,7 +426,7 @@ async def _compute_ehbp_actual_cost(
"cost_total_msats": cost.total_msats,
},
)
return _build_cost_info(max_cost_for_model, actual_model=actual_model)
return _build_cost_info(0, actual_model=actual_model)
if isinstance(cost, CostData):
actual = max(int(cost.total_msats), int(settings.min_request_msat))
clamped = min(actual, max_cost_for_model)
@@ -450,13 +450,13 @@ async def _compute_ehbp_actual_cost(
)
# CostDataError
logger.warning(
"EHBP usage cost calculation error, falling back to max cost",
"EHBP usage cost calculation error; releasing instead of charging max cost",
extra={
"model": pricing_model_id,
"error": getattr(cost, "message", str(cost)),
},
)
return _build_cost_info(max_cost_for_model, actual_model=actual_model)
return _build_cost_info(0, actual_model=actual_model)
def _extract_usage_from_response(
@@ -600,78 +600,25 @@ async def finalize_ehbp_max_cost_payment(
model_id: str,
reservation_snapshot: ReservationSnapshot | None = None,
) -> int:
"""Finalize an EHBP bearer request by charging the reserved max cost.
"""Release an unmeasured EHBP request without charging its reservation.
EHBP responses are encrypted, so Routstr cannot inspect token usage. Unlike
normal completion handlers, this intentionally charges the pre-reserved max
cost and releases the reservation.
The legacy name is retained for compatibility with internal callers. EHBP
responses are encrypted, so no local estimate is possible when the trusted
usage header/trailer is absent.
"""
reservation = reservation_snapshot or await get_reservation_snapshot(key, session)
await _validate_reservation_snapshot(key, reservation, session)
if not await _claim_reservation_for_charge(reservation, session):
return 0
max_cost_for_model = reservation.reserved_msats
billing_key = await get_billing_key(key, session)
key_hash = key.hashed_key
billing_key_hash = billing_key.hashed_key
total_cost_msats = max(0, int(max_cost_for_model))
now = int(time.time())
charged = await _charge_reservation_rows(
session,
billing_key_hash=billing_key_hash,
key_hash=key_hash,
reserved_msats=max_cost_for_model,
charge_msats=total_cost_msats,
)
if not charged:
logger.error(
"Failed to finalize EHBP max-cost payment",
extra={
"key_hash": key_hash[:8] + "...",
"billing_key_hash": billing_key_hash[:8] + "...",
"model": model_id,
"max_cost_for_model": max_cost_for_model,
},
)
await _release_failed_ehbp_charge(reservation, session)
return 0
await session.commit()
await _stop_reservation_heartbeat(reservation.release_id)
await session.refresh(billing_key)
if billing_key.hashed_key != key.hashed_key:
await session.refresh(key)
if total_cost_msats > 0 and ROUTSTR_FEE_PERCENT > 0:
fee_msats = math.ceil(total_cost_msats * ROUTSTR_FEE_PERCENT / 100)
try:
await accumulate_routstr_fee(session, fee_msats)
except Exception as e:
logger.warning(
"Failed to accumulate Routstr fee for EHBP request",
extra={"error": str(e), "fee_msats": fee_msats},
)
payments_logger.info(
"FINALIZE",
key_log_hash = key.hashed_key[:8] + "..."
await release_reservation(reservation, session, reservation.reserved_msats)
logger.warning(
"Released unmeasured EHBP reservation without charging max cost",
extra={
"event": "finalize",
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"key_hash": key_log_hash,
"model": model_id,
"cost_reserved": max_cost_for_model,
"cost_charged": total_cost_msats,
"input_tokens": 0,
"output_tokens": 0,
"balance": billing_key.balance,
"reserved_balance": billing_key.reserved_balance,
"total_spent": billing_key.total_spent,
"finalize_type": "ehbp_max_cost",
"finalized_at": now,
"max_cost_for_model": max_cost_for_model,
},
)
return total_cost_msats
return 0
async def send_cashu_refund(
@@ -846,8 +793,8 @@ async def forward_ehbp_request(
cost_data["computed_msats"] = computed_msats
else:
logger.warning(
"EHBP usage metrics not found in headers or trailers, "
"falling back to max-cost billing",
"EHBP usage metrics not found in headers or trailers; "
"releasing instead of charging the authorization ceiling",
extra={
"model": model_obj.id,
"provider": provider_type,
@@ -868,11 +815,9 @@ async def forward_ehbp_request(
"input_tokens": 0,
"output_tokens": 0,
}
if charged_msats != max_cost_for_model:
cost_data["computed_msats"] = max_cost_for_model
# Build the cost_info dict from what adjust_payment_for_tokens returned
# or from the max-cost fallback. Fields match CostData/MaxCostData.dict().
# Build the cost_info dict from measured usage or the unmeasured-release
# fallback. Fields match CostData/MaxCostData.dict().
cost_info = {
"total_msats": cost_data.get("total_msats", max_cost_for_model),
"input_tokens": cost_data.get("input_tokens", 0),
tests/integration/test_free_response_stale_reservation.py
+10 -8
View File
@@ -92,12 +92,15 @@ async def test_overrun_with_corrupted_aggregate_releases_without_charging(
@pytest.mark.asyncio
async def test_missing_usage_settles_at_reservation_not_zero(
async def test_missing_usage_never_turns_reservation_into_charge(
integration_session: AsyncSession,
) -> None:
"""A response with no usable usage data must settle at the reserved max
cost (bounded fallback), never at zero — otherwise the request is free
inference. Exercises the REAL calculate_cost, no patching."""
"""A reservation is an authorization ceiling, not evidence of usage.
Upstream handlers should provide locally estimated usage when possible. If
no measurement or estimate reaches settlement, release the reservation
rather than charging its full value.
"""
from routstr.auth import (
adjust_payment_for_tokens,
get_reservation_snapshot,
@@ -122,13 +125,12 @@ async def test_missing_usage_settles_at_reservation_not_zero(
reservation_snapshot=reservation,
)
# Charged the authorized max, not zero.
assert result["charged_msats"] == reserved
assert result["charged_msats"] == 0
integration_session.expunge_all()
key_row = await integration_session.get(ApiKey, key_hash)
assert key_row is not None
assert key_row.total_spent == reserved, "missing usage must not be free"
assert key_row.balance == 10_000 - reserved
assert key_row.total_spent == 0
assert key_row.balance == 10_000
assert key_row.reserved_balance == 0
tests/unit/test_count_tokens_local.py
+89 -1
View File
@@ -13,7 +13,7 @@ from unittest.mock import patch
from routstr.payment.models import Architecture, Model, Pricing
from routstr.upstream import count_tokens as count_tokens_module
from routstr.upstream.count_tokens import count_tokens_locally
from routstr.upstream.count_tokens import MissingUsageEstimator, count_tokens_locally
def _make_model(model_id: str = "anthropic/claude-3-5-sonnet") -> Model:
@@ -154,6 +154,94 @@ def test_supports_anthropic_system_block_list() -> None:
assert payload["input_tokens"] > 0
def test_missing_usage_estimator_prices_request_and_streamed_output() -> None:
model = _make_model()
request_body = _body(
{
"model": model.id,
"messages": [{"role": "user", "content": "price this prompt"}],
}
)
with (
patch.object(count_tokens_module, "_count_with_litellm", return_value=17),
patch.object(count_tokens_module, "_count_text_with_litellm", return_value=5),
):
estimator = MissingUsageEstimator(request_body, model)
estimator.observe(
{
"model": "provider/model",
"choices": [{"delta": {"content": "estimated output"}}],
}
)
response = estimator.response_data("provider/model")
assert response == {
"model": "provider/model",
"usage": {
"input_tokens": 17,
"output_tokens": 5,
"total_tokens": 22,
"estimated": True,
},
}
def test_missing_usage_estimator_skips_responses_api_done_events() -> None:
estimator = MissingUsageEstimator(b"{}", None)
estimator.observe({"type": "response.output_text.delta", "delta": "streamed"})
estimator.observe({"type": "response.output_text.done", "text": "streamed"})
estimator.observe(
{
"type": "response.content_part.done",
"part": {"type": "output_text", "text": "streamed"},
}
)
assert estimator.output_text == "streamed"
def test_missing_usage_estimator_openai_dialect() -> None:
model = _make_model()
request_body = _body(
{
"model": model.id,
"messages": [{"role": "user", "content": "price this prompt"}],
}
)
with (
patch.object(count_tokens_module, "_count_with_litellm", return_value=17),
patch.object(count_tokens_module, "_count_text_with_litellm", return_value=5),
):
estimator = MissingUsageEstimator(request_body, model)
estimator.observe({"choices": [{"delta": {"content": "estimated output"}}]})
response = estimator.openai_response_data("provider/model")
assert response == {
"model": "provider/model",
"usage": {
"prompt_tokens": 17,
"completion_tokens": 5,
"total_tokens": 22,
"estimated": True,
},
}
def test_missing_usage_estimator_does_not_count_response_metadata() -> None:
estimator = MissingUsageEstimator(b"{}", None)
estimator.observe(
{
"id": "chatcmpl-this-is-not-generated-text",
"model": "also-not-generated-text",
"choices": [{"delta": {"role": "assistant"}}],
}
)
assert estimator.output_text == ""
def test_uses_forwarded_model_id_when_present() -> None:
model = _make_model("anthropic/claude-3-5-sonnet")
model.forwarded_model_id = "claude-3-5-sonnet-20241022"
tests/unit/test_ehbp_finalize_payment.py
+10 -8
View File
@@ -107,7 +107,7 @@ async def test_finalize_actual_cost_payment_updates_balance_and_releases_reserve
@pytest.mark.asyncio
async def test_finalize_max_cost_payment_updates_parent_and_child_spend(
async def test_unmeasured_ehbp_releases_parent_and_child_reservation(
session: AsyncSession,
) -> None:
parent = ApiKey(hashed_key="ehbp-parent", balance=10_000)
@@ -126,19 +126,19 @@ async def test_finalize_max_cost_payment_updates_parent_and_child_spend(
reservation_snapshot=reservation,
)
assert charged == 3_000
assert charged == 0
updated_parent = await _api_key(session, "ehbp-parent")
updated_child = await _api_key(session, "ehbp-child")
assert updated_parent is not None
assert updated_child is not None
assert updated_parent.balance == 7_000
assert updated_parent.balance == 10_000
assert updated_parent.reserved_balance == 0
assert updated_parent.reserved_at is None
assert updated_parent.total_spent == 3_000
assert updated_parent.total_spent == 0
assert updated_child.balance == 0
assert updated_child.reserved_balance == 0
assert updated_child.reserved_at is None
assert updated_child.total_spent == 3_000
assert updated_child.total_spent == 0
@pytest.mark.asyncio
@@ -178,7 +178,7 @@ async def test_finalize_actual_cost_payment_rolls_back_when_parent_update_matche
@pytest.mark.asyncio
async def test_finalize_max_cost_payment_rolls_back_parent_when_child_update_matches_no_rows(
async def test_unmeasured_ehbp_release_is_safe_when_charge_update_would_fail(
session: AsyncSession,
monkeypatch: pytest.MonkeyPatch,
) -> None:
@@ -207,11 +207,13 @@ async def test_finalize_max_cost_payment_rolls_back_parent_when_child_update_mat
updated_parent = await _api_key(session, "ehbp-rollback-parent")
assert updated_parent is not None
assert updated_parent.balance == 10_000
assert updated_parent.reserved_balance == 0
# The injected partial-update failure rolls aggregate subtraction back;
# terminal fencing prevents a charge or retry from consuming those funds.
assert updated_parent.reserved_balance == 3_000
assert updated_parent.total_spent == 0
updated_child = await _api_key(session, "ehbp-missing-child")
assert updated_child is not None
assert updated_child.reserved_balance == 0
assert updated_child.reserved_balance == 3_000
assert updated_child.total_spent == 0
release = await session.get(ReservationRelease, reservation.release_id)
assert release is not None and release.status == "released"
tests/unit/test_pricing_rate_validation.py
+6 -8
View File
@@ -74,19 +74,17 @@ def _usage_response() -> dict[str, Any]:
ids=["nan", "inf", "negative"],
)
@pytest.mark.asyncio
async def test_unusable_token_rate_falls_back_to_max_cost(bad_rate: float) -> None:
"""An unusable configured rate must not be billed on.
It reached the token math, which raises after the response was already
served — where the streaming handlers swallow it and the request goes
unbilled.
"""
async def test_unusable_token_rate_never_charges_the_reservation(
bad_rate: float,
) -> None:
"""An unusable configured rate must not turn authorization into usage."""
model = _model(Pricing(prompt=bad_rate, completion=1.0))
cost = await calculate_cost(_usage_response(), max_cost=1234, model_obj=model)
assert isinstance(cost, MaxCostData)
assert cost.total_msats == 1234
assert cost.total_msats == 0
assert (cost.input_tokens, cost.output_tokens) == (1000, 500)
@pytest.mark.parametrize(
tests/unit/test_streaming_billing_finalization.py
+42 -12
View File
@@ -22,6 +22,7 @@ from routstr.auth import (
)
from routstr.core.db import ApiKey, ReservationRelease
from routstr.payment.cost_calculation import MaxCostData
from routstr.payment.models import Architecture, Model, Pricing
from routstr.upstream.base import BaseUpstreamProvider
@@ -671,16 +672,14 @@ async def test_cross_key_reservation_snapshot_is_rejected_without_mutation() ->
@pytest.mark.asyncio
async def test_client_disconnect_midstream_finalizes_and_stops_heartbeat() -> None:
"""A client that aborts the socket mid-stream must not leak its reservation.
async def test_client_disconnect_midstream_estimates_usage_and_stops_heartbeat() -> (
None
):
"""A client abort releases the hold after charging only estimated usage.
Starlette closes the response generator (``aclose``) on disconnect, whose
``finally`` schedules the background finalizer. That finalizer must settle
the reservation (charge the reserved max — usage is unknown), reach a
terminal durable state, and stop the lease heartbeat so the sweeper is not
needed. Driven against a real engine and the real finalizer; the socket
abort is modelled deterministically with ``aclose`` (the exact hook
Starlette invokes) to keep the test CI-stable.
Starlette closes the response generator (``aclose``) on disconnect. The
finalizer still has the request and streamed deltas, so it can estimate
usage without converting the reservation ceiling into the charge.
"""
engine = await _engine()
provider = BaseUpstreamProvider(
@@ -707,6 +706,26 @@ async def test_client_disconnect_midstream_finalizes_and_stops_heartbeat() -> No
)
upstream_response.aiter_bytes = aiter_bytes
model = Model(
id="test-model",
name="test-model",
created=0,
description="",
context_length=8_192,
architecture=Architecture(
modality="text",
input_modalities=["text"],
output_modalities=["text"],
tokenizer="unknown",
instruct_type=None,
),
pricing=Pricing(prompt=0.01, completion=0.02),
sats_pricing=Pricing(prompt=0.01, completion=0.02),
)
request_body = json.dumps(
{"model": model.id, "messages": [{"role": "user", "content": "hi"}]}
).encode()
background_tasks = BackgroundTasks()
try:
with (
@@ -718,13 +737,24 @@ async def test_client_disconnect_midstream_finalizes_and_stops_heartbeat() -> No
"routstr.upstream.base.adjust_payment_for_tokens",
auth_module.adjust_payment_for_tokens,
),
patch("routstr.upstream.count_tokens._count_with_litellm", return_value=3),
patch(
"routstr.upstream.count_tokens._count_text_with_litellm",
return_value=2,
),
patch(
"routstr.payment.cost_calculation.sats_usd_price",
return_value=5.0e-5,
),
):
response = await provider.handle_streaming_chat_completion(
response=upstream_response,
key=key,
max_cost_for_model=500,
background_tasks=background_tasks,
model_obj=model,
reservation_snapshot=snapshot,
request_body=request_body,
)
iterator = cast(AsyncGenerator[bytes, None], response.body_iterator)
await iterator.__anext__() # first chunk reaches the client
@@ -744,9 +774,9 @@ async def test_client_disconnect_midstream_finalizes_and_stops_heartbeat() -> No
# The reservation reached a single terminal outcome; funds are not locked.
assert record is not None and record.status in {"charged", "released"}
assert final_key.reserved_balance == 0
# Unknown usage settles at the reserved max, never free.
assert final_key.total_spent == 500
assert final_key.balance == 500
# 3 input tokens × 10 msats + 2 output tokens × 20 msats = 70 msats.
assert final_key.total_spent == 70
assert final_key.balance == 930
# The heartbeat is gone — no forever-renewing task on an abandoned request.
assert snapshot.release_id not in auth_module._reservation_heartbeats
await engine.dispose()
tests/unit/test_tinfoil_integration.py
+4 -4
View File
@@ -240,12 +240,12 @@ class TestResolveEhbpTargetUrl:
class TestComputeEhbpActualCost:
@pytest.mark.asyncio
async def test_no_usage_falls_back_to_max_cost(self) -> None:
async def test_no_usage_does_not_charge_authorization_ceiling(self) -> None:
model_obj = MagicMock()
model_obj.id = "llama3-3-70b"
model_obj.forwarded_model_id = "llama3-3-70b"
result = await _compute_ehbp_actual_cost(None, model_obj, 100_000)
assert result["total_msats"] == 100_000
assert result["total_msats"] == 0
assert result["input_tokens"] == 0
assert result["output_tokens"] == 0
@@ -285,7 +285,7 @@ class TestComputeEhbpActualCost:
assert result["output_msats"] == 20
@pytest.mark.asyncio
async def test_max_cost_data_falls_back(self) -> None:
async def test_unpriceable_usage_does_not_charge_authorization_ceiling(self) -> None:
model_obj = MagicMock()
model_obj.id = "llama3-3-70b"
model_obj.forwarded_model_id = "llama3-3-70b"
@@ -309,7 +309,7 @@ class TestComputeEhbpActualCost:
model_obj,
50_000,
)
assert result["total_msats"] == 50_000
assert result["total_msats"] == 0
assert result["input_tokens"] == 0
assert result["output_tokens"] == 0
tests/unit/test_x_cashu_responses_streaming_sse.py
+4 -4
View File
@@ -186,7 +186,7 @@ async def test_multiline_data_payload_is_parsed_and_reframed() -> None:
@pytest.mark.asyncio
async def test_missing_usage_settles_at_authorized_max() -> None:
async def test_missing_usage_refunds_instead_of_charging_authorized_max() -> None:
chunks = [
b'data: {"type":"response.created","response":{"model":"gpt-5-mini"}}\r\n\r\n',
b"data: [DONE]\r\n\r\n",
@@ -198,9 +198,9 @@ async def test_missing_usage_settles_at_authorized_max() -> None:
send_refund.assert_awaited_once()
assert send_refund.await_args is not None
assert send_refund.await_args.args[0] == 10_000 - 9_000
assert send_refund.await_args.args[0] == 10_000
assert response.headers["x-cashu"] == "cashuBrefundtoken0123456789"
assert response.headers["x-routstr-cost-msats"] == "9000"
assert response.headers["x-routstr-cost-msats"] == "0"
@pytest.mark.asyncio
@@ -215,7 +215,7 @@ async def test_malformed_events_do_not_retain_whole_token() -> None:
)
assert send_refund.await_args is not None
assert send_refund.await_args.args[0] == 1000
assert send_refund.await_args.args[0] == 10_000
body = await _collect(response)
assert b"\\n" not in body
assert body.endswith(b"\n\n")