diff --git a/routstr/core/exceptions.py b/routstr/core/exceptions.py index 4262f91c..8dcda2a5 100644 --- a/routstr/core/exceptions.py +++ b/routstr/core/exceptions.py @@ -7,11 +7,26 @@ logger = get_logger(__name__) class UpstreamError(Exception): - """Exception raised when an upstream provider fails.""" + """Exception raised when an upstream provider fails. - def __init__(self, message: str, status_code: int = 502): + ``code`` carries a stable, machine-readable classification (e.g. + ``UPSTREAM_RATE_LIMIT``) so callers can distinguish failure kinds without + string-matching the message. ``details`` holds optional structured, + redaction-safe context. Both default to ``None`` for backwards + compatibility. + """ + + def __init__( + self, + message: str, + status_code: int = 502, + code: str | None = None, + details: dict[str, object] | None = None, + ): self.message = message self.status_code = status_code + self.code = code + self.details = details super().__init__(message) diff --git a/routstr/core/logging.py b/routstr/core/logging.py index fac0100f..3886637c 100644 --- a/routstr/core/logging.py +++ b/routstr/core/logging.py @@ -51,6 +51,8 @@ from pythonjsonlogger import jsonlogger from rich.console import Console from rich.logging import RichHandler +from .redaction import redact_obj, redact_org_ids + # Only use RichHandler when stdout is a real TTY. In non-TTY contexts # (docker logs, pipes, CI) Rich pads every line to width and wraps long # records, producing visually-empty trailing whitespace and split records. @@ -180,6 +182,37 @@ class RequestIdFilter(logging.Filter): return True +# Standard ``LogRecord`` attributes that are never user-supplied ``extra`` +# fields; skipped when redacting structured extras (``msg``/``message`` are +# handled separately above). +_NON_EXTRA_RECORD_ATTRS = frozenset( + { + "name", + "msg", + "args", + "levelname", + "levelno", + "pathname", + "filename", + "module", + "exc_info", + "exc_text", + "stack_info", + "lineno", + "funcName", + "created", + "msecs", + "relativeCreated", + "thread", + "threadName", + "processName", + "process", + "taskName", + "message", + } +) + + class SecurityFilter(logging.Filter): """Filter to remove sensitive information from logs.""" @@ -203,6 +236,7 @@ class SecurityFilter(logging.Filter): """Filter out sensitive information from log records.""" try: message = record.getMessage() + message = redact_org_ids(message) standalone_patterns = [ r"Bearer\s+([a-zA-Z0-9_\-\.]{10,})", # Bearer token (must be 10 characters or more to reduce false-positives) r"cashu[A-Z]+([a-zA-Z0-9_\-\.=/+]+)", # Cashu tokens @@ -224,6 +258,16 @@ class SecurityFilter(logging.Filter): record.msg = message record.args = () + # Structured `extra={...}` fields are emitted by the JSON formatter + # straight from the record dict and never pass through the message + # formatting above. Redact organization IDs from any string-valued + # extra so they cannot leak via structured logs. + for attr, value in list(record.__dict__.items()): + if attr in _NON_EXTRA_RECORD_ATTRS: + continue + if isinstance(value, (str, dict, list, tuple)): + record.__dict__[attr] = redact_obj(value) + except Exception: pass diff --git a/routstr/core/redaction.py b/routstr/core/redaction.py new file mode 100644 index 00000000..c0519581 --- /dev/null +++ b/routstr/core/redaction.py @@ -0,0 +1,51 @@ +"""Redaction helpers for sensitive provider identifiers. + +Single source of truth for stripping account-scoped identifiers (e.g. OpenAI +organization IDs) from any text before it is logged, returned to a caller, or +written to an audit entry. +""" + +from __future__ import annotations + +import re +from typing import Any + +# OpenAI-style organization identifiers look like ``org-``. Require at +# least 6 trailing chars so the already-redacted literal ``org-[REDACTED]`` is +# never re-matched (``[`` is not in the character class). +_ORG_ID_PATTERN = re.compile(r"\borg-[A-Za-z0-9]{6,}\b") + +ORG_ID_PLACEHOLDER = "org-[REDACTED]" + + +def redact_org_ids(text: str) -> str: + """Replace OpenAI-style organization IDs with ``org-[REDACTED]``. + + Args: + text: Arbitrary text that may embed an ``org-*`` identifier. + + Returns: + The text with every organization ID replaced. Non-string input is + returned unchanged after coercion to ``str``. + """ + if not text: + return text + return _ORG_ID_PATTERN.sub(ORG_ID_PLACEHOLDER, text) + + +def redact_obj(obj: Any) -> Any: + """Recursively redact organization IDs in arbitrary nested structures. + + Strings are redacted in place; dicts and lists/tuples are walked so that + identifiers nested inside structured payloads (e.g. log ``extra`` fields or + error ``details``) are also stripped. Other types are returned unchanged. + """ + if isinstance(obj, str): + return redact_org_ids(obj) + if isinstance(obj, dict): + return {key: redact_obj(value) for key, value in obj.items()} + if isinstance(obj, list): + return [redact_obj(value) for value in obj] + if isinstance(obj, tuple): + return tuple(redact_obj(value) for value in obj) + return obj diff --git a/routstr/payment/helpers.py b/routstr/payment/helpers.py index f333de84..afd71219 100644 --- a/routstr/payment/helpers.py +++ b/routstr/payment/helpers.py @@ -11,6 +11,8 @@ from PIL import Image from sqlmodel.ext.asyncio.session import AsyncSession from ..core import get_logger +from ..core.exceptions import UpstreamError +from ..core.redaction import redact_org_ids from ..core.settings import settings from ..wallet import deserialize_token_from_string @@ -418,16 +420,27 @@ def create_error_response( status_code: int, request: Request, token: str | None = None, + code: str | int | None = None, + details: dict[str, object] | None = None, ) -> Response: - """Create a standardized error response.""" + """Create a standardized error response. + + ``code`` is a stable, machine-readable classification (e.g. + ``UPSTREAM_RATE_LIMIT``); when omitted it defaults to the HTTP status code + for backwards compatibility. ``details`` carries optional structured, + redaction-safe context. + """ + error_obj: dict[str, object] = { + "message": redact_org_ids(message), + "type": error_type, + "code": code if code is not None else status_code, + } + if details is not None: + error_obj["details"] = details return Response( content=json.dumps( { - "error": { - "message": message, - "type": error_type, - "code": status_code, - }, + "error": error_obj, "request_id": getattr(request.state, "request_id", "unknown"), } ), @@ -435,3 +448,20 @@ def create_error_response( media_type="application/json", headers={"X-Cashu": token} if token else {}, ) + + +def create_upstream_error_response( + error: UpstreamError, + request: Request, + fallback_status: int = 502, +) -> Response: + """Build an error response from an :class:`UpstreamError`, preserving its + structured ``code``, ``details``, and original ``status_code``.""" + return create_error_response( + "upstream_error", + str(error), + error.status_code or fallback_status, + request=request, + code=getattr(error, "code", None), + details=getattr(error, "details", None), + ) diff --git a/routstr/proxy.py b/routstr/proxy.py index 7a74f682..8f6564e7 100644 --- a/routstr/proxy.py +++ b/routstr/proxy.py @@ -24,6 +24,7 @@ from .payment.helpers import ( calculate_discounted_max_cost, check_token_balance, create_error_response, + create_upstream_error_response, get_max_cost_for_model, ) from .payment.models import Model @@ -211,9 +212,7 @@ async def proxy( e, ) if i == len(all_upstreams) - 1: - last_error_response = create_error_response( - "upstream_error", str(e), 502, request=request - ) + last_error_response = create_upstream_error_response(e, request) continue return last_error_response or create_error_response( "upstream_error", "All upstreams failed", 502, request=request @@ -283,11 +282,10 @@ async def proxy( last_error = e continue + if last_error is not None: + return create_upstream_error_response(last_error, request) return create_error_response( - "upstream_error", - str(last_error) if last_error else "All upstreams failed", - 502, - request=request, + "upstream_error", "All upstreams failed", 502, request=request ) elif auth := headers.get("authorization", None): @@ -343,9 +341,7 @@ async def proxy( except UpstreamError as e: logger.warning(f"Upstream {upstream.provider_type} failed (GET): {e}") if i == len(upstreams) - 1: - last_error_response = create_error_response( - "upstream_error", str(e), 502, request=request - ) + last_error_response = create_upstream_error_response(e, request) continue return last_error_response or create_error_response( "upstream_error", "All upstreams failed", 502, request=request @@ -525,9 +521,7 @@ async def proxy( # If this was the last provider if i == len(upstreams) - 1: await revert_pay_for_request(key, session, max_cost_for_model) - return create_error_response( - "upstream_error", str(e), 502, request=request - ) + return create_upstream_error_response(e, request) # Otherwise loop continues to next provider continue diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index dd9dd429..2467044c 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import json +import math import traceback import uuid from collections.abc import AsyncGenerator, AsyncIterator, Iterator @@ -23,6 +24,7 @@ from ..core.db import ( store_cashu_transaction, ) from ..core.exceptions import UpstreamError +from ..core.redaction import redact_org_ids from ..payment.cost_calculation import ( CostData, CostDataError, @@ -47,6 +49,7 @@ from .cache_breakpoints import ( ) from .count_tokens import count_tokens_locally from .litellm_routing import detect_litellm_prefix +from .rate_limit import UPSTREAM_RATE_LIMIT, classify_rate_limit logger = get_logger(__name__) @@ -582,7 +585,7 @@ class BaseUpstreamProvider: preview = body_bytes.decode("utf-8", errors="ignore").strip() if preview: message = preview[:500] - return message, upstream_code + return redact_org_ids(message), upstream_code async def on_upstream_error_redirect( self, status_code: int, error_message: str @@ -625,10 +628,23 @@ class BaseUpstreamProvider: body_bytes = b"" body_read_error = f"{type(exc).__name__}: {exc}" + # ``message`` is already redacted by ``_extract_upstream_error_message``; + # the raw body preview is redacted here before it reaches logs or the + # forwarded envelope so provider account identifiers never leak. message, upstream_code = self._extract_upstream_error_message(body_bytes) - body_preview = body_bytes.decode("utf-8", errors="ignore").strip()[:500] + body_preview = redact_org_ids( + body_bytes.decode("utf-8", errors="ignore").strip()[:500] + ) is_json_body = _is_json_content_type(content_type) + # Classify upstream rate-limit failures into a stable, structured error. + rate_limit = classify_rate_limit(status_code, message, headers) + error_code: str | int = upstream_code or status_code + error_details: dict[str, object] | None = None + if rate_limit is not None: + error_code = UPSTREAM_RATE_LIMIT + error_details = rate_limit.as_details() + logger.warning( "Upstream %s returned %s for model=%s path=%s: %s", self.provider_type, @@ -642,6 +658,7 @@ class BaseUpstreamProvider: "model": model_id or "unknown", "upstream_status": status_code, "upstream_code": upstream_code, + "error_code": error_code, "upstream_content_type": content_type, "upstream_request_id": upstream_request_id, "message_preview": message[:200], @@ -676,13 +693,41 @@ class BaseUpstreamProvider: ): headers.pop(header_name, None) + # Propagate a usable retry hint to the caller when the upstream supplied + # one but did not echo a ``Retry-After`` header. RFC 7231 delta-seconds + # is an integer, so round sub-second hints up to a usable ``1``. + if ( + rate_limit is not None + and rate_limit.retry_after_seconds is not None + and "retry-after" not in {k.lower() for k in headers} + ): + headers["Retry-After"] = str(max(1, math.ceil(rate_limit.retry_after_seconds))) + if is_json_body: if not content_type: headers.pop("content-type", None) headers.pop("Content-Type", None) media_type = content_type or None + # Re-serialise the body with organization IDs stripped. The narrow + # ``org-*`` regex preserves the surrounding JSON structure. + redacted_text = redact_org_ids(body_bytes.decode("utf-8", errors="ignore")) + redacted_body = redacted_text.encode() + # Surface the stable rate-limit classification on the forwarded + # body so callers can switch on ``error.code`` without parsing the + # provider-specific message. Fall back to the redacted bytes if the + # body is not a JSON object with an ``error`` mapping. + if rate_limit is not None: + try: + parsed = json.loads(redacted_text) + err = parsed.get("error") if isinstance(parsed, dict) else None + if isinstance(err, dict): + err["code"] = UPSTREAM_RATE_LIMIT + err["details"] = error_details + redacted_body = json.dumps(parsed).encode() + except (ValueError, AttributeError): + pass return Response( - content=body_bytes, + content=redacted_body, status_code=status_code, headers=headers, media_type=media_type, @@ -693,15 +738,18 @@ class BaseUpstreamProvider: for header_name in ("content-type", "Content-Type"): headers.pop(header_name, None) + error_obj: dict[str, object] = { + "message": message or "Upstream returned a non-JSON error response", + "type": "upstream_error", + "code": error_code, + "upstream_status": status_code, + "upstream_content_type": content_type or None, + "upstream_body_preview": body_preview or None, + } + if error_details is not None: + error_obj["details"] = error_details envelope = { - "error": { - "message": message or "Upstream returned a non-JSON error response", - "type": "upstream_error", - "code": upstream_code or status_code, - "upstream_status": status_code, - "upstream_content_type": content_type or None, - "upstream_body_preview": body_preview or None, - }, + "error": error_obj, "request_id": getattr(request.state, "request_id", None), } @@ -2521,9 +2569,16 @@ class BaseUpstreamProvider: body_bytes = await response.aread() except Exception: body_bytes = b"" - body_preview = body_bytes.decode( - "utf-8", errors="ignore" - ).strip()[:500] + # Redact provider account identifiers before the body text + # reaches logs or the raised error. + body_preview = redact_org_ids( + body_bytes.decode("utf-8", errors="ignore").strip()[:500] + ) + rate_limit = classify_rate_limit( + response.status_code, + body_preview, + dict(response.headers), + ) logger.error( "Upstream %s returned %s for model=%s path=%s: %s", self.provider_type, @@ -2535,6 +2590,7 @@ class BaseUpstreamProvider: "provider": self.provider_type, "model": original_model_id or "unknown", "status_code": response.status_code, + "error_code": rate_limit.code if rate_limit else None, "reason_phrase": response.reason_phrase, "path": path, "body_preview": body_preview, @@ -2547,6 +2603,8 @@ class BaseUpstreamProvider: f"for model {original_model_id or 'unknown'}: " f"{body_preview[:200] or ''}", status_code=response.status_code, + code=rate_limit.code if rate_limit else None, + details=rate_limit.as_details() if rate_limit else None, ) try: @@ -2843,9 +2901,16 @@ class BaseUpstreamProvider: body_bytes = await response.aread() except Exception: body_bytes = b"" - body_preview = body_bytes.decode( - "utf-8", errors="ignore" - ).strip()[:500] + # Redact provider account identifiers before the body text + # reaches logs or the raised error. + body_preview = redact_org_ids( + body_bytes.decode("utf-8", errors="ignore").strip()[:500] + ) + rate_limit = classify_rate_limit( + response.status_code, + body_preview, + dict(response.headers), + ) logger.error( "Upstream %s returned %s for model=%s path=%s: %s", self.provider_type, @@ -2857,6 +2922,7 @@ class BaseUpstreamProvider: "provider": self.provider_type, "model": original_model_id or "unknown", "status_code": response.status_code, + "error_code": rate_limit.code if rate_limit else None, "path": path, "body_preview": body_preview, }, @@ -2868,6 +2934,8 @@ class BaseUpstreamProvider: f"for model {original_model_id or 'unknown'}: " f"{body_preview[:200] or ''}", status_code=response.status_code, + code=rate_limit.code if rate_limit else None, + details=rate_limit.as_details() if rate_limit else None, ) try: diff --git a/routstr/upstream/messages_dispatch.py b/routstr/upstream/messages_dispatch.py index ac8fb2fc..3d689922 100644 --- a/routstr/upstream/messages_dispatch.py +++ b/routstr/upstream/messages_dispatch.py @@ -29,7 +29,9 @@ import litellm from ..core import get_logger from ..core.exceptions import UpstreamError +from ..core.redaction import redact_org_ids from ..payment.models import Model +from .rate_limit import classify_rate_limit logger = get_logger(__name__) @@ -505,23 +507,33 @@ async def dispatch_anthropic_messages( try: result = await litellm.anthropic.messages.acreate(**kwargs) except Exception as exc: - exc_message = getattr(exc, "message", None) or str(exc) or repr(exc) + raw_message = getattr(exc, "message", None) or str(exc) or repr(exc) + # Redact provider account identifiers before the message reaches logs + # or the surfaced error. + exc_message = redact_org_ids(raw_message) exc_status = getattr(exc, "status_code", None) exc_response = getattr(exc, "response", None) response_text = None if exc_response is not None: try: - response_text = getattr(exc_response, "text", str(exc_response)) + response_text = redact_org_ids( + getattr(exc_response, "text", str(exc_response)) + ) except Exception: response_text = "" + status_for_classify = exc_status if isinstance(exc_status, int) else 502 + rate_limit = classify_rate_limit( + status_for_classify, exc_message, getattr(exc, "headers", None) + ) logger.error( "litellm dispatch failed", extra={ "error": exc_message, "error_type": type(exc).__name__, "status_code": exc_status, + "error_code": rate_limit.code if rate_limit else None, "llm_provider": getattr(exc, "llm_provider", None), - "body": getattr(exc, "body", None), + "body": redact_org_ids(str(getattr(exc, "body", "") or "")) or None, "response_text": response_text, "model": litellm_model, "api_base": base_url, @@ -529,7 +541,9 @@ async def dispatch_anthropic_messages( ) raise UpstreamError( f"Upstream error via litellm: {exc_message}", - status_code=exc_status if isinstance(exc_status, int) else 502, + status_code=status_for_classify, + code=rate_limit.code if rate_limit else None, + details=rate_limit.as_details() if rate_limit else None, ) from exc if not client_stream and hasattr(result, "__aiter__"): diff --git a/routstr/upstream/rate_limit.py b/routstr/upstream/rate_limit.py new file mode 100644 index 00000000..dca1ba5b --- /dev/null +++ b/routstr/upstream/rate_limit.py @@ -0,0 +1,136 @@ +"""Detection and parsing of upstream provider rate-limit errors. + +Upstream OpenAI-compatible providers signal rate limits via HTTP 429 and/or a +human-readable message such as:: + + Rate limit reached for gpt-5.5-2026-04-23 (for limit gpt-5.5) in organization + org-XXXX on tokens per min (TPM): Limit 180000000, Used 180000000, + Requested 8929. Please try again in 2ms. + +This module classifies those failures into a stable :data:`UPSTREAM_RATE_LIMIT` +code and extracts useful debugging fields. All retained text is redacted of +organization IDs first. +""" + +from __future__ import annotations + +import re +from dataclasses import asdict, dataclass + +from ..core.redaction import redact_org_ids + +# Stable error code callers can switch on to distinguish upstream rate limits +# from generic request failures. The literal value matches the identifier named +# in issue #555 ("UPSTREAM_RATE_LIMIT") so the public API contract is exact. +UPSTREAM_RATE_LIMIT = "UPSTREAM_RATE_LIMIT" + +# Message fragments that indicate a rate-limit even when the status code is not +# 429 (some providers wrap it in a 400/500 envelope). +_RATE_LIMIT_MARKERS = ( + "rate limit reached", + "rate_limit_exceeded", + "rate limit exceeded", + "too many requests", +) + +_MODEL_RE = re.compile(r"Rate limit reached for ([^\s(]+)", re.IGNORECASE) +_LIMIT_NAME_RE = re.compile(r"\(for limit ([^)]+)\)", re.IGNORECASE) +_METRIC_RE = re.compile(r"on ([a-z ]+\((?:TPM|RPM|TPD|RPD|IPM)\))", re.IGNORECASE) +_LIMIT_RE = re.compile(r"Limit (\d+)", re.IGNORECASE) +_USED_RE = re.compile(r"Used (\d+)", re.IGNORECASE) +_REQUESTED_RE = re.compile(r"Requested (\d+)", re.IGNORECASE) +_RETRY_RE = re.compile(r"try again in ([\d.]+)\s*(ms|s)", re.IGNORECASE) + + +@dataclass +class RateLimitInfo: + """Structured, redaction-safe view of an upstream rate-limit error.""" + + code: str + message: str + model: str | None = None + limit_name: str | None = None + metric: str | None = None + limit: int | None = None + used: int | None = None + requested: int | None = None + retry_after_seconds: float | None = None + + def as_details(self) -> dict[str, object]: + """Return a JSON-serialisable dict for embedding in an error envelope.""" + return {k: v for k, v in asdict(self).items() if v is not None} + + +def _looks_like_rate_limit(status_code: int, message: str) -> bool: + if status_code == 429: + return True + lowered = message.lower() + return any(marker in lowered for marker in _RATE_LIMIT_MARKERS) + + +def _parse_retry_after_header(headers: dict[str, str] | None) -> float | None: + """Parse a ``Retry-After`` header (delta-seconds form) into seconds.""" + if not headers: + return None + raw = headers.get("retry-after") or headers.get("Retry-After") + if raw is None: + return None + try: + return float(str(raw).strip()) + except (TypeError, ValueError): + return None + + +def _int_or_none(match: re.Match[str] | None) -> int | None: + if match is None: + return None + try: + return int(match.group(1)) + except (TypeError, ValueError): + return None + + +def classify_rate_limit( + status_code: int, + message: str, + headers: dict[str, str] | None = None, +) -> RateLimitInfo | None: + """Classify an upstream error as a rate-limit and extract its fields. + + Args: + status_code: HTTP status code from the upstream response. + message: Upstream error message (may contain sensitive identifiers). + headers: Optional upstream response headers, used for ``Retry-After``. + + Returns: + A :class:`RateLimitInfo` when the error is a rate-limit, else ``None``. + """ + message = message or "" + if not _looks_like_rate_limit(status_code, message): + return None + + redacted = redact_org_ids(message) + + model_match = _MODEL_RE.search(redacted) + metric_match = _METRIC_RE.search(redacted) + + retry_after = _parse_retry_after_header(headers) + if retry_after is None: + retry_match = _RETRY_RE.search(redacted) + if retry_match is not None: + value = float(retry_match.group(1)) + retry_after = value / 1000.0 if retry_match.group(2).lower() == "ms" else value + + limit_name_match = _LIMIT_NAME_RE.search(redacted) + + return RateLimitInfo( + code=UPSTREAM_RATE_LIMIT, + message=redacted, + model=model_match.group(1) if model_match else None, + limit_name=limit_name_match.group(1).strip() if limit_name_match else None, + metric=metric_match.group(1).strip() if metric_match else None, + limit=_int_or_none(_LIMIT_RE.search(redacted)), + used=_int_or_none(_USED_RE.search(redacted)), + requested=_int_or_none(_REQUESTED_RE.search(redacted)), + retry_after_seconds=retry_after, + ) diff --git a/tests/unit/test_upstream_rate_limit.py b/tests/unit/test_upstream_rate_limit.py new file mode 100644 index 00000000..ea336c14 --- /dev/null +++ b/tests/unit/test_upstream_rate_limit.py @@ -0,0 +1,396 @@ +"""Tests for upstream rate-limit detection, classification, and org-ID redaction. + +Covers issue #555: upstream OpenAI-compatible providers return rate-limit +errors that embed a sensitive organization ID. The proxy must classify these +distinctly (``UPSTREAM_RATE_LIMIT``), preserve useful debugging fields, and +never emit a raw ``org-*`` identifier in logs, errors, or returned bodies. +""" + +from __future__ import annotations + +import json +from typing import Any +from unittest.mock import AsyncMock, MagicMock, Mock, patch + +import httpx +import pytest + +from routstr.core.redaction import redact_org_ids +from routstr.upstream.base import BaseUpstreamProvider +from routstr.upstream.rate_limit import ( + UPSTREAM_RATE_LIMIT, + RateLimitInfo, + classify_rate_limit, +) + +# The exact scenario from the issue, with a realistic (fake) org identifier. +RAW_ORG_ID = "org-abc123XYZ456def" +RATE_LIMIT_MESSAGE = ( + f"Rate limit reached for gpt-5.5-2026-04-23 (for limit gpt-5.5) in " + f"organization {RAW_ORG_ID} on tokens per min (TPM): Limit 180000000, " + f"Used 180000000, Requested 8929. Please try again in 2ms. Visit " + f"https://platform.openai.com/account/rate-limits to learn more." +) + + +def _make_request(request_id: str = "req-123") -> Mock: + request = Mock(spec=["method", "state"]) + request.method = "POST" + request.state = Mock() + request.state.request_id = request_id + return request + + +def _make_upstream_response( + *, + body: bytes, + status_code: int = 429, + content_type: str | None = "application/json", + extra_headers: dict[str, str] | None = None, +) -> httpx.Response: + headers: dict[str, str] = {} + if content_type is not None: + headers["content-type"] = content_type + if extra_headers: + headers.update(extra_headers) + return httpx.Response(status_code=status_code, headers=headers, content=body) + + +@pytest.fixture +def provider() -> BaseUpstreamProvider: + return BaseUpstreamProvider( + base_url="https://privateprovider.xyz", api_key="k", provider_fee=1.0 + ) + + +# --------------------------------------------------------------------------- # +# Redaction +# --------------------------------------------------------------------------- # + + +def test_redact_org_ids_replaces_identifier() -> None: + assert RAW_ORG_ID not in redact_org_ids(RATE_LIMIT_MESSAGE) + assert "org-[REDACTED]" in redact_org_ids(RATE_LIMIT_MESSAGE) + + +def test_redact_org_ids_is_idempotent() -> None: + once = redact_org_ids(RATE_LIMIT_MESSAGE) + assert redact_org_ids(once) == once + + +def test_redact_org_ids_leaves_unrelated_text() -> None: + assert redact_org_ids("organize the org-chart") == "organize the org-chart" + assert redact_org_ids("") == "" + + +# --------------------------------------------------------------------------- # +# Classification +# --------------------------------------------------------------------------- # + + +def test_classify_exact_scenario() -> None: + info = classify_rate_limit(429, RATE_LIMIT_MESSAGE) + assert isinstance(info, RateLimitInfo) + assert info.code == UPSTREAM_RATE_LIMIT + assert info.model == "gpt-5.5-2026-04-23" + assert info.limit_name == "gpt-5.5" + assert info.metric == "tokens per min (TPM)" + assert info.limit == 180000000 + assert info.used == 180000000 + assert info.requested == 8929 + assert info.retry_after_seconds == pytest.approx(0.002) + # Redaction-safe: no raw org id survives into the structured view. + assert RAW_ORG_ID not in info.message + assert RAW_ORG_ID not in json.dumps(info.as_details()) + + +def test_classify_by_status_code_without_marker() -> None: + info = classify_rate_limit(429, "slow down") + assert info is not None + assert info.code == UPSTREAM_RATE_LIMIT + + +def test_classify_by_message_marker_without_429() -> None: + info = classify_rate_limit(400, "rate_limit_exceeded for this key") + assert info is not None + + +def test_retry_after_header_takes_precedence() -> None: + info = classify_rate_limit(429, RATE_LIMIT_MESSAGE, {"Retry-After": "12"}) + assert info is not None + assert info.retry_after_seconds == pytest.approx(12.0) + + +def test_non_rate_limit_error_is_not_classified() -> None: + assert classify_rate_limit(400, "invalid request: missing field 'model'") is None + assert classify_rate_limit(500, "internal server error") is None + + +# --------------------------------------------------------------------------- # +# forward_upstream_error_response integration +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_json_rate_limit_body_is_redacted_and_forwarded( + provider: BaseUpstreamProvider, +) -> None: + body = json.dumps( + {"error": {"message": RATE_LIMIT_MESSAGE, "type": "rate_limit_exceeded"}} + ).encode() + upstream = _make_upstream_response(body=body, status_code=429) + + response = await provider.forward_upstream_error_response( + _make_request(), "v1/chat/completions", upstream + ) + + assert response.status_code == 429 + raw = bytes(response.body).decode() + # No raw organization id may survive in the forwarded body. + assert RAW_ORG_ID not in raw + assert "org-[REDACTED]" in raw + # Body remains valid JSON; the original type is preserved while a stable + # rate-limit code is injected so callers can switch on it. + payload: dict[str, Any] = json.loads(raw) + assert payload["error"]["type"] == "rate_limit_exceeded" + assert payload["error"]["code"] == UPSTREAM_RATE_LIMIT + assert payload["error"]["details"]["model"] == "gpt-5.5-2026-04-23" + # A retry hint extracted from the message is surfaced as a header. + assert "retry-after" in {k.lower() for k in response.headers} + + +@pytest.mark.asyncio +async def test_non_json_rate_limit_envelope_uses_stable_code( + provider: BaseUpstreamProvider, +) -> None: + upstream = _make_upstream_response( + body=RATE_LIMIT_MESSAGE.encode(), + status_code=429, + content_type="text/plain", + ) + + response = await provider.forward_upstream_error_response( + _make_request(), "v1/chat/completions", upstream + ) + + payload: dict[str, Any] = json.loads(bytes(response.body)) + assert payload["error"]["code"] == UPSTREAM_RATE_LIMIT + assert payload["error"]["details"]["model"] == "gpt-5.5-2026-04-23" + serialized = json.dumps(payload) + assert RAW_ORG_ID not in serialized + assert "org-[REDACTED]" in serialized + + +@pytest.mark.asyncio +async def test_non_rate_limit_json_error_unchanged( + provider: BaseUpstreamProvider, +) -> None: + body = json.dumps( + {"error": {"message": "missing field 'model'", "type": "invalid_request"}} + ).encode() + upstream = _make_upstream_response(body=body, status_code=400) + + response = await provider.forward_upstream_error_response( + _make_request(), "v1/chat/completions", upstream + ) + + assert response.status_code == 400 + payload: dict[str, Any] = json.loads(bytes(response.body)) + assert payload["error"]["type"] == "invalid_request" + assert "retry-after" not in {k.lower() for k in response.headers} + + +# --------------------------------------------------------------------------- # +# UpstreamError -> proxy response (preserves code/details/status) +# --------------------------------------------------------------------------- # + + +def test_create_upstream_error_response_preserves_structure() -> None: + from routstr.core.exceptions import UpstreamError + from routstr.payment.helpers import create_upstream_error_response + + info = classify_rate_limit(429, RATE_LIMIT_MESSAGE) + assert info is not None + err = UpstreamError( + f"Upstream error via litellm: {RATE_LIMIT_MESSAGE}", + status_code=429, + code=info.code, + details=info.as_details(), + ) + + response = create_upstream_error_response(err, _make_request()) + + # Original upstream status is preserved (not flattened to 502). + assert response.status_code == 429 + payload: dict[str, Any] = json.loads(bytes(response.body)) + assert payload["error"]["type"] == "upstream_error" + assert payload["error"]["code"] == UPSTREAM_RATE_LIMIT + assert payload["error"]["details"]["requested"] == 8929 + serialized = json.dumps(payload) + assert RAW_ORG_ID not in serialized + assert "org-[REDACTED]" in serialized + + +def test_generic_upstream_error_still_defaults_to_502() -> None: + from routstr.core.exceptions import UpstreamError + from routstr.payment.helpers import create_upstream_error_response + + err = UpstreamError("connection refused") # status_code defaults to 502 + + response = create_upstream_error_response(err, _make_request()) + + assert response.status_code == 502 + payload: dict[str, Any] = json.loads(bytes(response.body)) + assert payload["error"]["type"] == "upstream_error" + assert payload["error"]["code"] == 502 + assert "details" not in payload["error"] + + +# --------------------------------------------------------------------------- # +# Structured log-extra redaction +# --------------------------------------------------------------------------- # + + +def test_security_filter_redacts_org_id_in_extra() -> None: + import logging + + from routstr.core.logging import SecurityFilter + + record = logging.LogRecord( + name="test", + level=logging.ERROR, + pathname=__file__, + lineno=1, + msg="upstream failed", + args=(), + exc_info=None, + ) + # Simulate an ``extra={"body_preview": ...}`` field carrying an org id. + setattr(record, "body_preview", RATE_LIMIT_MESSAGE) + + assert SecurityFilter().filter(record) is True + redacted: str = getattr(record, "body_preview") + assert RAW_ORG_ID not in redacted + assert "org-[REDACTED]" in redacted + + +def test_security_filter_redacts_org_id_in_nested_extra() -> None: + import logging + + from routstr.core.logging import SecurityFilter + + record = logging.LogRecord( + name="test", + level=logging.ERROR, + pathname=__file__, + lineno=1, + msg="upstream failed", + args=(), + exc_info=None, + ) + # Nested structures: dict containing a list containing the org id. + setattr(record, "body", {"error": {"messages": [RATE_LIMIT_MESSAGE]}}) + + assert SecurityFilter().filter(record) is True + serialized = json.dumps(getattr(record, "body")) + assert RAW_ORG_ID not in serialized + assert "org-[REDACTED]" in serialized + + +# --------------------------------------------------------------------------- # +# 5xx-wrapped rate limit through forward_upstream_error_response +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_5xx_wrapped_rate_limit_is_classified( + provider: BaseUpstreamProvider, +) -> None: + # Some providers wrap a rate-limit in a 5xx envelope; classification must + # key off the message marker, not only the 429 status. + body = json.dumps({"error": {"message": RATE_LIMIT_MESSAGE}}).encode() + upstream = _make_upstream_response(body=body, status_code=500) + + response = await provider.forward_upstream_error_response( + _make_request(), "v1/chat/completions", upstream + ) + + assert response.status_code == 500 + payload: dict[str, Any] = json.loads(bytes(response.body)) + assert payload["error"]["code"] == UPSTREAM_RATE_LIMIT + serialized = json.dumps(payload) + assert RAW_ORG_ID not in serialized + assert "org-[REDACTED]" in serialized + + +# --------------------------------------------------------------------------- # +# Real proxy loop: structured error surfaced + reservation reverted once +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_proxy_loop_surfaces_rate_limit_and_reverts_once() -> None: + from routstr import proxy as proxy_module + from routstr.core.db import ApiKey + from routstr.core.exceptions import UpstreamError + + info = classify_rate_limit(429, RATE_LIMIT_MESSAGE) + assert info is not None + + key = ApiKey(hashed_key="rlkey", balance=10_000) + + request = MagicMock() + request.method = "POST" + request.headers = {"authorization": "Bearer sk-rlkey"} + request.body = AsyncMock(return_value=b'{"model": "test-model"}') + request.state = MagicMock() + request.state.request_id = "req-rl" + + upstream = MagicMock() + upstream.provider_type = "test" + upstream.prepare_headers = MagicMock(side_effect=lambda h: h) + upstream.forward_request = AsyncMock( + side_effect=UpstreamError( + f"Upstream error via litellm: {RATE_LIMIT_MESSAGE}", + status_code=429, + code=info.code, + details=info.as_details(), + ) + ) + + session = MagicMock() + revert_mock = AsyncMock(return_value=True) + + with ( + patch.object(proxy_module, "get_model_instance", return_value=MagicMock()), + patch.object(proxy_module, "get_provider_for_model", return_value=[upstream]), + patch.object( + proxy_module, "get_max_cost_for_model", AsyncMock(return_value=1_000) + ), + patch.object( + proxy_module, + "calculate_discounted_max_cost", + AsyncMock(return_value=1_000), + ), + patch.object(proxy_module, "check_token_balance", MagicMock()), + patch.object( + proxy_module, "get_bearer_token_key", AsyncMock(return_value=key) + ), + patch.object(proxy_module, "pay_for_request", AsyncMock(return_value=1_000)), + patch.object(proxy_module, "revert_pay_for_request", revert_mock), + ): + response = await proxy_module.proxy( + request, "v1/chat/completions", session=session + ) + + # Original 429 status and the stable code/details survive to the client. + assert response.status_code == 429 + payload: dict[str, Any] = json.loads(bytes(response.body)) + assert payload["error"]["type"] == "upstream_error" + assert payload["error"]["code"] == UPSTREAM_RATE_LIMIT + assert payload["error"]["details"]["requested"] == 8929 + serialized = json.dumps(payload) + assert RAW_ORG_ID not in serialized + assert "org-[REDACTED]" in serialized + # Single upstream failed -> reservation reverted exactly once (no double-charge). + revert_mock.assert_awaited_once_with(key, session, 1_000)