codexLB-up 1.25.0b10__py3-none-any.whl

app/core/middleware/inflight.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+from importlib import import_module
+from types import ModuleType
+
+from starlette.responses import JSONResponse
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from app.core.resilience.overload import (
+    deny_websocket_with_http_response,
+    is_proxy_path,
+    local_unavailable_error,
+    merge_retry_after_headers,
+)
+
+_DRAIN_MESSAGE = "Server is draining"
+
+_DRAIN_ALLOWED_HTTP_PATHS = frozenset(
+    {
+        "/health/live",
+        "/internal/drain/start",
+        "/internal/drain/stop",
+        "/internal/drain/status",
+        "/internal/bridge/responses",
+    }
+)
+
+_IN_FLIGHT_EXCLUDED_HTTP_PATHS = frozenset(
+    {
+        "/health/live",
+        "/health/ready",
+        "/health/startup",
+        "/internal/drain/start",
+        "/internal/drain/stop",
+        "/internal/drain/status",
+    }
+)
+
+_IN_FLIGHT_WEBSOCKET_PATHS = frozenset(
+    {
+        "/backend-api/codex/responses",
+        "/v1/responses",
+    }
+)
+
+
+class InFlightMiddleware:
+    def __init__(self, app: ASGIApp) -> None:
+        self.app = app
+        # Resolved lazily on the first request (import-cycle safety) and cached
+        # so the hot path skips the per-request sys.modules lookup.
+        self._shutdown_state: ModuleType | None = None
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        """Track admitted work and reject new requests while the server drains."""
+        scope_type = scope["type"]
+        if scope_type not in {"http", "websocket"}:
+            await self.app(scope, receive, send)
+            return
+
+        shutdown_state = self._shutdown_state
+        if shutdown_state is None:
+            shutdown_state = self._shutdown_state = import_module("app.core.shutdown")
+
+        if scope_type == "websocket":
+            # Register before checking the drain barrier. A synchronous signal
+            # can interleave anywhere between Python bytecodes; increment-first
+            # guarantees that the scope is either rejected and unregistered or
+            # already visible to shutdown as pre-barrier work.
+            shutdown_state.increment_in_flight()
+            if shutdown_state.is_draining():
+                shutdown_state.decrement_in_flight()
+                # ``websocket.close`` before the handshake is surfaced by ASGI
+                # servers as an opaque HTTP 403, which SDK clients treat as a
+                # terminal access error. Deny the upgrade with the same
+                # retryable 503 + Retry-After the overload bulkhead uses so the
+                # drain reads as transient on every client.
+                path = scope.get("path", "")
+                await deny_websocket_with_http_response(
+                    receive,
+                    send,
+                    status_code=503,
+                    payload=(
+                        local_unavailable_error(_DRAIN_MESSAGE) if is_proxy_path(path) else {"detail": _DRAIN_MESSAGE}
+                    ),
+                    headers=merge_retry_after_headers(),
+                )
+                return
+
+            if scope.get("path", "") not in _IN_FLIGHT_WEBSOCKET_PATHS:
+                # Admission is global, but only Responses sockets participate
+                # in turn-aware drain. Other long-lived protocols are closed by
+                # Uvicorn after the pre-connection drain barrier.
+                shutdown_state.decrement_in_flight()
+                await self.app(scope, receive, send)
+                return
+
+            try:
+                await self.app(scope, receive, send)
+            finally:
+                shutdown_state.decrement_in_flight()
+            return
+
+        path = scope.get("path", "")
+        track_in_flight = path not in _IN_FLIGHT_EXCLUDED_HTTP_PATHS
+        if track_in_flight:
+            # As with Responses WebSockets, register before checking the
+            # barrier so a synchronous signal cannot admit invisible work.
+            shutdown_state.increment_in_flight()
+        try:
+            if shutdown_state.is_draining() and path not in _DRAIN_ALLOWED_HTTP_PATHS:
+                response = JSONResponse(
+                    status_code=503,
+                    content={
+                        "error": {
+                            "type": "service_unavailable",
+                            "message": _DRAIN_MESSAGE,
+                        }
+                    },
+                )
+                await response(scope, receive, send)
+                return
+            await self.app(scope, receive, send)
+        finally:
+            if track_in_flight:
+                shutdown_state.decrement_in_flight()

