streaming not stoped

routstr/upstream/base.py

@@ -853,6 +853,7 @@ class BaseUpstreamProvider:
             last_model_seen: str | None = None
             usage_chunk_data: dict | None = None
             done_seen: bool = False
+            pending_finish_event: bytes | None = None
 
             async def finalize_db_only() -> None:
                 nonlocal usage_finalized
@@ -895,6 +896,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:
@@ -970,16 +972,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 = encoded
+                            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
+                        # client cost metadata are ready so cancellation cannot
+                        # make the cost trailer disappear.
+                        pending_finish_event = encoded
+                    else:
+                        yield encoded
                 else:
                     if final:
                         # Final flush of a truncated tail: the upstream closed
@@ -1022,6 +1044,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:
@@ -1098,8 +1121,14 @@ 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()
+                        )
 
+                if pending_finish_event is not None:
+                    yield pending_finish_event
+                if cost_trailer is not None:
+                    yield cost_trailer
                 if done_seen:
                     yield b"data: [DONE]\n\n"
 

tests/unit/test_streaming_sse_providers.py

@@ -254,6 +254,62 @@ 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"}],'
+        b'"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}\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)
+
+    assert b'"finish_reason": "stop"' in bytes(first)
+    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
+    )
+
+
 @pytest.mark.asyncio
 async def test_gemini_combined_content_and_usage_chunk() -> None:
     """Gemini thinking models pack usage into the final *content* chunk.