Merge branch 'main' into with-price-calculation

routstr/payment/cost_calculation.py

@@ -111,6 +111,13 @@ async def calculate_cost(
     if not isinstance(usage_data, dict):
         usage_data = {}
 
+    logger.warning(
+        "Upstream response cost and usage: model=%s cost=%s usage=%s",
+        response_data.get("model", "unknown"),
+        usage_data.get("cost"),
+        usage_data,
+    )
+
     input_tokens = usage.input_tokens
     output_tokens = usage.output_tokens
     cache_read_tokens = usage.cache_read_tokens
@@ -297,6 +304,7 @@ def _get_pricing_rates(
     ):
         return None
 
+    from .models import litellm_cost_entry
     from ..proxy import get_model_instance
     from .models import litellm_cost_entry
 

routstr/upstream/base.py

@@ -274,6 +274,39 @@ class BaseUpstreamProvider:
             total_usd = cost_data.total_usd
             cost_dict = cost_data.dict()
 
+        input_msats = int(cost_dict.get("input_msats", 0) or 0)
+        output_msats = int(cost_dict.get("output_msats", 0) or 0)
+        input_tokens = int(cost_dict.get("input_tokens", 0) or 0)
+        output_tokens = int(cost_dict.get("output_tokens", 0) or 0)
+        cache_tokens = int(cost_dict.get("cache_read_input_tokens", 0) or 0) + int(
+            cost_dict.get("cache_creation_input_tokens", 0) or 0
+        )
+        priced_tokens = input_tokens + cache_tokens + output_tokens
+        if (
+            total_msats > 0
+            and input_msats == 0
+            and output_msats == 0
+            and priced_tokens > 0
+        ):
+            logger.warning(
+                "Repairing missing client cost breakdown before response: "
+                "model=%s cost=%s",
+                response_json.get("model", "unknown"),
+                cost_dict,
+            )
+            output_msats = total_msats * output_tokens // priced_tokens
+            input_msats = total_msats - output_msats
+            cost_dict = cost_dict.copy()
+            cost_dict["base_msats"] = 0
+            cost_dict["input_msats"] = input_msats
+            cost_dict["output_msats"] = output_msats
+
+        logger.warning(
+            "Client-facing cost metadata: model=%s cost=%s",
+            response_json.get("model", "unknown"),
+            cost_dict,
+        )
+
         sats_cost = total_msats // 1000
 
         # Inject into top-level usage block (OpenAI/Anthropic style)
@@ -820,6 +853,48 @@ class BaseUpstreamProvider:
             last_model_seen: str | None = None
             usage_chunk_data: dict | None = None
             done_seen: bool = False
+            pending_finish_event: tuple[bytes, dict] | None = None
+            observed_event_shapes: set[tuple] = set()
+
+            def warn_event_shape(
+                raw_event: bytes,
+                obj: dict | None = None,
+            ) -> None:
+                """Log each distinct SSE framing/semantic shape once per stream."""
+                fields = tuple(
+                    line.split(b":", 1)[0].decode("ascii", errors="replace")
+                    for line in raw_event.split(b"\n")
+                    if line
+                )
+                choices = obj.get("choices") if isinstance(obj, dict) else None
+                finish_reasons = tuple(
+                    choice.get("finish_reason")
+                    for choice in choices or []
+                    if isinstance(choice, dict)
+                    and choice.get("finish_reason") is not None
+                )
+                shape = (
+                    fields,
+                    tuple(sorted(obj.keys())) if isinstance(obj, dict) else (),
+                    isinstance(obj.get("usage"), dict)
+                    if isinstance(obj, dict)
+                    else False,
+                    len(choices) if isinstance(choices, list) else None,
+                    finish_reasons,
+                )
+                if shape in observed_event_shapes or len(observed_event_shapes) >= 12:
+                    return
+                observed_event_shapes.add(shape)
+                logger.warning(
+                    "Streaming SSE event shape: model=%s fields=%s "
+                    "json_keys=%s has_usage=%s choices=%s finish_reasons=%s",
+                    (obj or {}).get("model", last_model_seen or "unknown"),
+                    fields,
+                    shape[1],
+                    shape[2],
+                    shape[3],
+                    finish_reasons,
+                )
 
             async def finalize_db_only() -> None:
                 nonlocal usage_finalized
@@ -862,6 +937,7 @@ class BaseUpstreamProvider:
                   end of stream.
                 """
                 nonlocal last_model_seen, usage_chunk_data, done_seen
+                nonlocal pending_finish_event
 
                 event = raw_event.strip(b"\r\n")
                 if not event:
@@ -882,6 +958,7 @@ class BaseUpstreamProvider:
                         field_lines.append(line)
 
                 if not data_lines:
+                    warn_event_shape(event)
                     return
 
                 data = b"\n".join(data_lines)
@@ -902,6 +979,8 @@ class BaseUpstreamProvider:
                 except Exception:
                     obj = None
 
+                warn_event_shape(event, obj if isinstance(obj, dict) else None)
+
                 if isinstance(obj, dict):
                     self._apply_provider_field(obj)
                     if obj.get("model"):
@@ -937,16 +1016,36 @@ class BaseUpstreamProvider:
                             # Forward the content now, without usage, so token
                             # usage is reported exactly once (in the trailer).
                             forward = {k: v for k, v in obj.items() if k != "usage"}
-                            yield (
+                            encoded = (
                                 prefix
                                 + b"data: "
                                 + json.dumps(forward).encode()
                                 + b"\n\n"
                             )
+                            if any(
+                                choice.get("finish_reason") is not None
+                                for choice in obj.get("choices", [])
+                                if isinstance(choice, dict)
+                            ):
+                                pending_finish_event = (prefix, forward)
+                            else:
+                                yield encoded
                             return
                         usage_chunk_data = obj
                         return
-                    yield prefix + b"data: " + json.dumps(obj).encode() + b"\n\n"
+                    encoded = prefix + b"data: " + json.dumps(obj).encode() + b"\n\n"
+                    if any(
+                        choice.get("finish_reason") is not None
+                        for choice in obj.get("choices", [])
+                        if isinstance(choice, dict)
+                    ):
+                        # Some consumers stop reading immediately at
+                        # finish_reason. Hold it until payment finalization and
+                        # attach cost metadata to this same event so clients do
+                        # not need to consume a later proprietary trailer.
+                        pending_finish_event = (prefix, obj)
+                    else:
+                        yield encoded
                 else:
                     if final:
                         # Final flush of a truncated tail: the upstream closed
@@ -989,6 +1088,7 @@ class BaseUpstreamProvider:
                     for out in _process_event(buffer, final=True):
                         yield out
 
+                cost_trailer: bytes | None = None
                 async with create_session() as session:
                     fresh_key = await session.get(key.__class__, key.hashed_key)
                     if fresh_key:
@@ -1002,11 +1102,39 @@ class BaseUpstreamProvider:
                                     "usage": None,
                                 }
                             )
-                            cost_data = await adjust_payment_for_tokens(
-                                fresh_key,
-                                adjustment_input,
-                                session,
-                                max_cost_for_model,
+                            logger.warning(
+                                "Streaming cost finalization starting: model=%s "
+                                "has_usage=%s",
+                                adjustment_input.get("model", "unknown"),
+                                isinstance(adjustment_input.get("usage"), dict),
+                            )
+                            finalization_task = asyncio.create_task(
+                                adjust_payment_for_tokens(
+                                    fresh_key,
+                                    adjustment_input,
+                                    session,
+                                    max_cost_for_model,
+                                )
+                            )
+                            try:
+                                cost_data = await asyncio.shield(finalization_task)
+                            except asyncio.CancelledError:
+                                # The upstream response is already complete, so
+                                # abandoning payment finalization here loses the
+                                # client cost trailer and leaves the reservation
+                                # unsettled. Finish the shielded task before
+                                # allowing stream teardown to continue.
+                                logger.warning(
+                                    "Streaming cost finalization cancellation "
+                                    "deferred: model=%s",
+                                    adjustment_input.get("model", "unknown"),
+                                )
+                                cost_data = await finalization_task
+                            logger.warning(
+                                "Streaming cost finalization returned: model=%s "
+                                "cost=%s",
+                                adjustment_input.get("model", "unknown"),
+                                cost_data,
                             )
                             usage_finalized = True
                         except Exception as e:
@@ -1018,16 +1146,36 @@ class BaseUpstreamProvider:
                                 },
                             )
 
-                            # Fall back so we still emit a non-zero sats cost downstream.
-                            cost_data = {
-                                "base_msats": 0,
-                                "input_msats": 0,
-                                "output_msats": 0,
-                                "total_msats": 0,
-                                "total_usd": 0.0,
-                                "input_tokens": 0,
-                                "output_tokens": 0,
-                            }
+                            # Payment settlement can fail after calculate_cost()
+                            # has already produced a valid breakdown (for example
+                            # because of a concurrent reservation/database race).
+                            # Never replace that client-visible calculation with
+                            # the old all-zero fallback. Recalculate without DB
+                            # mutation so the stream still carries accurate cost
+                            # metadata while the exception remains fully logged.
+                            fallback_cost = await calculate_cost(
+                                adjustment_input,
+                                max_cost_for_model,
+                            )
+                            if isinstance(fallback_cost, CostDataError):
+                                cost_data = {
+                                    "base_msats": max_cost_for_model,
+                                    "input_msats": 0,
+                                    "output_msats": 0,
+                                    "total_msats": max_cost_for_model,
+                                    "total_usd": 0.0,
+                                    "input_tokens": 0,
+                                    "output_tokens": 0,
+                                }
+                            else:
+                                cost_data = fallback_cost.dict()
+                            logger.warning(
+                                "Streaming payment finalization failed; using "
+                                "calculated client cost: model=%s error=%s cost=%s",
+                                adjustment_input.get("model", "unknown"),
+                                str(e),
+                                cost_data,
+                            )
 
                         if usage_chunk_data is None:
                             if not hasattr(self, "_current_stream_id"):
@@ -1065,8 +1213,30 @@ class BaseUpstreamProvider:
                                 },
                             )
 
-                        yield f"data: {json.dumps(usage_chunk_data)}\n\n".encode()
+                        cost_trailer = (
+                            f"data: {json.dumps(usage_chunk_data)}\n\n".encode()
+                        )
+                    else:
+                        logger.warning(
+                            "Streaming cost finalization skipped: API key not found; "
+                            "model=%s",
+                            last_model_seen or "unknown",
+                        )
 
+                if pending_finish_event is not None:
+                    finish_prefix, finish_obj = pending_finish_event
+                    if usage_chunk_data is not None:
+                        for field in ("provider", "metadata", "cost"):
+                            if field in usage_chunk_data:
+                                finish_obj[field] = usage_chunk_data[field]
+                    yield (
+                        finish_prefix
+                        + b"data: "
+                        + json.dumps(finish_obj).encode()
+                        + b"\n\n"
+                    )
+                if cost_trailer is not None:
+                    yield cost_trailer
                 if done_seen:
                     yield b"data: [DONE]\n\n"
 
@@ -1152,31 +1322,12 @@ class BaseUpstreamProvider:
             )
 
             await session.refresh(key)
-            remaining_balance_msats = key.balance
 
-            # Merge cost into usage for OpenCode
-            if "usage" in response_json:
-                response_json["usage"]["cost"] = cost_data.get("total_usd", 0.0)
-                response_json["usage"]["cost_sats"] = (
-                    cost_data.get("total_msats", 0) // 1000
-                )
-                response_json["usage"]["remaining_balance_msats"] = (
-                    remaining_balance_msats
-                )
-                self._fold_cache_into_input_tokens(response_json["usage"])
-
-            # Keep detailed cost
-            response_json["metadata"] = response_json.get("metadata", {})
-            response_json["metadata"]["routstr"] = {"cost": cost_data}
-            response_json["metadata"]["routstr"]["cost"]["sats_cost"] = (
-                cost_data.get("total_msats", 0) // 1000
-            )
-            response_json["metadata"]["routstr"]["cost"]["remaining_balance_msats"] = (
-                remaining_balance_msats
-            )
-            response_json["cost"] = cost_data
-            response_json["cost"]["sats_cost"] = cost_data.get("total_msats", 0) // 1000
-            response_json["cost"]["remaining_balance_msats"] = remaining_balance_msats
+            # Use the same metadata path as streaming responses. The previous
+            # non-streaming-only serializer bypassed component repair and the
+            # client-facing diagnostic log, allowing a valid total to be sent
+            # with zero input/output costs.
+            self.inject_cost_metadata(response_json, cost_data, key)
 
             logger.debug(
                 "Payment adjustment completed for non-streaming",
@@ -1209,6 +1360,18 @@ class BaseUpstreamProvider:
 
             if requested_model:
                 response_json["model"] = requested_model
+            logger.warning(
+                "Non-streaming response final cost before serialization: "
+                "model=%s top_level=%s metadata=%s usage_cost=%s",
+                response_json.get("model", "unknown"),
+                response_json.get("cost"),
+                response_json.get("metadata", {})
+                .get("routstr", {})
+                .get("cost"),
+                response_json.get("usage", {}).get("cost")
+                if isinstance(response_json.get("usage"), dict)
+                else None,
+            )
             return Response(
                 content=json.dumps(response_json).encode(),
                 status_code=response.status_code,
@@ -2548,6 +2711,64 @@ class BaseUpstreamProvider:
 
         transformed_body = self.prepare_request_body(request_body, model_obj)
 
+        def request_shape(body: bytes | None) -> dict[str, Any]:
+            """Return diagnostic request metadata without prompt/tool contents."""
+            if not body:
+                return {"body_present": False}
+            try:
+                parsed = json.loads(body)
+            except (json.JSONDecodeError, UnicodeDecodeError):
+                return {"body_present": True, "valid_json": False}
+            if not isinstance(parsed, dict):
+                return {
+                    "body_present": True,
+                    "valid_json": True,
+                    "json_type": type(parsed).__name__,
+                }
+            messages = parsed.get("messages")
+            roles = (
+                [
+                    message.get("role", "unknown")
+                    for message in messages
+                    if isinstance(message, dict)
+                ]
+                if isinstance(messages, list)
+                else []
+            )
+            tools = parsed.get("tools")
+            return {
+                "body_present": True,
+                "valid_json": True,
+                "keys": sorted(parsed.keys()),
+                "model": parsed.get("model", "unknown"),
+                "stream": parsed.get("stream", False),
+                "stream_options": parsed.get("stream_options"),
+                "message_count": len(messages)
+                if isinstance(messages, list)
+                else 0,
+                "message_roles": roles,
+                "tool_count": len(tools) if isinstance(tools, list) else 0,
+                "tool_choice": parsed.get("tool_choice"),
+                "max_tokens": parsed.get("max_tokens"),
+                "max_completion_tokens": parsed.get("max_completion_tokens"),
+                "reasoning_effort": parsed.get("reasoning_effort"),
+                "thinking_type": (
+                    parsed.get("thinking", {}).get("type")
+                    if isinstance(parsed.get("thinking"), dict)
+                    else None
+                ),
+            }
+
+        logger.warning(
+            "Upstream request shape: provider=%s path=%s forwarded_model=%s "
+            "original=%s transformed=%s",
+            self.provider_type,
+            path,
+            original_model_id or "unknown",
+            request_shape(request_body),
+            request_shape(transformed_body),
+        )
+
         logger.debug(
             "Forwarding request to upstream",
             extra={
@@ -2726,15 +2947,18 @@ class BaseUpstreamProvider:
                     upstream_is_streaming = "text/event-stream" in content_type
                     is_streaming = client_wants_streaming and upstream_is_streaming
 
-                    logger.debug(
-                        "Response type analysis",
-                        extra={
-                            "is_streaming": is_streaming,
-                            "client_wants_streaming": client_wants_streaming,
-                            "upstream_is_streaming": upstream_is_streaming,
-                            "content_type": content_type,
-                            "key_hash": key.hashed_key[:8] + "...",
-                        },
+                    logger.warning(
+                        "Upstream response route: provider=%s path=%s model=%s "
+                        "client_stream=%s upstream_stream=%s selected_stream=%s "
+                        "content_type=%s status=%s",
+                        self.provider_type,
+                        path,
+                        original_model_id or "unknown",
+                        client_wants_streaming,
+                        upstream_is_streaming,
+                        is_streaming,
+                        content_type,
+                        response.status_code,
                     )
 
                     if is_streaming and response.status_code == 200:

tests/unit/test_non_streaming_cost_metadata.py

@@ -0,0 +1,73 @@
+import json
+from unittest.mock import AsyncMock, MagicMock
+
+import httpx
+import pytest
+
+from routstr.core.db import ApiKey
+from routstr.upstream import base
+from routstr.upstream.base import BaseUpstreamProvider
+
+
+@pytest.mark.asyncio
+async def test_non_streaming_deepseek_uses_shared_cost_metadata_path(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Non-streaming responses repair and expose their component costs."""
+    provider = BaseUpstreamProvider(
+        base_url="https://api.deepseek.com", api_key="test-key"
+    )
+    key = MagicMock(spec=ApiKey)
+    key.hashed_key = "test-hash"
+    key.balance = 10_000
+    session = MagicMock()
+    session.refresh = AsyncMock()
+    response = httpx.Response(
+        200,
+        json={
+            "id": "chatcmpl-deepseek",
+            "model": "deepseek-v4-flash",
+            "choices": [{"message": {"role": "assistant", "content": "ok"}}],
+            "usage": {
+                "prompt_tokens": 82,
+                "completion_tokens": 121,
+                "total_tokens": 203,
+                "prompt_tokens_details": {"cached_tokens": 0},
+                "completion_tokens_details": {"reasoning_tokens": 115},
+                "prompt_cache_hit_tokens": 0,
+                "prompt_cache_miss_tokens": 82,
+            },
+        },
+    )
+    monkeypatch.setattr(
+        base,
+        "adjust_payment_for_tokens",
+        AsyncMock(
+            return_value={
+                "base_msats": 100,
+                "input_msats": 0,
+                "output_msats": 0,
+                "total_msats": 100,
+                "total_usd": 0.00006,
+                "input_tokens": 82,
+                "output_tokens": 121,
+                "cache_read_input_tokens": 0,
+                "cache_creation_input_tokens": 0,
+            }
+        ),
+    )
+
+    result = await provider.handle_non_streaming_chat_completion(
+        response=response,
+        key=key,
+        session=session,
+        deducted_max_cost=100,
+        requested_model="deepseek-v4-flash",
+    )
+    payload = json.loads(result.body)
+
+    assert payload["cost"]["input_msats"] == 41
+    assert payload["cost"]["output_msats"] == 59
+    assert payload["cost"]["total_msats"] == 100
+    assert payload["metadata"]["routstr"]["cost"]["input_msats"] == 41
+    assert payload["usage"]["cost_sats"] == 0

tests/unit/test_streaming_sse_providers.py

@@ -25,6 +25,7 @@ from unittest.mock import AsyncMock, MagicMock
 import pytest
 
 from routstr.core.db import ApiKey
+from routstr.payment.cost_calculation import CostData
 from routstr.upstream import base
 from routstr.upstream.base import BaseUpstreamProvider
 
@@ -254,6 +255,139 @@ async def test_openrouter_mid_stream_error_event() -> None:
     assert any("error" in o for o in objs), "error event must be forwarded intact"
 
 
+@pytest.mark.asyncio
+async def test_cost_metadata_is_ready_before_finish_reason_is_exposed(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A finish-aware client must not cancel before cost finalization runs."""
+    provider = BaseUpstreamProvider(
+        base_url="https://api.example.com", api_key="test_key"
+    )
+    key = MagicMock(spec=ApiKey)
+    key.hashed_key = "test_hash"
+    key.balance = 1000
+    adjustment = AsyncMock(
+        return_value={
+            "base_msats": 0,
+            "input_msats": 40,
+            "output_msats": 60,
+            "total_msats": 100,
+            "total_usd": 0.0001,
+            "input_tokens": 3,
+            "output_tokens": 2,
+        }
+    )
+    mock_session = MagicMock()
+    mock_session.get = AsyncMock(return_value=key)
+    mock_ctx = MagicMock()
+    mock_ctx.__aenter__ = AsyncMock(return_value=mock_session)
+    mock_ctx.__aexit__ = AsyncMock(return_value=None)
+    warning = MagicMock()
+    monkeypatch.setattr(base, "adjust_payment_for_tokens", adjustment)
+    monkeypatch.setattr(base, "create_session", MagicMock(return_value=mock_ctx))
+    monkeypatch.setattr(base.logger, "warning", warning)
+
+    chunks = [
+        b'data: {"id":"g","model":"deepseek-v4-flash",'
+        b'"choices":[{"delta":{"content":"done"},"finish_reason":"stop"}]}\n\n',
+        b'data: {"id":"g","model":"deepseek-v4-flash","choices":[],'
+        b'"usage":{"prompt_tokens":90,"completion_tokens":233,'
+        b'"total_tokens":323,"prompt_tokens_details":{"cached_tokens":0},'
+        b'"completion_tokens_details":{"reasoning_tokens":227},'
+        b'"prompt_cache_hit_tokens":0,"prompt_cache_miss_tokens":90}}\n\n',
+        b"data: [DONE]\n\n",
+    ]
+    response = await provider.handle_streaming_chat_completion(
+        response=_make_response(chunks),
+        key=key,
+        max_cost_for_model=100,
+        background_tasks=MagicMock(),
+    )
+
+    first = await anext(response.body_iterator)
+
+    first_payload = json.loads(bytes(first).split(b"data: ", 1)[1])
+    assert first_payload["choices"][0]["finish_reason"] == "stop"
+    assert first_payload["cost"]["input_msats"] == 40
+    assert first_payload["cost"]["output_msats"] == 60
+    assert first_payload["metadata"]["routstr"]["cost"]["total_msats"] == 100
+    adjustment.assert_awaited_once()
+    assert any(
+        call.args
+        and call.args[0] == "Client-facing cost metadata: model=%s cost=%s"
+        for call in warning.call_args_list
+    )
+    assert any(
+        call.args
+        and str(call.args[0]).startswith("Streaming SSE event shape:")
+        for call in warning.call_args_list
+    )
+    assert any(
+        call.args
+        and str(call.args[0]).startswith("Streaming cost finalization returned:")
+        for call in warning.call_args_list
+    )
+
+
+@pytest.mark.asyncio
+async def test_payment_failure_preserves_calculated_client_cost(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A DB finalization failure must not replace valid costs with zeros."""
+    provider = BaseUpstreamProvider(
+        base_url="https://api.example.com", api_key="test_key"
+    )
+    key = MagicMock(spec=ApiKey)
+    key.hashed_key = "test_hash"
+    key.balance = 1000
+    mock_session = MagicMock()
+    mock_session.get = AsyncMock(return_value=key)
+    mock_ctx = MagicMock()
+    mock_ctx.__aenter__ = AsyncMock(return_value=mock_session)
+    mock_ctx.__aexit__ = AsyncMock(return_value=None)
+    calculated = CostData(
+        base_msats=0,
+        input_msats=23,
+        output_msats=61,
+        total_msats=84,
+        total_usd=0.00005,
+        input_tokens=89,
+        output_tokens=194,
+    )
+    monkeypatch.setattr(
+        base,
+        "adjust_payment_for_tokens",
+        AsyncMock(side_effect=RuntimeError("concurrent reservation failure")),
+    )
+    monkeypatch.setattr(
+        base,
+        "calculate_cost",
+        AsyncMock(return_value=calculated),
+    )
+    monkeypatch.setattr(base, "create_session", MagicMock(return_value=mock_ctx))
+
+    chunks = [
+        b'data: {"id":"d","model":"deepseek-v4-flash",'
+        b'"choices":[{"delta":{},"finish_reason":"stop"}],'
+        b'"usage":{"prompt_tokens":89,"completion_tokens":194,'
+        b'"total_tokens":283,"prompt_cache_hit_tokens":0,'
+        b'"prompt_cache_miss_tokens":89}}\n\n',
+        b"data: [DONE]\n\n",
+    ]
+    response = await provider.handle_streaming_chat_completion(
+        response=_make_response(chunks),
+        key=key,
+        max_cost_for_model=100,
+        background_tasks=MagicMock(),
+    )
+    first = await anext(response.body_iterator)
+    payload = json.loads(bytes(first).split(b"data: ", 1)[1])
+
+    assert payload["cost"]["input_msats"] == 23
+    assert payload["cost"]["output_msats"] == 61
+    assert payload["cost"]["total_msats"] == 84
+
+
 @pytest.mark.asyncio
 async def test_gemini_combined_content_and_usage_chunk() -> None:
     """Gemini thinking models pack usage into the final *content* chunk.

tests/unit/test_x_cashu_cost_sats.py

@@ -1,5 +1,6 @@
 import json
 import os
+from types import SimpleNamespace
 from unittest.mock import AsyncMock, patch
 
 import httpx
@@ -32,6 +33,44 @@ def _make_cost_data(total_msats: int = 5000) -> CostData:
     )
 
 
+def test_inject_cost_metadata_repairs_zero_components_with_reported_tokens() -> None:
+    provider = _make_provider()
+    response = {
+        "usage": {
+            "prompt_tokens": 83,
+            "completion_tokens": 114,
+            "total_tokens": 197,
+        }
+    }
+    broken_cost = {
+        "base_msats": 5000,
+        "input_msats": 0,
+        "output_msats": 0,
+        "total_msats": 5000,
+        "total_usd": 0.00025,
+        "input_tokens": 83,
+        "output_tokens": 114,
+        "cache_read_input_tokens": 0,
+        "cache_creation_input_tokens": 0,
+    }
+
+    provider.inject_cost_metadata(
+        response,
+        broken_cost,
+        SimpleNamespace(balance=10000),
+    )
+
+    client_cost = response["metadata"]["routstr"]["cost"]
+    assert client_cost["input_msats"] == 2107
+    assert client_cost["output_msats"] == 2893
+    assert client_cost["input_msats"] + client_cost["output_msats"] == 5000
+    assert response["cost"] == {
+        **client_cost,
+        "sats_cost": 5,
+        "remaining_balance_msats": 10000,
+    }
+
+
 # ---------------------------------------------------------------------------
 # Non-streaming (chat completions)
 # ---------------------------------------------------------------------------