app/core/middleware/multipart_content_encoding.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import logging
+
+from fastapi import FastAPI
+from starlette._utils import get_route_path
+from starlette.datastructures import Headers
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from app.core.errors import dashboard_error, openai_error
+from app.core.runtime_logging import log_error_response
+
+_LIMITED_MULTIPART_PATHS = frozenset(
+    {
+        "/api/accounts/import",
+        "/backend-api/transcribe",
+        "/v1/audio/transcriptions",
+        "/v1/images/edits",
+    }
+)
+_OPENAI_MULTIPART_PATH_PREFIXES = ("/backend-api", "/v1")
+_CONTENT_ENCODING_GATE_HANDLED_STATE = "_codex_lb_multipart_content_encoding_gate_handled"
+_UNSUPPORTED_CONTENT_ENCODING_STATE = "_codex_lb_unsupported_multipart_content_encoding"
+
+logger = logging.getLogger(__name__)
+
+
+def _canonical_route_path(scope: Scope) -> str:
+    path = get_route_path(scope)
+    if path != "/":
+        path = path.rstrip("/")
+    return path
+
+
+class UnsupportedMultipartContentEncoding(Exception):
+    pass
+
+
+def is_route_owned_multipart_operation(scope: Scope) -> bool:
+    return scope.get("method") == "POST" and _canonical_route_path(scope) in _LIMITED_MULTIPART_PATHS
+
+
+def _uses_openai_multipart_errors(scope: Scope) -> bool:
+    path = _canonical_route_path(scope)
+    return any(path == prefix or path.startswith(f"{prefix}/") for prefix in _OPENAI_MULTIPART_PATH_PREFIXES)
+
+
+def multipart_error_response(
+    request: Request,
+    *,
+    status_code: int,
+    code: str,
+    message: str,
+    param: str | None = None,
+) -> JSONResponse:
+    uses_openai_errors = _uses_openai_multipart_errors(request.scope)
+    response_code = "invalid_request_error" if uses_openai_errors and code == "invalid_request" else code
+    log_error_response(
+        logger,
+        request,
+        status_code,
+        response_code,
+        message,
+        category="openai_error_response" if uses_openai_errors else "dashboard_error_response",
+    )
+    if uses_openai_errors:
+        error = openai_error(response_code, message, error_type="invalid_request_error")
+        if param is not None:
+            error["error"]["param"] = param
+        return JSONResponse(status_code=status_code, content=error)
+    return JSONResponse(status_code=status_code, content=dashboard_error(response_code, message))
+
+
+def _without_content_encoding(scope: Scope, *, unsupported: bool = False) -> Scope:
+    copied = dict(scope)
+    copied["headers"] = [
+        (name, value) for name, value in scope.get("headers", []) if name.lower() != b"content-encoding"
+    ]
+    state = dict(scope.get("state", {}))
+    state[_CONTENT_ENCODING_GATE_HANDLED_STATE] = True
+    if unsupported:
+        state[_UNSUPPORTED_CONTENT_ENCODING_STATE] = True
+    copied["state"] = state
+    return copied
+
+
+def multipart_content_encoding_gate_was_applied(scope: Scope) -> bool:
+    return scope.get("state", {}).get(_CONTENT_ENCODING_GATE_HANDLED_STATE) is True
+
+
+def raise_for_unsupported_multipart_content_encoding(request: Request) -> None:
+    if getattr(request.state, _UNSUPPORTED_CONTENT_ENCODING_STATE, False) is True:
+        raise UnsupportedMultipartContentEncoding
+
+
+class MultipartContentEncodingMiddleware:
+    def __init__(self, app: ASGIApp) -> None:
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        headers = Headers(scope=scope)
+        if not is_route_owned_multipart_operation(scope):
+            await self.app(scope, receive, send)
+            return
+
+        content_encoding_values = headers.getlist("content-encoding")
+        if not content_encoding_values:
+            await self.app(scope, receive, send)
+            return
+        encodings = [
+            encoding.strip().lower()
+            for value in content_encoding_values
+            for encoding in value.split(",")
+            if encoding.strip()
+        ]
+        if not encodings or all(value == "identity" for value in encodings):
+            await self.app(_without_content_encoding(scope), receive, send)
+            return
+        await self.app(_without_content_encoding(scope, unsupported=True), receive, send)
+
+
+def add_multipart_content_encoding_middleware(app: FastAPI) -> None:
+    app.add_middleware(MultipartContentEncodingMiddleware)

app/core/middleware/path_rewrite.py

