From 4287f038cf66295ff892f8dfec0c80e02bf9115a Mon Sep 17 00:00:00 2001 From: redshift <213178690+1ftredsh@users.noreply.github.com> Date: Wed, 8 Jul 2026 21:39:33 +0530 Subject: [PATCH] feat(ehbp): extract & use model name from Tinfoil usage metrics header Tinfoil PR #385 added model= to the X-Tinfoil-Usage-Metrics header/trailer. This commit uses that field for accurate billing. Changes in routstr/upstream/ehbp.py: - parse_tinfoil_usage_metrics(): extract the model= field as a string alongside the existing token counts (previously silently discarded because int() failed on it). - _build_cost_info(): accept optional actual_model parameter propagated through to callers when a real mismatch is detected. - _compute_ehbp_actual_cost(): compare the served model against model_obj.forwarded_model_id (the expected upstream ID) rather than model_obj.id (the client-facing alias). This prevents spurious mismatches when a node runner maps e.g. tinfoil-glm-5-2 -> glm-5-2 and the header correctly reports glm-5-2. On a genuine mismatch (failover to a different upstream model), look up the actual model's pricing via get_model_instance() (forwarded_model_id values are registered as routable aliases in the global model map). - forward_ehbp_request() / forward_ehbp_x_cashu_request(): use the actual served model for payment finalization and logging when a mismatch is detected. Tests: 6 new scenarios (alias match, real mismatch with alias, unknown model fallback, old-format compat, cache token details + model), plus forwarded_model_id set on all existing mock model objects to keep them passing. All 49 Tinfoil/EHBP unit tests pass. --- docs/tinfoil-direct-integration.md | 45 ++++- routstr/upstream/ehbp.py | 121 +++++++++++-- tests/unit/test_tinfoil_integration.py | 225 +++++++++++++++++++++++++ 3 files changed, 367 insertions(+), 24 deletions(-) diff --git a/docs/tinfoil-direct-integration.md b/docs/tinfoil-direct-integration.md index ce689203..80d104c1 100644 --- a/docs/tinfoil-direct-integration.md +++ b/docs/tinfoil-direct-integration.md @@ -395,8 +395,10 @@ and `routstr/upstream/ehbp.py`. `TINFOIL_API_KEY` env var. - `routstr/upstream/ehbp.py`: - - `parse_tinfoil_usage_metrics()` parses `prompt=N,completion=N[,total=N]` - into an OpenAI-style usage dict. + - `parse_tinfoil_usage_metrics()` parses + `prompt=N,completion=N[,total=N][,model=]` into an OpenAI-style + usage dict. The `model` field (added in tinfoilsh/confidential-model-router + PR #385) is extracted as a string. - `_resolve_ehbp_target_url()` overrides the forwarding URL with `X-Tinfoil-Enclave-Url` when the SDK sends it. - `_strip_proxy_headers()` removes `X-Routstr-Model`, @@ -404,11 +406,15 @@ and `routstr/upstream/ehbp.py`. forwarding to the enclave. - `_compute_ehbp_actual_cost()` converts the usage header into msats via `calculate_cost()`, clamped to `[min_request_msat, max_cost_for_model]`. + When the header's `model=` differs from the requested model, the + 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. + for exact billing; otherwise falls back to max-cost. 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. + refund from actual cost instead of max cost, using the actual served + model's pricing when applicable. - `routstr/proxy.py`: `/attestation` and `/tee/attestation` paths are forwarded to Tinfoil upstreams without model/cost/auth lookups. @@ -448,10 +454,37 @@ TINFOIL_API_KEY=your-tinfoil-api-key The provider is auto-seeded on first startup. +### Usage metrics header format + +Tinfoil returns usage metrics in the `X-Tinfoil-Usage-Metrics` response header +(non-streaming) or HTTP trailer (streaming) when `X-Tinfoil-Request-Usage-Metrics: +true` is sent. As of tinfoilsh/confidential-model-router PR #385, the format is: + +``` +prompt=,completion=,total=,model= +``` + +The `model` field carries the actual model name served by the enclave. +Routstr uses this to: + +- Verify the served model matches the expected upstream model. The comparison + uses ``model_obj.forwarded_model_id`` (the actual upstream ID, e.g. + ``glm-5-2``) rather than ``model_obj.id`` (the client-facing alias, e.g. + ``tinfoil-glm-5-2``), so aliased models don't trigger a spurious mismatch. +- When they genuinely differ (Tinfoil served a different upstream model than + expected), look up the actual served model's pricing and use it for billing. + The reverse lookup uses ``get_model_instance``, which resolves + ``forwarded_model_id`` values registered as routable aliases. +- Log the discrepancy for observability. + +If the actual model is not found in Routstr's model registry, billing falls +back to the requested model's pricing. + ### What still needs verification -- End-to-end test with a real Tinfoil SDK client against a Routstr node with - `TINFOIL_API_KEY` set. +- ~~End-to-end test with a real Tinfoil SDK client against a Routstr node with + `TINFOIL_API_KEY` set.~~ Verified: both non-streaming (header) and streaming + (trailer) responses include `model=`. - Streaming requests: usage is delivered as an HTTP trailer. Currently the bearer path finalizes max-cost before streaming begins. Supporting streaming usage would require buffering the response (for X-Cashu) or a deferred diff --git a/routstr/upstream/ehbp.py b/routstr/upstream/ehbp.py index e5c073d3..3f643a74 100644 --- a/routstr/upstream/ehbp.py +++ b/routstr/upstream/ehbp.py @@ -63,30 +63,48 @@ _PROXY_ONLY_HEADERS = frozenset( def parse_tinfoil_usage_metrics(header_value: str | None) -> dict | None: """Parse ``X-Tinfoil-Usage-Metrics`` into an OpenAI-style usage dict. - The header format is ``prompt=,completion=,total=``. Returns a dict - like ``{"prompt_tokens": n, "completion_tokens": n}`` suitable for - :func:`calculate_cost`, or ``None`` when the header is absent or malformed. + The header format is:: + + prompt=,completion=,total=[,model=] + + The ``model`` field (added in tinfoilsh/confidential-model-router PR #385) + is extracted as a string and included in the returned dict under the + ``"model"`` key so callers can compare the served model against the + requested one and adjust pricing. + + Returns a dict like ``{"prompt_tokens": n, "completion_tokens": n, + "model": ""}`` suitable for :func:`calculate_cost` (which ignores + the extra ``model`` key in the usage sub-dict), or ``None`` when the + header is absent or malformed. """ if not header_value: return None parts: dict[str, int] = {} + model: str | None = None for item in header_value.split(","): key, sep, value = item.partition("=") if not sep: continue + key = key.strip() + value = value.strip() + if key == "model": + model = value + continue try: - parts[key.strip()] = int(value.strip()) + parts[key] = int(value) except (ValueError, TypeError): continue prompt = parts.get("prompt") completion = parts.get("completion") if prompt is not None and completion is not None: - result: dict[str, int] = { + result: dict[str, int | str] = { "prompt_tokens": prompt, "completion_tokens": completion, } if "total" in parts: result["total_tokens"] = parts["total"] + if model: + result["model"] = model return result logger.warning( "Failed to parse X-Tinfoil-Usage-Metrics header", @@ -242,9 +260,15 @@ def _build_cost_info( output_tokens: int = 0, input_msats: int = 0, output_msats: int = 0, + actual_model: str | None = None, ) -> dict: - """Build a cost-info dict with token counts and per-token-type costs.""" - return { + """Build a cost-info dict with token counts and per-token-type costs. + + When ``actual_model`` is set (the served model differs from the requested + one), it is included in the returned dict so callers can use it for billing + finalization and logging. + """ + result: dict[str, int | str | None] = { "total_msats": total_msats, "input_tokens": input_tokens, "output_tokens": output_tokens, @@ -252,6 +276,9 @@ def _build_cost_info( "input_msats": input_msats, "output_msats": output_msats, } + if actual_model: + result["actual_model"] = actual_model + return result def _inject_cost_response_headers( @@ -280,48 +307,98 @@ async def _compute_ehbp_actual_cost( max_cost_for_model]`` so the refund never exceeds the reservation and is never zero. + When the usage-metrics header includes ``model=`` and it differs + from ``model_obj.id``, the actual served model's pricing is used for the + cost calculation. The returned dict includes an ``"actual_model"`` key + in that case so callers can use it for billing finalization. + Returns a dict with ``total_msats``, ``input_tokens``, ``output_tokens``, - ``total_tokens``, ``input_msats``, and ``output_msats``. + ``total_tokens``, ``input_msats``, and ``output_msats`` (and optionally + ``actual_model``). """ usage_dict = parse_tinfoil_usage_metrics(usage_header) if usage_dict is None: return _build_cost_info(max_cost_for_model) + # The enclave may serve a different model than the one requested (e.g. + # due to failover). The usage-metrics header's ``model=`` carries + # the actual upstream model ID (e.g. ``glm-5-2``), which may differ from + # the client-facing ``model_obj.id`` (e.g. ``tinfoil-glm-5-2``) even when + # the correct model was served — the alias is resolved through + # ``model_obj.forwarded_model_id``. Only when the served model differs + # from the expected upstream ID do we treat it as a real mismatch and + # look up the actual model's pricing. + actual_model: str | None = usage_dict.pop("model", None) # type: ignore[arg-type] + pricing_model_id = model_obj.id + expected_upstream_model = model_obj.forwarded_model_id or model_obj.id + if actual_model and actual_model != expected_upstream_model: + from ..proxy import get_model_instance + + # ``forwarded_model_id`` values are registered as routable aliases in + # the global model map, so ``get_model_instance`` will find a model + # whose upstream ID matches the actually-served model. + actual_model_obj = get_model_instance(actual_model) + if actual_model_obj: + logger.info( + "EHBP served model differs from requested, using actual " + "model for pricing", + extra={ + "requested_model": model_obj.id, + "expected_upstream_model": expected_upstream_model, + "actual_model": actual_model, + }, + ) + pricing_model_id = actual_model_obj.id + else: + logger.warning( + "EHBP served model not found in registry, falling back to " + "requested model for pricing", + extra={ + "requested_model": model_obj.id, + "expected_upstream_model": expected_upstream_model, + "actual_model": actual_model, + }, + ) + actual_model = None # do not propagate unknown model + else: + # Models match or no model in header — use requested model's pricing. + actual_model = None + try: cost = await calculate_cost( - {"model": model_obj.id, "usage": usage_dict}, + {"model": pricing_model_id, "usage": usage_dict}, max_cost_for_model, ) except Exception as e: logger.warning( "EHBP usage cost calculation failed, falling back to max cost", extra={ - "model": model_obj.id, + "model": pricing_model_id, "error": str(e), "usage": usage_dict, }, ) - return _build_cost_info(max_cost_for_model) + return _build_cost_info(max_cost_for_model, actual_model=actual_model) if isinstance(cost, MaxCostData): logger.warning( "EHBP calculate_cost returned MaxCostData (no model pricing), " "falling back to max cost", extra={ - "model": model_obj.id, + "model": pricing_model_id, "max_cost_for_model": max_cost_for_model, "usage": usage_dict, "cost_total_msats": cost.total_msats, }, ) - return _build_cost_info(max_cost_for_model) + return _build_cost_info(max_cost_for_model, 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) logger.info( "EHBP actual cost computed from usage metrics", extra={ - "model": model_obj.id, + "model": pricing_model_id, "usage": usage_dict, "cost_total_msats": cost.total_msats, "clamped_msats": clamped, @@ -334,16 +411,17 @@ async def _compute_ehbp_actual_cost( output_tokens=cost.output_tokens, input_msats=cost.input_msats, output_msats=cost.output_msats, + actual_model=actual_model, ) # CostDataError logger.warning( "EHBP usage cost calculation error, falling back to max cost", extra={ - "model": model_obj.id, + "model": pricing_model_id, "error": getattr(cost, "message", str(cost)), }, ) - return _build_cost_info(max_cost_for_model) + return _build_cost_info(max_cost_for_model, actual_model=actual_model) def _extract_usage_from_response( @@ -771,8 +849,11 @@ async def forward_ehbp_request( cost_info = await _compute_ehbp_actual_cost( usage_header, model_obj, max_cost_for_model ) + # Use the actual served model for billing when it differs from + # the requested model. + billing_model = cost_info.pop("actual_model", None) or model_obj.id await finalize_ehbp_actual_cost_payment( - key, session, max_cost_for_model, model_obj.id, cost_info + key, session, max_cost_for_model, billing_model, cost_info ) cost_data = {**cost_info, "total_usd": 0.0} else: @@ -972,11 +1053,15 @@ async def forward_ehbp_x_cashu_request( usage_header, model_obj, max_cost_for_model ) actual_cost_msats = cost_info["total_msats"] + actual_model = cost_info.get("actual_model") + billing_model = actual_model or model_obj.id refund_amount = amount - _msats_to_unit_amount(actual_cost_msats, unit) logger.info( "EHBP X-Cashu refund computed", extra={ - "model": model_obj.id, + "model": billing_model, + "requested_model": model_obj.id, + "actual_model": actual_model, "redeemed_amount": amount, "actual_cost_msats": actual_cost_msats, "refund_amount": refund_amount, diff --git a/tests/unit/test_tinfoil_integration.py b/tests/unit/test_tinfoil_integration.py index 7183b49a..3548b4f3 100644 --- a/tests/unit/test_tinfoil_integration.py +++ b/tests/unit/test_tinfoil_integration.py @@ -64,6 +64,56 @@ class TestParseTinfoilUsageMetrics: "total_tokens": 300, } + def test_with_model_field(self) -> None: + result = parse_tinfoil_usage_metrics( + "prompt=42,completion=10,total=52,model=llama3-3-70b" + ) + assert result == { + "prompt_tokens": 42, + "completion_tokens": 10, + "total_tokens": 52, + "model": "llama3-3-70b", + } + + def test_with_model_no_total(self) -> None: + result = parse_tinfoil_usage_metrics( + "prompt=67,completion=42,model=gpt-oss-120b" + ) + assert result == { + "prompt_tokens": 67, + "completion_tokens": 42, + "model": "gpt-oss-120b", + } + + def test_model_with_dashes_and_numbers(self) -> None: + result = parse_tinfoil_usage_metrics( + "prompt=1,completion=1,total=2,model=kimi-k2-6" + ) + assert result["model"] == "kimi-k2-6" + + def test_model_with_extra_fields(self) -> None: + result = parse_tinfoil_usage_metrics( + "prompt=69,completion=20,total=89," + "cached_prompt_tokens=64,uncached_prompt_tokens=5," + "model=kimi-k2-6" + ) + assert result["prompt_tokens"] == 69 + assert result["completion_tokens"] == 20 + assert result["total_tokens"] == 89 + assert result["model"] == "kimi-k2-6" + + def test_old_format_still_works(self) -> None: + """Headers without the model field (pre-PR #385) still parse.""" + result = parse_tinfoil_usage_metrics( + "prompt=67,completion=42,total=109" + ) + assert result == { + "prompt_tokens": 67, + "completion_tokens": 42, + "total_tokens": 109, + } + assert "model" not in result + # --------------------------------------------------------------------------- # _strip_proxy_headers @@ -191,6 +241,7 @@ class TestComputeEhbpActualCost: async def test_no_usage_falls_back_to_max_cost(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["input_tokens"] == 0 @@ -200,6 +251,7 @@ class TestComputeEhbpActualCost: async def test_usage_parsed_and_clamped(self) -> None: model_obj = MagicMock() model_obj.id = "llama3-3-70b" + model_obj.forwarded_model_id = "llama3-3-70b" # The actual cost from calculate_cost will be small; we just verify # it's clamped to min_request_msat at minimum. with patch( @@ -234,6 +286,7 @@ class TestComputeEhbpActualCost: async def test_max_cost_data_falls_back(self) -> None: model_obj = MagicMock() model_obj.id = "llama3-3-70b" + model_obj.forwarded_model_id = "llama3-3-70b" with patch( "routstr.upstream.ehbp.calculate_cost", new_callable=AsyncMock, @@ -258,6 +311,178 @@ class TestComputeEhbpActualCost: assert result["input_tokens"] == 0 assert result["output_tokens"] == 0 + @pytest.mark.asyncio + async def test_model_match_no_actual_model_key(self) -> None: + """When the served model matches the requested one, no actual_model key.""" + model_obj = MagicMock() + model_obj.id = "llama3-3-70b" + model_obj.forwarded_model_id = "llama3-3-70b" + with patch( + "routstr.upstream.ehbp.calculate_cost", + new_callable=AsyncMock, + ) as mock_calc: + from routstr.payment.cost_calculation import CostData + + mock_calc.return_value = CostData( + base_msats=0, + input_msats=5, + output_msats=10, + total_msats=15, + total_usd=0.0, + input_tokens=42, + output_tokens=10, + ) + result = await _compute_ehbp_actual_cost( + "prompt=42,completion=10,total=52,model=llama3-3-70b", + model_obj, + 100_000, + ) + assert "actual_model" not in result + # calculate_cost called with requested model + call_args = mock_calc.call_args + assert call_args[0][0]["model"] == "llama3-3-70b" + + @pytest.mark.asyncio + async def test_alias_match_no_actual_model_key(self) -> None: + """When the served upstream model matches forwarded_model_id through + a client-facing alias, no actual_model key is set.""" + model_obj = MagicMock() + model_obj.id = "tinfoil-glm-5-2" # client-facing alias + model_obj.forwarded_model_id = "glm-5-2" # actual upstream ID + with patch( + "routstr.upstream.ehbp.calculate_cost", + new_callable=AsyncMock, + ) as mock_calc: + from routstr.payment.cost_calculation import CostData + + mock_calc.return_value = CostData( + base_msats=0, + input_msats=5, + output_msats=10, + total_msats=15, + total_usd=0.0, + input_tokens=42, + output_tokens=10, + ) + # Tinfoil header returns the actual upstream model ID + result = await _compute_ehbp_actual_cost( + "prompt=42,completion=10,total=52,model=glm-5-2", + model_obj, + 100_000, + ) + assert "actual_model" not in result + # calculate_cost called with the client-facing model ID (whose + # pricing includes the correct upstream rates) + call_args = mock_calc.call_args + assert call_args[0][0]["model"] == "tinfoil-glm-5-2" + + @pytest.mark.asyncio + async def test_real_mismatch_uses_actual_model_for_pricing(self) -> None: + """When the served model differs from the expected upstream model, + the actual model's pricing is used.""" + model_obj = MagicMock() + model_obj.id = "tinfoil-gpt-oss-120b" # client-facing alias + model_obj.forwarded_model_id = "gpt-oss-120b" # expected upstream + + actual_model_obj = MagicMock() + actual_model_obj.id = "tinfoil-llama3-3-70b" # client-facing of actual + actual_model_obj.forwarded_model_id = "llama3-3-70b" + + with patch( + "routstr.proxy.get_model_instance", + return_value=actual_model_obj, + ), patch( + "routstr.upstream.ehbp.calculate_cost", + new_callable=AsyncMock, + ) as mock_calc: + from routstr.payment.cost_calculation import CostData + + mock_calc.return_value = CostData( + base_msats=0, + input_msats=20, + output_msats=40, + total_msats=60, + total_usd=0.0, + input_tokens=42, + output_tokens=10, + ) + # Tinfoil served llama3-3-70b instead of gpt-oss-120b + result = await _compute_ehbp_actual_cost( + "prompt=42,completion=10,total=52,model=llama3-3-70b", + model_obj, + 100_000, + ) + assert result["actual_model"] == "llama3-3-70b" + assert result["total_msats"] == 60 + # calculate_cost called with the actual model's client-facing ID + call_args = mock_calc.call_args + assert call_args[0][0]["model"] == "tinfoil-llama3-3-70b" + + @pytest.mark.asyncio + async def test_model_mismatch_unknown_model_falls_back(self) -> None: + """When the served model is not in the registry, use requested model.""" + model_obj = MagicMock() + model_obj.id = "gpt-oss-120b" + model_obj.forwarded_model_id = "gpt-oss-120b" + + with patch( + "routstr.proxy.get_model_instance", + return_value=None, + ), patch( + "routstr.upstream.ehbp.calculate_cost", + new_callable=AsyncMock, + ) as mock_calc: + from routstr.payment.cost_calculation import CostData + + mock_calc.return_value = CostData( + base_msats=0, + input_msats=5, + output_msats=10, + total_msats=15, + total_usd=0.0, + input_tokens=42, + output_tokens=10, + ) + result = await _compute_ehbp_actual_cost( + "prompt=42,completion=10,total=52,model=nonexistent", + model_obj, + 100_000, + ) + assert "actual_model" not in result + # calculate_cost called with the requested model (fallback) + call_args = mock_calc.call_args + assert call_args[0][0]["model"] == "gpt-oss-120b" + + @pytest.mark.asyncio + async def test_old_format_no_model_uses_requested(self) -> None: + """Old format without model field uses requested model for pricing.""" + model_obj = MagicMock() + model_obj.id = "llama3-3-70b" + model_obj.forwarded_model_id = "llama3-3-70b" + with patch( + "routstr.upstream.ehbp.calculate_cost", + new_callable=AsyncMock, + ) as mock_calc: + from routstr.payment.cost_calculation import CostData + + mock_calc.return_value = CostData( + base_msats=0, + input_msats=5, + output_msats=10, + total_msats=15, + total_usd=0.0, + input_tokens=67, + output_tokens=42, + ) + result = await _compute_ehbp_actual_cost( + "prompt=67,completion=42,total=109", + model_obj, + 100_000, + ) + assert "actual_model" not in result + call_args = mock_calc.call_args + assert call_args[0][0]["model"] == "llama3-3-70b" + # --------------------------------------------------------------------------- # TinfoilUpstreamProvider