Compare commits

...

5 Commits

Author SHA1 Message Date
9qeklajc
0696362b23 better random ids 2026-06-07 11:33:31 +02:00
9qeklajc
c8b4b52424 Merge branch 'main' into ss-refactoring-multi-providers 2026-06-07 11:28:40 +02:00
9qeklajc
3bc30c37e0 retry to update update req. body to fit provider req. 2026-06-07 11:25:35 +02:00
9qeklajc
d4296d6087 fix usage caturing 2026-06-03 17:04:48 +02:00
9qeklajc
320efe2e85 update sse impl. for supported providers 2026-06-02 14:26:54 +02:00
6 changed files with 881 additions and 137 deletions

View File

@@ -28,6 +28,7 @@ from .payment.helpers import (
from .payment.models import Model
from .upstream import BaseUpstreamProvider
from .upstream.helpers import init_upstreams
from .upstream.request_correction import correct_request, extract_error_message
logger = get_logger(__name__)
proxy_router = APIRouter()
@@ -352,50 +353,89 @@ async def proxy(
if request_body_dict:
await pay_for_request(key, max_cost_for_model, session)
# Tracks request params already removed in response to upstream rejections,
# shared across providers so a stripped param stays stripped on failover and
# the reactive retry can never loop unboundedly.
already_stripped: set[str] = set()
for i, upstream in enumerate(upstreams):
headers = upstream.prepare_headers(dict(request.headers))
try:
try:
if is_responses_api:
response = await upstream.forward_responses_request(
request,
path,
headers,
request_body,
key,
max_cost_for_model,
session,
model_obj,
while True:
try:
if is_responses_api:
response = await upstream.forward_responses_request(
request,
path,
headers,
request_body,
key,
max_cost_for_model,
session,
model_obj,
)
else:
response = await upstream.forward_request(
request,
path,
headers,
request_body,
key,
max_cost_for_model,
session,
model_obj,
)
except UpstreamError:
# Let the outer UpstreamError handler manage retry/revert
raise
except Exception as e:
# Unexpected error (not an upstream failure) — revert and propagate
logger.error(
"Unexpected error in upstream request, reverting payment",
extra={
"error": str(e),
"error_type": type(e).__name__,
"path": path,
"key_hash": key.hashed_key[:8] + "...",
"max_cost_for_model": max_cost_for_model,
},
)
else:
response = await upstream.forward_request(
request,
path,
headers,
await revert_pay_for_request(key, session, max_cost_for_model)
raise
# Reactive recovery: some models reject one specific request
# param (e.g. newer Anthropic models deprecating `temperature`).
# When the upstream 400s naming such a param, strip it from the
# body and retry the SAME upstream. ``already_stripped`` bounds
# this to one retry per distinct param so it always terminates.
if response.status_code == 400:
correction = correct_request(
request_body,
key,
max_cost_for_model,
session,
model_obj,
extract_error_message(response),
already_stripped,
)
except UpstreamError:
# Let the outer UpstreamError handler manage retry/revert
raise
except Exception as e:
# Unexpected error (not an upstream failure) — revert and propagate
logger.error(
"Unexpected error in upstream request, reverting payment",
extra={
"error": str(e),
"error_type": type(e).__name__,
"path": path,
"key_hash": key.hashed_key[:8] + "...",
"max_cost_for_model": max_cost_for_model,
},
)
await revert_pay_for_request(key, session, max_cost_for_model)
raise
if correction is not None:
request_body, bad_param = correction.body, correction.label
already_stripped.add(bad_param)
request_body_dict = parse_request_body_json(
request_body, path
)
logger.warning(
"Upstream %s rejected param '%s' for model=%s; "
"stripping and retrying same upstream",
upstream.provider_type,
bad_param,
model_id,
extra={
"provider": upstream.provider_type,
"model": model_id,
"stripped_param": bad_param,
"path": path,
},
)
continue
break
if response.status_code != 200:
# Check if we should retry (502 Upstream Error or 429 Rate Limit)

View File

@@ -2,10 +2,9 @@ from __future__ import annotations
import asyncio
import json
import re
import traceback
import uuid
from collections.abc import AsyncGenerator, AsyncIterator
from collections.abc import AsyncGenerator, AsyncIterator, Iterator
from typing import Any, Mapping, cast
import httpx
@@ -716,56 +715,134 @@ class BaseUpstreamProvider:
except Exception:
pass
def _process_event(raw_event: bytes) -> Iterator[bytes]:
"""Process one complete SSE event block (lines up to a blank line).
Handles arbitrary upstream framing across every supported
provider:
* ``data:`` lines are gathered and concatenated per the SSE
spec, so a payload split across network chunks is reassembled
before parsing.
* Comment/keepalive lines (those beginning with ``:`` such as
OpenRouter's ``: OPENROUTER PROCESSING``) are dropped. They
carry no JSON and forwarding them downstream breaks naive SSE
clients; the keepalive only matters for the upstream hop.
* Other SSE fields (``event:``/``id:``/``retry:``) are preserved
and kept attached to the event's ``data:`` line, which the
OpenAI Responses API and Anthropic-style streams rely on.
* ``[DONE]`` is swallowed so it can be re-emitted exactly once at
end of stream.
"""
nonlocal last_model_seen, usage_chunk_data, done_seen
event = raw_event.strip(b"\r\n")
if not event:
return
field_lines: list[bytes] = []
data_lines: list[bytes] = []
for line in event.split(b"\n"):
line = line.rstrip(b"\r")
if line.startswith(b"data:"):
# Strip the field name and a single optional leading space.
data_lines.append(line[len(b"data:") :].lstrip(b" "))
elif line.startswith(b":"):
# SSE comment / keepalive - drop.
continue
elif line:
# Other SSE field (event:/id:/retry:) - preserve in order.
field_lines.append(line)
if not data_lines:
return
data = b"\n".join(data_lines)
if not data.strip():
return
# Re-emit preserved SSE fields immediately before the data line so
# event/data framing stays intact (single trailing newline each;
# the blank-line terminator is appended to the data line below).
prefix = b"".join(fl + b"\n" for fl in field_lines)
if data.strip() == b"[DONE]":
done_seen = True
return
try:
obj = json.loads(data)
except Exception:
obj = None
if isinstance(obj, dict):
self._apply_provider_field(obj)
if obj.get("model"):
last_model_seen = str(obj.get("model"))
if requested_model:
obj["model"] = requested_model
if (
"id" not in obj
or not isinstance(obj["id"], str)
or obj["id"] == "existing-id"
):
if not hasattr(self, "_current_stream_id"):
self._current_stream_id = f"chatcmpl-{uuid.uuid4()}"
obj["id"] = self._current_stream_id
if isinstance(obj.get("usage"), dict):
# Capture usage for end-of-stream cost reconciliation.
# Some models (e.g. Gemini thinking models over the
# OpenAI-compat endpoint) attach ``usage`` to the SAME
# chunk that carries the final content/finish_reason
# rather than sending a separate ``choices: []`` usage
# chunk. Only swallow the chunk when it is a pure usage
# chunk (no choices); otherwise the content would be
# silently dropped and the client would receive no
# assistant message at all.
if obj.get("choices"):
# Capture usage (with model) for the cost trailer,
# but with choices stripped so the trailer never
# re-emits this chunk's content.
usage_chunk_data = {
k: v for k, v in obj.items() if k != "choices"
}
usage_chunk_data["choices"] = []
# 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 (
prefix
+ b"data: "
+ json.dumps(forward).encode()
+ b"\n\n"
)
return
usage_chunk_data = obj
return
yield prefix + b"data: " + json.dumps(obj).encode() + b"\n\n"
else:
# Non-JSON data payload (partial fragment already reassembled
# by buffering, or a provider control string) - forward as-is.
yield prefix + b"data: " + data + b"\n\n"
try:
# Buffer bytes across network chunks and dispatch only on the SSE
# event delimiter (a blank line). ``aiter_bytes`` yields arbitrary
# byte boundaries, so a single event's JSON can span chunks and
# multiple events can arrive together; buffering makes parsing
# boundary-independent for every provider.
buffer = b""
async for chunk in response.aiter_bytes():
# Split chunk into SSE events
parts = re.split(b"data: ", chunk)
for i, part in enumerate(parts):
if not part:
continue
buffer += chunk.replace(b"\r\n", b"\n")
while b"\n\n" in buffer:
raw_event, buffer = buffer.split(b"\n\n", 1)
for out in _process_event(raw_event):
yield out
stripped_part = part.strip()
if not stripped_part:
continue
if stripped_part == b"[DONE]":
done_seen = True
continue
try:
# Only parse if it looks like a JSON object to avoid SSE control messages or partials
if part.strip().startswith(b"{") and part.strip().endswith(
b"}"
):
obj = json.loads(part)
if isinstance(obj, dict):
self._apply_provider_field(obj)
if obj.get("model"):
last_model_seen = str(obj.get("model"))
if requested_model:
obj["model"] = requested_model
if (
"id" not in obj
or not isinstance(obj["id"], str)
or obj["id"] == "existing-id"
):
if not hasattr(self, "_current_stream_id"):
self._current_stream_id = (
f"chatcmpl-{uuid.uuid4()}"
)
obj["id"] = self._current_stream_id
if isinstance(obj.get("usage"), dict):
usage_chunk_data = obj
continue
yield b"data: " + json.dumps(obj).encode() + b"\n\n"
continue
except Exception:
pass
prefix = (
b"data: " if (i > 0 or chunk.startswith(b"data: ")) else b""
)
yield prefix + part
# Flush any trailing event that lacked a final blank line.
if buffer.strip():
for out in _process_event(buffer):
yield out
async with create_session() as session:
fresh_key = await session.get(key.__class__, key.hashed_key)
@@ -1067,56 +1144,92 @@ class BaseUpstreamProvider:
except Exception:
pass
def _process_event(raw_event: bytes) -> Iterator[bytes]:
"""Process one complete SSE event block for the Responses API.
Buffers full events (delimited by a blank line) so parsing is
boundary-independent, gathers ``data:`` lines, drops comment/
keepalive lines (e.g. OpenRouter's ``: OPENROUTER PROCESSING``),
and preserves ``event:``/``id:`` fields attached to their data
line so Responses API event framing stays intact.
"""
nonlocal last_model_seen, usage_chunk_data, done_seen
nonlocal reasoning_tokens
event = raw_event.strip(b"\r\n")
if not event:
return
field_lines: list[bytes] = []
data_lines: list[bytes] = []
for line in event.split(b"\n"):
line = line.rstrip(b"\r")
if line.startswith(b"data:"):
data_lines.append(line[len(b"data:") :].lstrip(b" "))
elif line.startswith(b":"):
# SSE comment / keepalive - drop.
continue
elif line:
# Preserve event:/id:/retry: (Responses API event names).
field_lines.append(line)
if not data_lines:
return
data = b"\n".join(data_lines)
if not data.strip():
return
prefix = b"".join(fl + b"\n" for fl in field_lines)
if data.strip() == b"[DONE]":
done_seen = True
return
try:
obj = json.loads(data)
except json.JSONDecodeError:
obj = None
if isinstance(obj, dict):
self._apply_provider_field(obj)
if obj.get("model"):
last_model_seen = str(obj.get("model"))
if requested_model:
obj["model"] = requested_model
# Track reasoning tokens for Responses API
if usage := obj.get("usage", {}):
if isinstance(usage, dict) and "reasoning_tokens" in usage:
reasoning_tokens += usage.get("reasoning_tokens", 0)
# Responses API usage is in response.completed/incomplete events
chunk_type = obj.get("type", "")
if chunk_type in (
"response.completed",
"response.incomplete",
):
usage_chunk_data = obj
return
yield prefix + b"data: " + json.dumps(obj).encode() + b"\n\n"
else:
yield prefix + b"data: " + data + b"\n\n"
try:
# Buffer across network chunks; dispatch only on the SSE event
# delimiter so parsing is independent of byte boundaries.
buffer = b""
async for chunk in response.aiter_bytes():
# Split chunk into SSE events
parts = re.split(b"data: ", chunk)
for i, part in enumerate(parts):
if not part:
continue
buffer += chunk.replace(b"\r\n", b"\n")
while b"\n\n" in buffer:
raw_event, buffer = buffer.split(b"\n\n", 1)
for out in _process_event(raw_event):
yield out
stripped_part = part.strip()
if not stripped_part:
continue
if stripped_part == b"[DONE]":
done_seen = True
continue
try:
obj = json.loads(part)
if isinstance(obj, dict):
self._apply_provider_field(obj)
if obj.get("model"):
last_model_seen = str(obj.get("model"))
if requested_model:
obj["model"] = requested_model
# Track reasoning tokens for Responses API
if usage := obj.get("usage", {}):
if (
isinstance(usage, dict)
and "reasoning_tokens" in usage
):
reasoning_tokens += usage.get(
"reasoning_tokens", 0
)
# Responses API usage is in response.completed/incomplete events
chunk_type = obj.get("type", "")
if chunk_type in (
"response.completed",
"response.incomplete",
):
usage_chunk_data = obj
continue
except json.JSONDecodeError:
pass
prefix = (
b"data: " if (i > 0 or chunk.startswith(b"data: ")) else b""
)
yield prefix + part
if buffer.strip():
for out in _process_event(buffer):
yield out
# Always emit a cost-bearing data chunk
async with create_session() as session:

View File

@@ -0,0 +1,142 @@
"""Reactive request-correction layer.
When an upstream rejects a request with a recoverable 4xx error, this layer
tries to *fix* the request body and let the caller retry the same upstream
instead of failing outright. It is provider-agnostic: correctors key off the
upstream's own error wording, so the same recovery works across every provider.
The layer is a small pipeline of :data:`Corrector` callables. Each corrector
inspects the parsed request body and the upstream error message and either
returns a corrected body (plus a short label identifying the fix) or declines
by returning ``None``. Adding a new reactive fix means writing one corrector
and adding it to :data:`DEFAULT_CORRECTORS` — no changes to the proxy loop.
All corrections are immutable: a corrector never mutates the body it is given,
it returns a new ``dict``. The proxy threads an ``applied`` set of fix labels
through retries so each distinct fix is applied at most once, guaranteeing the
retry loop always terminates.
"""
from __future__ import annotations
import json
import re
from collections.abc import Callable, Sequence
from dataclasses import dataclass
from fastapi.responses import Response
from ..core import get_logger
logger = get_logger(__name__)
# Matches upstream error text that names a single rejected request parameter,
# e.g. "`temperature` is deprecated for this model." or
# "parameter 'top_p' is not supported". Keys off the upstream's own wording so
# a 400 about an unsupported sampling/option field can be recovered by stripping
# that field and retrying the same upstream.
_UNSUPPORTED_PARAM_RE = re.compile(
r"[`'\"]?(?P<param>[a-zA-Z_][a-zA-Z0-9_]*)[`'\"]?\s+is\s+"
r"(?:deprecated|not\s+supported|unsupported|no\s+longer\s+supported)",
re.IGNORECASE,
)
# A corrector inspects the parsed request body and the upstream error message
# and returns ``(new_body_dict, label)`` for a fix it can apply, or ``None`` to
# decline. ``label`` identifies the fix so it is applied at most once per request.
Corrector = Callable[[dict, str], "tuple[dict, str] | None"]
@dataclass(frozen=True)
class Correction:
"""A successful request correction ready to retry.
``body`` is the corrected JSON body (encoded), ``label`` identifies the fix
that was applied (e.g. the stripped param name) so the caller can guard
against applying the same fix twice.
"""
body: bytes
label: str
def extract_error_message(response: Response) -> str:
"""Best-effort extraction of an error message string from a proxy Response."""
body_bytes = getattr(response, "body", None)
if not body_bytes:
return ""
try:
data = json.loads(body_bytes)
except Exception:
return body_bytes.decode("utf-8", errors="ignore")[:500]
if isinstance(data, dict):
err = data.get("error")
if isinstance(err, dict):
msg = err.get("message") or err.get("detail")
if isinstance(msg, str):
return msg
elif isinstance(err, str):
return err
if isinstance(data.get("message"), str):
return data["message"]
return ""
def strip_unsupported_param(
body: dict, error_message: str
) -> tuple[dict, str] | None:
"""Drop a top-level param the upstream named as unsupported/deprecated.
Returns ``(new_body, param)`` (a new dict, original untouched) when the
error names a top-level param present in the body, otherwise ``None``.
"""
match = _UNSUPPORTED_PARAM_RE.search(error_message)
if not match:
return None
param = match.group("param")
if param not in body:
return None
new_body = {k: v for k, v in body.items() if k != param}
return new_body, param
# Ordered pipeline of correctors tried on each recoverable rejection.
DEFAULT_CORRECTORS: tuple[Corrector, ...] = (strip_unsupported_param,)
def correct_request(
request_body: bytes,
error_message: str,
applied: set[str],
correctors: Sequence[Corrector] = DEFAULT_CORRECTORS,
) -> Correction | None:
"""Try to correct a rejected request body so it can be retried.
Runs each corrector in order against the parsed body and ``error_message``.
The first corrector that proposes a fix whose ``label`` is not already in
``applied`` wins; its result is returned as a :class:`Correction`. Returns
``None`` when nothing parses, nothing matches, or every proposed fix was
already applied — the caller then treats the response as a normal failure.
``applied`` is read-only here; the caller records the returned ``label`` to
bound retries and guarantee forward progress.
"""
if not request_body or not error_message:
return None
try:
data = json.loads(request_body)
except Exception:
return None
if not isinstance(data, dict):
return None
for corrector in correctors:
result = corrector(data, error_message)
if result is None:
continue
new_body, label = result
if label in applied:
continue
return Correction(body=json.dumps(new_body).encode(), label=label)
return None

View File

@@ -174,17 +174,20 @@ class TestmintWallet:
"""Fallback method to create a basic test token"""
import base64
import json
import random
import time
import secrets
unique_id = int(time.time() * 1000000) + random.randint(1000, 9999)
# Use a cryptographically random id/secret so every minted token is
# guaranteed unique. Time/PRNG-based ids can collide on hosts with
# coarse clock resolution, producing byte-identical tokens that hash to
# the same api_key and silently dedupe (flaky concurrency tests).
unique_id = secrets.token_hex(16)
token_data = {
"token": [
{
"mint": self.mint_url,
"proofs": [
{
"id": f"009a1f293253e41e{unique_id % 100000000:08d}",
"id": f"009a1f293253e41e{unique_id[:8]}",
"amount": amount,
"secret": f"test-secret-{amount}-{unique_id}",
"C": "02194603ffa36356f4a56b7df9371fc3192472351453ec7398b8da8117e7c3e104",

View File

@@ -0,0 +1,130 @@
"""Unit tests for the reactive request-correction layer.
Covers the recovery path that lets a request survive a 400 where the upstream
names a single unsupported request param (e.g. newer Anthropic models
deprecating ``temperature``): the param is stripped from the JSON body and the
same upstream is retried, provider-agnostically, keyed off the error text.
"""
from __future__ import annotations
import json
from fastapi.responses import Response
from routstr.upstream.request_correction import (
Correction,
correct_request,
extract_error_message,
strip_unsupported_param,
)
def _body(**kwargs: object) -> bytes:
return json.dumps(kwargs).encode()
class TestCorrectRequest:
def test_strips_deprecated_temperature(self) -> None:
body = _body(model="claude-opus-4-8", temperature=1, messages=[])
result = correct_request(
body, "`temperature` is deprecated for this model.", set()
)
assert isinstance(result, Correction)
assert result.label == "temperature"
decoded = json.loads(result.body)
assert "temperature" not in decoded
assert decoded["model"] == "claude-opus-4-8"
def test_strips_not_supported_param(self) -> None:
body = _body(model="m", top_p=0.9, messages=[])
result = correct_request(body, "Parameter 'top_p' is not supported", set())
assert result is not None
assert result.label == "top_p"
assert "top_p" not in json.loads(result.body)
def test_returns_none_when_label_already_applied(self) -> None:
body = _body(model="m", temperature=1)
assert (
correct_request(body, "`temperature` is deprecated", {"temperature"})
is None
)
def test_returns_none_when_param_absent_from_body(self) -> None:
body = _body(model="m", messages=[])
assert correct_request(body, "`temperature` is deprecated", set()) is None
def test_returns_none_when_message_does_not_match(self) -> None:
body = _body(model="m", temperature=1)
assert correct_request(body, "Insufficient balance", set()) is None
def test_returns_none_on_empty_inputs(self) -> None:
assert correct_request(b"", "`temperature` is deprecated", set()) is None
assert correct_request(_body(temperature=1), "", set()) is None
def test_returns_none_on_non_object_body(self) -> None:
assert correct_request(b"[1, 2, 3]", "`temperature` is deprecated", set()) is None
class TestStripUnsupportedParam:
def test_does_not_mutate_input(self) -> None:
body = {"model": "m", "temperature": 1}
result = strip_unsupported_param(body, "`temperature` is deprecated")
assert result is not None
new_body, param = result
assert param == "temperature"
assert "temperature" not in new_body
# original untouched (immutability)
assert body == {"model": "m", "temperature": 1}
def test_declines_when_no_match(self) -> None:
assert strip_unsupported_param({"temperature": 1}, "nope") is None
class TestExtractErrorMessage:
def test_extracts_nested_error_message(self) -> None:
resp = Response(
content=json.dumps(
{"error": {"message": "`temperature` is deprecated", "type": "x"}}
).encode(),
status_code=400,
)
assert extract_error_message(resp) == "`temperature` is deprecated"
def test_extracts_string_error(self) -> None:
resp = Response(
content=json.dumps({"error": "bad request"}).encode(), status_code=400
)
assert extract_error_message(resp) == "bad request"
def test_extracts_top_level_message(self) -> None:
resp = Response(
content=json.dumps({"message": "nope"}).encode(), status_code=400
)
assert extract_error_message(resp) == "nope"
def test_empty_body_returns_empty_string(self) -> None:
assert extract_error_message(Response(status_code=400)) == ""
def test_non_json_body_returns_preview(self) -> None:
resp = Response(content=b"plain text error", status_code=400)
assert extract_error_message(resp) == "plain text error"
class TestEndToEndChaining:
def test_two_distinct_params_corrected_sequentially(self) -> None:
"""Simulates the proxy loop: each 400 fixes one param, set guards reuse."""
body = _body(model="m", temperature=1, top_p=0.5, messages=[])
applied: set[str] = set()
first = correct_request(body, "`temperature` is deprecated", applied)
assert first is not None
body, applied = first.body, applied | {first.label}
second = correct_request(body, "`top_p` is not supported", applied)
assert second is not None
body, applied = second.body, applied | {second.label}
decoded = json.loads(body)
assert "temperature" not in decoded and "top_p" not in decoded
assert applied == {"temperature", "top_p"}

View File

@@ -0,0 +1,316 @@
"""Battle-test the streaming SSE parser against real per-provider framing.
Each test drives the *actual* ``handle_streaming_chat_completion`` generator
with a mock upstream response whose ``aiter_bytes`` emits byte sequences that
mirror what each supported provider sends on the wire (captured from the
providers' own streaming docs):
* OpenAI / Groq / Fireworks / xAI / Perplexity / Azure - plain
``data: {json}\\n\\n`` + ``data: [DONE]``.
* OpenRouter - same, but with ``: OPENROUTER PROCESSING`` keepalive comments
interleaved (the framing that produced the original
``Unexpected token ':'`` client crash).
* Gemini (native ``alt=sse``) - ``data:`` payloads framed with CRLF.
The invariant every provider must satisfy: every line the proxy emits that
starts with ``data: `` either equals ``[DONE]`` or is valid JSON, and no SSE
comment ever reaches the client. That invariant is exactly what the buggy
``re.split(b"data: ")`` parser violated for OpenRouter.
"""
import json
from collections.abc import AsyncGenerator
from unittest.mock import AsyncMock, MagicMock
import pytest
from routstr.core.db import ApiKey
from routstr.upstream import base
from routstr.upstream.base import BaseUpstreamProvider
def _make_response(chunks: list[bytes]) -> MagicMock:
async def aiter_bytes() -> AsyncGenerator[bytes, None]:
for chunk in chunks:
yield chunk
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {"content-type": "text/event-stream"}
mock_response.aiter_bytes = aiter_bytes
return mock_response
async def _drive(chunks: list[bytes], requested_model: str | None = None) -> list[bytes]:
"""Run the real streaming generator over ``chunks`` and collect output bytes."""
provider = BaseUpstreamProvider(
base_url="https://api.example.com", api_key="test_key"
)
key = MagicMock(spec=ApiKey)
key.hashed_key = "test_hash"
key.balance = 1000
base.adjust_payment_for_tokens = AsyncMock(
return_value={"total_usd": 0.1, "total_msats": 100}
)
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)
base.create_session = MagicMock(return_value=mock_ctx)
streaming_response = await provider.handle_streaming_chat_completion(
response=_make_response(chunks),
key=key,
max_cost_for_model=100,
background_tasks=MagicMock(),
requested_model=requested_model,
)
out: list[bytes] = []
async for chunk in streaming_response.body_iterator:
if isinstance(chunk, str):
out.append(chunk.encode())
else:
out.append(bytes(chunk))
return out
def _data_payloads(out: list[bytes]) -> list[bytes]:
"""Return the raw payload of every ``data: `` line across all emitted bytes."""
payloads: list[bytes] = []
for chunk in out:
for line in chunk.split(b"\n"):
if line.startswith(b"data: "):
payloads.append(line[len(b"data: ") :])
return payloads
def _assert_clean(out: list[bytes]) -> list[dict]:
"""Core invariant: every data line is [DONE] or valid JSON; no comments leak."""
blob = b"".join(out)
# No SSE comment line must ever reach the client.
for line in blob.split(b"\n"):
assert not line.startswith(b":"), f"comment leaked to client: {line!r}"
# The original bug signature: a data line whose value is itself a comment.
assert not line.startswith(b"data: :"), f"mangled comment frame: {line!r}"
objs: list[dict] = []
for payload in _data_payloads(out):
stripped = payload.strip()
if stripped == b"[DONE]":
continue
obj = json.loads(stripped) # raises if the proxy emitted non-JSON data
objs.append(obj)
return objs
@pytest.mark.asyncio
async def test_openai_style_plain_stream() -> None:
"""OpenAI / Groq / Fireworks / xAI / Perplexity: plain data + [DONE]."""
chunks = [
b'data: {"id":"x","choices":[{"delta":{"content":"Hello"}}]}\n\n',
b'data: {"id":"x","choices":[{"delta":{"content":" world"}}]}\n\n',
b"data: [DONE]\n\n",
]
out = await _drive(chunks)
objs = _assert_clean(out)
assert any(o.get("choices") for o in objs)
assert b"data: [DONE]\n\n" in b"".join(out)
@pytest.mark.asyncio
async def test_openrouter_keepalive_comments() -> None:
"""OpenRouter ``: OPENROUTER PROCESSING`` keepalives must never crash clients.
This is the exact regression: the old parser emitted
``data: : OPENROUTER PROCESSING`` which made downstream
``JSON.parse`` throw ``Unexpected token ':'``.
"""
chunks = [
b": OPENROUTER PROCESSING\n\n",
b": OPENROUTER PROCESSING\n\n",
b'data: {"id":"x","choices":[{"delta":{"content":"Hi"}}]}\n\n',
b": OPENROUTER PROCESSING\n\n",
b'data: {"id":"x","choices":[{"delta":{"content":"!"}}]}\n\n',
b"data: [DONE]\n\n",
]
out = await _drive(chunks)
objs = _assert_clean(out)
# The keepalive must be gone entirely.
assert b"OPENROUTER PROCESSING" not in b"".join(out)
# Real content survived.
contents = [
c["delta"]["content"]
for o in objs
for c in o.get("choices", [])
if "delta" in c
]
assert "Hi" in contents and "!" in contents
@pytest.mark.asyncio
async def test_openrouter_comment_glued_to_data_chunk() -> None:
"""Keepalive packed into the same TCP chunk as data (the harder case)."""
chunks = [
b'data: {"id":"x","choices":[{"delta":{"content":"a"}}]}\n\n'
b": OPENROUTER PROCESSING\n\n"
b'data: {"id":"x","choices":[{"delta":{"content":"b"}}]}\n\n',
b"data: [DONE]\n\n",
]
out = await _drive(chunks)
objs = _assert_clean(out)
contents = [
c["delta"]["content"]
for o in objs
for c in o.get("choices", [])
if "delta" in c
]
assert contents == ["a", "b"]
@pytest.mark.asyncio
async def test_json_split_across_chunk_boundary() -> None:
"""A single event's JSON arriving in two TCP reads must reassemble."""
chunks = [
b'data: {"id":"x","choices":[{"delta":{"con',
b'tent":"split"}}]}\n\n',
b"data: [DONE]\n\n",
]
out = await _drive(chunks)
objs = _assert_clean(out)
contents = [
c["delta"]["content"]
for o in objs
for c in o.get("choices", [])
if "delta" in c
]
assert contents == ["split"]
@pytest.mark.asyncio
async def test_byte_by_byte_fragmentation() -> None:
"""Pathological framing: one byte per chunk. Must still parse cleanly."""
raw = (
b'data: {"id":"x","choices":[{"delta":{"content":"drip"}}]}\n\n'
b": OPENROUTER PROCESSING\n\n"
b"data: [DONE]\n\n"
)
chunks = [raw[i : i + 1] for i in range(len(raw))]
out = await _drive(chunks)
objs = _assert_clean(out)
assert objs and objs[0]["choices"][0]["delta"]["content"] == "drip"
@pytest.mark.asyncio
async def test_gemini_crlf_framing() -> None:
"""Gemini native (alt=sse) frames events with CRLF."""
chunks = [
b'data: {"id":"g","choices":[{"delta":{"content":"hej"}}]}\r\n\r\n',
b'data: {"id":"g","choices":[{"delta":{"content":"!"}}]}\r\n\r\n',
]
out = await _drive(chunks)
objs = _assert_clean(out)
contents = [
c["delta"]["content"]
for o in objs
for c in o.get("choices", [])
if "delta" in c
]
assert contents == ["hej", "!"]
@pytest.mark.asyncio
async def test_azure_leading_role_chunk() -> None:
"""Azure OpenAI opens with a content-filter / role-only chunk."""
chunks = [
b'data: {"id":"az","choices":[],"prompt_filter_results":[]}\n\n',
b'data: {"id":"az","choices":[{"delta":{"role":"assistant"}}]}\n\n',
b'data: {"id":"az","choices":[{"delta":{"content":"ok"}}]}\n\n',
b"data: [DONE]\n\n",
]
out = await _drive(chunks)
_assert_clean(out)
@pytest.mark.asyncio
async def test_openrouter_mid_stream_error_event() -> None:
"""OpenRouter mid-stream errors arrive as a normal data JSON event."""
err = {
"id": "x",
"object": "chat.completion.chunk",
"model": "openai/gpt-4o",
"error": {"code": "server_error", "message": "Provider disconnected"},
"choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": "error"}],
}
chunks = [
b'data: {"id":"x","choices":[{"delta":{"content":"partial"}}]}\n\n',
b"data: " + json.dumps(err).encode() + b"\n\n",
]
out = await _drive(chunks)
objs = _assert_clean(out)
assert any("error" in o for o in objs), "error event must be forwarded intact"
@pytest.mark.asyncio
async def test_gemini_combined_content_and_usage_chunk() -> None:
"""Gemini thinking models pack usage into the final *content* chunk.
Regression: the parser swallowed any chunk carrying a ``usage`` dict, so
when content + usage arrived together the assistant text was dropped and
the client saw "no assistant messages" despite a 200 + token accounting.
"""
chunks = [
b'data: {"id":"g","choices":[{"delta":{"content":"the answer"},'
b'"finish_reason":"stop"}],"usage":{"prompt_tokens":3,'
b'"completion_tokens":2,"total_tokens":5}}\n\n',
b"data: [DONE]\n\n",
]
out = await _drive(chunks)
objs = _assert_clean(out)
contents = [
c["delta"]["content"]
for o in objs
for c in o.get("choices", [])
if "delta" in c
]
# Content delivered exactly once (not dropped, not duplicated by the trailer).
assert contents == ["the answer"]
@pytest.mark.asyncio
async def test_separate_usage_chunk_not_forwarded_as_content() -> None:
"""A pure usage chunk (choices: []) is still swallowed, content intact."""
chunks = [
b'data: {"id":"x","choices":[{"delta":{"content":"hello"}}]}\n\n',
b'data: {"id":"x","choices":[],"usage":{"total_tokens":4}}\n\n',
b"data: [DONE]\n\n",
]
out = await _drive(chunks)
objs = _assert_clean(out)
contents = [
c["delta"]["content"]
for o in objs
for c in o.get("choices", [])
if "delta" in c
]
assert contents == ["hello"]
@pytest.mark.asyncio
async def test_requested_model_override_applied() -> None:
"""Model rewriting still works through the buffered parser."""
chunks = [
b'data: {"id":"x","model":"upstream-model","choices":[{"delta":{"content":"hi"}}]}\n\n',
b"data: [DONE]\n\n",
]
out = await _drive(chunks, requested_model="routstr-model")
objs = _assert_clean(out)
# The upstream content chunk carried model "upstream-model"; the parser must
# rewrite it to the requested model. (The trailing routstr-generated usage
# chunk is excluded - it is not an upstream-forwarded chunk.)
content_chunks = [o for o in objs if o.get("choices")]
assert content_chunks, "expected at least one forwarded content chunk"
assert all(o.get("model") == "routstr-model" for o in content_chunks)