@@ -0,0 +1,148 @@
+"""Redact Live WebSocket server scopes and rewrite duplicated Codex paths.
+
+Some OpenAI-compatible clients unconditionally append ``/v1/`` to a configured
+``/backend-api/codex`` base URL. For non-Live requests this middleware collapses
+``/backend-api/codex/v1/<rest>`` to the canonical
+``/backend-api/codex/<rest>`` route for both HTTP and WebSocket scopes.
+
+Live WebSocket handshakes require the inverse ownership rule: Uvicorn retains
+the original ASGI scope for accepted and rejected handshake logs, while the
+application still needs the original path and query for routing. The middleware
+therefore copies the original routing scope for downstream use, then redacts the
+server-owned scope before any downstream accept or rejection can be emitted.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Any
+
+from fastapi import FastAPI
+from starlette.types import Receive, Scope, Send
+
+from app.core.clients.proxy_websocket import REALTIME_LIVE_CALL_ID_ROUTE_REGEX
+
+# Alias rewriting stays scoped to the duplicated Codex prefix. Live WebSocket
+# redaction separately covers every routed ingress family plus the unsupported
+# duplicated Live alias.
+_CODEX_V1_PREFIX = "/backend-api/codex/v1/"
+_CODEX_V1_PREFIX_BYTES = _CODEX_V1_PREFIX.encode("ascii")
+_CODEX_CANONICAL_PREFIX = "/backend-api/codex/"
+_CODEX_CANONICAL_PREFIX_BYTES = _CODEX_CANONICAL_PREFIX.encode("ascii")
+_V1_LIVE_PREFIX = "/v1/live/"
+_V1_LIVE_PREFIX_BYTES = _V1_LIVE_PREFIX.encode("ascii")
+_V1_REALTIME_PATH = "/v1/realtime"
+_V1_REALTIME_RAW_PATH = _V1_REALTIME_PATH.encode("ascii")
+_REDACTED_CALL_ID = "<redacted>"
+_REDACTED_CALL_ID_BYTES = b"%3Credacted%3E"
+_REALTIME_LIVE_CALL_ID_PATTERN = re.compile(rf"{REALTIME_LIVE_CALL_ID_ROUTE_REGEX}\Z")
+_REALTIME_LIVE_PATH_REDACTIONS = (
+    (
+        _CODEX_V1_PREFIX,
+        f"{_CODEX_V1_PREFIX}{_REDACTED_CALL_ID}",
+        _CODEX_V1_PREFIX_BYTES + _REDACTED_CALL_ID_BYTES,
+    ),
+    (
+        _CODEX_CANONICAL_PREFIX,
+        f"{_CODEX_CANONICAL_PREFIX}{_REDACTED_CALL_ID}",
+        _CODEX_CANONICAL_PREFIX_BYTES + _REDACTED_CALL_ID_BYTES,
+    ),
+)
+
+
+def _canonicalize_backend_api_codex_path(path: str) -> str:
+    """Collapse ``/backend-api/codex/v1/<rest>`` -> ``/backend-api/codex/<rest>``.
+
+    Returns the input unchanged for any path that is not the duplicated
+    Codex ``/v1/`` shape. In particular, ``/backend-api/codex`` (no
+    rest) and ``/backend-api/codex/v1`` (no further rest) are left
+    alone -- those are legal request paths a future contributor might
+    register, and collapsing them would silently change routing
+    semantics.
+    """
+    if not path.startswith(_CODEX_V1_PREFIX):
+        return path
+    return _CODEX_CANONICAL_PREFIX + path[len(_CODEX_V1_PREFIX) :]
+
+
+def _realtime_live_scope_redaction(path: str) -> tuple[str, bytes] | None:
+    if path == _V1_REALTIME_PATH:
+        return _V1_REALTIME_PATH, _V1_REALTIME_RAW_PATH
+
+    # The v3 ingress owns this whole namespace, including requests whose
+    # suffix cannot match a route. Rejected handshakes must therefore redact
+    # malformed, empty, and overlong suffixes before Uvicorn can log them.
+    if path.startswith(_V1_LIVE_PREFIX):
+        return (
+            f"{_V1_LIVE_PREFIX}{_REDACTED_CALL_ID}",
+            _V1_LIVE_PREFIX_BYTES + _REDACTED_CALL_ID_BYTES,
+        )
+
+    # Generic Codex paths share their namespace with non-Live routes, so their
+    # redaction boundary remains the exact Live call-id grammar.
+    for prefix, redacted_path, redacted_raw_path in _REALTIME_LIVE_PATH_REDACTIONS:
+        if path.startswith(prefix) and _REALTIME_LIVE_CALL_ID_PATTERN.fullmatch(path[len(prefix) :]) is not None:
+            return redacted_path, redacted_raw_path
+    return None
+
+
+def redact_realtime_live_path(path: str) -> str:
+    """Return the credential-safe path used by pre-routing diagnostics."""
+
+    redaction = _realtime_live_scope_redaction(path)
+    return path if redaction is None else redaction[0]
+
+
+def _canonicalize_raw_path(raw_path: bytes) -> bytes:
+    if not raw_path.startswith(_CODEX_V1_PREFIX_BYTES):
+        return raw_path
+    return _CODEX_CANONICAL_PREFIX_BYTES + raw_path[len(_CODEX_V1_PREFIX_BYTES) :]
+
+
+class BackendApiCodexV1AliasMiddleware:
+    """Keep Live handshake logs private and canonicalise the duplicated Codex prefix.
+
+    This runs immediately inside trusted-proxy projection. For a Live WebSocket,
+    the downstream application receives a copy containing the projected client
+    metadata and original routing values while the server-owned scope is
+    redacted in place. Lifespan and non-Live scopes preserve the existing
+    alias-only behavior.
+    """
+
+    def __init__(self, app: Any) -> None:
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        scope_type = scope.get("type")
+        path = scope.get("path")
+
+        if scope_type == "websocket" and isinstance(path, str):
+            redaction = _realtime_live_scope_redaction(path)
+            if redaction is not None:
+                routing_scope = dict(scope)
+                routing_scope["path"] = _canonicalize_backend_api_codex_path(path)
+                if isinstance(raw_path := routing_scope.get("raw_path"), bytes):
+                    routing_scope["raw_path"] = _canonicalize_raw_path(raw_path)
+                redacted_path, redacted_raw_path = redaction
+                scope["path"] = redacted_path
+                scope["raw_path"] = redacted_raw_path
+                scope["query_string"] = b""
+                await self.app(routing_scope, receive, send)
+                return
+
+        if scope_type in {"http", "websocket"} and isinstance(path, str) and path.startswith(_CODEX_V1_PREFIX):
+            rewritten = _canonicalize_backend_api_codex_path(path)
+            if rewritten != path:
+                # Preserve the server-owned scope when only downstream routing changes.
+                routing_scope = dict(scope)
+                routing_scope["path"] = rewritten
+                if isinstance(raw_path := routing_scope.get("raw_path"), bytes):
+                    routing_scope["raw_path"] = _canonicalize_raw_path(raw_path)
+                scope = routing_scope
+        await self.app(scope, receive, send)
+
+
+def add_backend_api_codex_v1_alias_middleware(app: FastAPI) -> None:
+    """Register path rewriting and Live scope redaction inside trusted proxying."""
+
+    app.add_middleware(BackendApiCodexV1AliasMiddleware)

app/core/middleware/request_body_limit.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+import logging
+
+from fastapi import FastAPI
+from starlette._utils import get_route_path
+from starlette.datastructures import Headers
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+from app.core.errors import SCIM_CONTENT_TYPE, dashboard_error, openai_error, scim_error
+from app.core.ingress_limits import MAX_DECOMPRESSED_BODY_BYTES, MAX_DECOMPRESSED_RESPONSES_BODY_BYTES
+from app.core.middleware.multipart_content_encoding import (
+    is_route_owned_multipart_operation,
+    multipart_content_encoding_gate_was_applied,
+)
+from app.core.runtime_logging import log_error_response
+
+REQUEST_BODY_TOO_LARGE_MESSAGE = "Request body exceeds the maximum allowed size"
+
+_RESPONSES_INGRESS_PATHS = frozenset(
+    {
+        "/backend-api/codex/responses",
+        "/v1/responses",
+    }
+)
+_OPENAI_INGRESS_PATH_PREFIXES = (
+    "/api/codex",
+    "/backend-api",
+    "/internal/bridge",
+    "/v1",
+)
+_REQUEST_BODY_TOO_LARGE_STATE = "_codex_lb_request_body_too_large"
+
+logger = logging.getLogger(__name__)
+
+
+class _RequestBodyTooLarge(Exception):
+    pass
+
+
+def request_body_limit_for_path(path: str) -> int:
+    if path.rstrip("/") in _RESPONSES_INGRESS_PATHS:
+        return max(MAX_DECOMPRESSED_BODY_BYTES, MAX_DECOMPRESSED_RESPONSES_BODY_BYTES)
+    return MAX_DECOMPRESSED_BODY_BYTES
+
+
+def _path_belongs_to(path: str, prefix: str) -> bool:
+    return path == prefix or path.startswith(f"{prefix}/")
+
+
+def _uses_openai_ingress_errors(path: str) -> bool:
+    return any(_path_belongs_to(path, prefix) for prefix in _OPENAI_INGRESS_PATH_PREFIXES)
+
+
+def _uses_scim_ingress_errors(path: str) -> bool:
+    """``/scim`` answers RFC 7644's envelope even before routing.
+
+    The ingress guard runs ahead of the router, so a refusal here would
+    otherwise be the dashboard's shape on a surface no dashboard client reads.
+    """
+
+    return _path_belongs_to(path, "/scim")
+
+
+def request_ingress_error_response(
+    request: Request,
+    *,
+    status_code: int,
+    code: str,
+    message: str,
+) -> JSONResponse:
+    path = get_route_path(request.scope)
+    uses_openai_errors = _uses_openai_ingress_errors(path)
+    uses_scim_errors = _uses_scim_ingress_errors(path)
+    response_code = "invalid_request_error" if uses_openai_errors and code == "invalid_request" else code
+    category = "dashboard_error_response"
+    if uses_openai_errors:
+        category = "openai_error_response"
+    elif uses_scim_errors:
+        category = "scim_error_response"
+    log_error_response(
+        logger,
+        request,
+        status_code,
+        response_code,
+        message,
+        category=category,
+    )
+    if uses_scim_errors:
+        return JSONResponse(
+            status_code=status_code,
+            content=scim_error(status_code, message),
+            media_type=SCIM_CONTENT_TYPE,
+        )
+    if uses_openai_errors:
+        return JSONResponse(
+            status_code=status_code,
+            content=openai_error(response_code, message, error_type="invalid_request_error"),
+        )
+    return JSONResponse(
+        status_code=status_code,
+        content=dashboard_error(response_code, message),
+    )
+
+
+def request_body_limit_was_exceeded(request: Request) -> bool:
+    return getattr(request.state, _REQUEST_BODY_TOO_LARGE_STATE, False) is True
+
+
+def _mark_request_body_limit_exceeded(scope: Scope) -> None:
+    state = scope.setdefault("state", {})
+    state[_REQUEST_BODY_TOO_LARGE_STATE] = True
+
+
+def _is_unencoded_multipart(headers: Headers) -> bool:
+    content_type = headers.get("content-type", "")
+    media_type = content_type.partition(";")[0].strip().lower()
+    return media_type == "multipart/form-data" and headers.get("content-encoding") is None
+
+
+def _declared_content_length(headers: Headers) -> int | None:
+    value = headers.get("content-length")
+    if value is None:
+        return None
+    try:
+        parsed = int(value)
+    except ValueError:
+        return None
+    return parsed if parsed >= 0 else None
+
+
+class RequestBodyLimitMiddleware:
+    def __init__(self, app: ASGIApp) -> None:
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        headers = Headers(scope=scope)
+        route_owned_multipart = is_route_owned_multipart_operation(scope) and _is_unencoded_multipart(headers)
+        if multipart_content_encoding_gate_was_applied(scope) or route_owned_multipart:
+            await self.app(scope, receive, send)
+            return
+
+        path = get_route_path(scope)
+        limit = request_body_limit_for_path(path)
+        declared_length = _declared_content_length(headers)
+        if declared_length is not None and declared_length > limit:
+            response = request_ingress_error_response(
+                Request(scope),
+                status_code=413,
+                code="payload_too_large",
+                message=REQUEST_BODY_TOO_LARGE_MESSAGE,
+            )
+            await response(scope, receive, send)
+            return
+
+        received = 0
+        response_started = False
+
+        async def limited_receive() -> Message:
+            nonlocal received
+            message = await receive()
+            if message["type"] == "http.request":
+                received += len(message.get("body", b""))
+                if received > limit:
+                    _mark_request_body_limit_exceeded(scope)
+                    raise _RequestBodyTooLarge
+            return message
+
+        async def tracked_send(message: Message) -> None:
+            nonlocal response_started
+            if message["type"] == "http.response.start":
+                response_started = True
+            await send(message)
+
+        try:
+            await self.app(scope, limited_receive, tracked_send)
+        except _RequestBodyTooLarge:
+            if response_started:
+                raise
+            response = request_ingress_error_response(
+                Request(scope),
+                status_code=413,
+                code="payload_too_large",
+                message=REQUEST_BODY_TOO_LARGE_MESSAGE,
+            )
+            await response(scope, receive, send)
+
+
+def add_request_body_limit_middleware(app: FastAPI) -> None:
+    app.add_middleware(RequestBodyLimitMiddleware)