better rate limit forwarding

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)
 
 

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
 

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-<base62>``. 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

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": {
+                "error": error_obj,
-                    "message": message,
-                    "type": error_type,
-                    "code": status_code,
-                },
                 "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),
+    )

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(
+                    last_error_response = create_upstream_error_response(e, request)
-                        "upstream_error", str(e), 502, request=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",
+            "upstream_error", "All upstreams failed", 502, request=request
-            str(last_error) if last_error else "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(
+                    last_error_response = create_upstream_error_response(e, request)
-                        "upstream_error", str(e), 502, request=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(
+                return create_upstream_error_response(e, request)
-                    "upstream_error", str(e), 502, request=request
-                )
 
             # Otherwise loop continues to next provider
             continue

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": {
+            "error": error_obj,
-                "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,
-            },
             "request_id": getattr(request.state, "request_id", None),
         }
 
@@ -2494,9 +2542,16 @@ class BaseUpstreamProvider:
                         body_bytes = await response.aread()
                     except Exception:
                         body_bytes = b""
-                    body_preview = body_bytes.decode(
+                    # Redact provider account identifiers before the body text
-                        "utf-8", errors="ignore"
+                    # reaches logs or the raised error.
-                    ).strip()[:500]
+                    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,
@@ -2508,6 +2563,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,
@@ -2520,6 +2576,8 @@ class BaseUpstreamProvider:
                         f"for model {original_model_id or 'unknown'}: "
                         f"{body_preview[:200] or '<empty>'}",
                         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:
@@ -2816,9 +2874,16 @@ class BaseUpstreamProvider:
                         body_bytes = await response.aread()
                     except Exception:
                         body_bytes = b""
-                    body_preview = body_bytes.decode(
+                    # Redact provider account identifiers before the body text
-                        "utf-8", errors="ignore"
+                    # reaches logs or the raised error.
-                    ).strip()[:500]
+                    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,
@@ -2830,6 +2895,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,
                         },
@@ -2841,6 +2907,8 @@ class BaseUpstreamProvider:
                         f"for model {original_model_id or 'unknown'}: "
                         f"{body_preview[:200] or '<empty>'}",
                         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:

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 = "<unreadable>"
+        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__"):

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,
+    )

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)