codexLB-up 1.25.0b10__py3-none-any.whl

app/modules/proxy/ring_membership.py

@@ -0,0 +1,251 @@
+from __future__ import annotations
+
+import json
+import uuid
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from datetime import datetime
+from hashlib import sha256
+from typing import TYPE_CHECKING, Any, cast
+
+from sqlalchemy import CursorResult, delete, select, update
+from sqlalchemy.dialects.postgresql import insert as pg_insert
+from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.core.utils.time import utcnow
+from app.db.models import BridgeRingMember
+from app.db.session import close_session
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+RING_HEARTBEAT_INTERVAL_SECONDS = 10
+RING_STALE_THRESHOLD_SECONDS = 30
+RING_STALE_GRACE_SECONDS = RING_HEARTBEAT_INTERVAL_SECONDS + 5
+RING_MEMBER_RETENTION_SECONDS = 24 * 60 * 60
+
+
+class RingMembershipService:
+    """Manages pod registration in the bridge ring.
+
+    This service stores and retrieves active pod memberships from the DB,
+    ensuring all pods see the same ring view (solving the split-brain problem).
+    """
+
+    def __init__(self, session_factory: Callable[[], AsyncSession]) -> None:
+        self._session_factory = session_factory
+
+    async def register(self, instance_id: str, *, endpoint_base_url: str | None = None) -> None:
+        """Upsert pod into ring. Safe to call multiple times."""
+        async with self._session() as session:
+            # Dialect-specific upsert
+            dialect = session.get_bind().dialect.name
+            metadata_json = _bridge_ring_metadata_json(endpoint_base_url)
+            if dialect == "postgresql":
+                stmt = (
+                    pg_insert(BridgeRingMember)
+                    .values(
+                        id=str(uuid.uuid4()),
+                        instance_id=instance_id,
+                        registered_at=utcnow(),
+                        last_heartbeat_at=utcnow(),
+                        metadata_json=metadata_json,
+                    )
+                    .on_conflict_do_update(
+                        index_elements=["instance_id"],
+                        set_={
+                            "last_heartbeat_at": utcnow(),
+                            "registered_at": utcnow(),
+                            "metadata_json": metadata_json,
+                        },
+                    )
+                )
+            elif dialect == "sqlite":
+                stmt = (
+                    sqlite_insert(BridgeRingMember)
+                    .values(
+                        id=str(uuid.uuid4()),
+                        instance_id=instance_id,
+                        registered_at=utcnow(),
+                        last_heartbeat_at=utcnow(),
+                        metadata_json=metadata_json,
+                    )
+                    .on_conflict_do_update(
+                        index_elements=["instance_id"],
+                        set_={
+                            "last_heartbeat_at": utcnow(),
+                            "registered_at": utcnow(),
+                            "metadata_json": metadata_json,
+                        },
+                    )
+                )
+            else:
+                raise RuntimeError(f"RingMembershipService unsupported for dialect={dialect!r}")
+            await session.execute(stmt)
+            await session.commit()
+
+    async def heartbeat(self, instance_id: str, *, endpoint_base_url: str | None = None) -> None:
+        """Upsert heartbeat — recovers from mark_stale or unregister by sibling workers."""
+        async with self._session() as session:
+            dialect = session.get_bind().dialect.name
+            now = utcnow()
+            metadata_json = _bridge_ring_metadata_json(endpoint_base_url)
+            if dialect == "postgresql":
+                stmt = (
+                    pg_insert(BridgeRingMember)
+                    .values(
+                        id=str(uuid.uuid4()),
+                        instance_id=instance_id,
+                        registered_at=now,
+                        last_heartbeat_at=now,
+                        metadata_json=metadata_json,
+                    )
+                    .on_conflict_do_update(
+                        index_elements=["instance_id"],
+                        set_={"last_heartbeat_at": now, "metadata_json": metadata_json},
+                    )
+                )
+            elif dialect == "sqlite":
+                stmt = (
+                    sqlite_insert(BridgeRingMember)
+                    .values(
+                        id=str(uuid.uuid4()),
+                        instance_id=instance_id,
+                        registered_at=now,
+                        last_heartbeat_at=now,
+                        metadata_json=metadata_json,
+                    )
+                    .on_conflict_do_update(
+                        index_elements=["instance_id"],
+                        set_={"last_heartbeat_at": now, "metadata_json": metadata_json},
+                    )
+                )
+            else:
+                stmt = (
+                    update(BridgeRingMember)
+                    .where(BridgeRingMember.instance_id == instance_id)
+                    .values(last_heartbeat_at=now, metadata_json=metadata_json)
+                )
+            await session.execute(stmt)
+            await session.commit()
+
+    async def unregister(self, instance_id: str) -> None:
+        """Remove pod from ring."""
+        async with self._session() as session:
+            stmt = delete(BridgeRingMember).where(BridgeRingMember.instance_id == instance_id)
+            await session.execute(stmt)
+            await session.commit()
+
+    async def mark_stale(
+        self,
+        instance_id: str,
+        *,
+        stale_threshold_seconds: int = RING_STALE_THRESHOLD_SECONDS,
+        grace_seconds: int = RING_STALE_GRACE_SECONDS,
+    ) -> None:
+        """Age the heartbeat close to expiry without deleting the shared row.
+
+        A short grace window lets sibling workers refresh the shared row on
+        their next heartbeat, while a fully terminating pod still ages out far
+        faster than the normal stale threshold.
+        """
+        from datetime import timedelta
+
+        active_for_seconds = max(grace_seconds, 0)
+        age_seconds = max(stale_threshold_seconds - active_for_seconds, 0)
+        stale_time = utcnow() - timedelta(seconds=age_seconds)
+        async with self._session() as session:
+            stmt = (
+                update(BridgeRingMember)
+                .where(BridgeRingMember.instance_id == instance_id)
+                .values(last_heartbeat_at=stale_time)
+            )
+            await session.execute(stmt)
+            await session.commit()
+
+    async def purge_stale_before(self, cutoff: datetime) -> int:
+        """Delete ring rows whose heartbeat predates the cutoff (dead replicas)."""
+
+        async with self._session() as session:
+            result = cast(
+                CursorResult[Any],
+                await session.execute(delete(BridgeRingMember).where(BridgeRingMember.last_heartbeat_at < cutoff)),
+            )
+            await session.commit()
+        return int(result.rowcount or 0)
+
+    async def list_active(
+        self,
+        stale_threshold_seconds: int = RING_STALE_THRESHOLD_SECONDS,
+        *,
+        require_endpoint: bool = False,
+    ) -> list[str]:
+        """Return sorted list of pods whose heartbeat is within threshold."""
+        from datetime import timedelta
+
+        cutoff = utcnow() - timedelta(seconds=stale_threshold_seconds)
+        statement = select(BridgeRingMember.instance_id).where(BridgeRingMember.last_heartbeat_at >= cutoff)
+        if require_endpoint:
+            statement = statement.where(BridgeRingMember.metadata_json.is_not(None))
+        async with self._session() as session:
+            result = await session.execute(statement.order_by(BridgeRingMember.instance_id))
+            return list(result.scalars().all())
+
+    async def resolve_endpoint(
+        self,
+        instance_id: str,
+        *,
+        stale_threshold_seconds: int = RING_STALE_THRESHOLD_SECONDS,
+    ) -> str | None:
+        from datetime import timedelta
+
+        cutoff = utcnow() - timedelta(seconds=stale_threshold_seconds)
+        async with self._session() as session:
+            result = await session.execute(
+                select(BridgeRingMember.metadata_json)
+                .where(
+                    BridgeRingMember.instance_id == instance_id,
+                    BridgeRingMember.last_heartbeat_at >= cutoff,
+                )
+                .limit(1)
+            )
+            metadata_json = result.scalar_one_or_none()
+        return _bridge_ring_endpoint_from_metadata(metadata_json)
+
+    async def ring_fingerprint(self, stale_threshold_seconds: int = RING_STALE_THRESHOLD_SECONDS) -> str:
+        """sha256 of sorted active member list. Same for all pods with same membership."""
+        members = await self.list_active(stale_threshold_seconds)
+        data = ",".join(sorted(members))
+        return sha256(data.encode()).hexdigest()
+
+    @asynccontextmanager
+    async def _session(self) -> AsyncIterator[AsyncSession]:
+        session = self._session_factory()
+        try:
+            yield session
+        finally:
+            await close_session(session)
+
+
+def _bridge_ring_metadata_json(endpoint_base_url: str | None) -> str | None:
+    if endpoint_base_url is None:
+        return None
+    return json.dumps({"endpoint_base_url": endpoint_base_url}, ensure_ascii=True, separators=(",", ":"))
+
+
+def _bridge_ring_endpoint_from_metadata(metadata_json: str | None) -> str | None:
+    if metadata_json is None:
+        return None
+    try:
+        payload = json.loads(metadata_json)
+    except json.JSONDecodeError:
+        return None
+    if not isinstance(payload, dict):
+        return None
+    endpoint = payload.get("endpoint_base_url")
+    if not isinstance(endpoint, str):
+        return None
+    stripped = endpoint.strip().rstrip("/")
+    return stripped or None

app/modules/proxy/schemas.py

@@ -0,0 +1,345 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from app.core.clients.files import OPENAI_FILE_UPLOAD_LIMIT_BYTES, OPENAI_FILE_USE_CASE
+from app.core.types import JsonValue
+from app.modules.proxy.types import (
+    AdditionalRateLimitData,
+    CreditStatusDetailsData,
+    RateLimitResetCreditsData,
+    RateLimitStatusDetailsData,
+    RateLimitStatusPayloadData,
+    RateLimitWindowSnapshotData,
+)
+
+
+class FileCreateRequest(BaseModel):
+    """Inbound payload for ``POST /backend-api/files``.
+
+    Mirrors upstream ChatGPT ``/backend-api/files`` registration body. The
+    proxy enforces the upstream-side ``OPENAI_FILE_UPLOAD_LIMIT_BYTES``
+    (512 MiB) at the edge so a misbehaving client does not allocate an
+    oversized SAS upload URL on a shared account.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    file_name: str = Field(min_length=1)
+    file_size: int = Field(gt=0, le=OPENAI_FILE_UPLOAD_LIMIT_BYTES)
+    use_case: str = Field(default=OPENAI_FILE_USE_CASE, min_length=1)
+
+
+class RateLimitWindowSnapshot(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    used_percent: int
+    limit_window_seconds: int | None = None
+    reset_after_seconds: int | None = None
+    reset_at: int | None = None
+
+    @classmethod
+    def from_data(cls, data: RateLimitWindowSnapshotData) -> "RateLimitWindowSnapshot":
+        return cls(
+            used_percent=data.used_percent,
+            limit_window_seconds=data.limit_window_seconds,
+            reset_after_seconds=data.reset_after_seconds,
+            reset_at=data.reset_at,
+        )
+
+
+class RateLimitStatusDetails(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    allowed: bool
+    limit_reached: bool
+    primary_window: RateLimitWindowSnapshot | None = None
+    secondary_window: RateLimitWindowSnapshot | None = None
+    monthly_window: RateLimitWindowSnapshot | None = None
+
+    @classmethod
+    def from_data(cls, data: RateLimitStatusDetailsData) -> "RateLimitStatusDetails":
+        return cls(
+            allowed=data.allowed,
+            limit_reached=data.limit_reached,
+            primary_window=RateLimitWindowSnapshot.from_data(data.primary_window) if data.primary_window else None,
+            secondary_window=RateLimitWindowSnapshot.from_data(data.secondary_window)
+            if data.secondary_window
+            else None,
+            monthly_window=RateLimitWindowSnapshot.from_data(data.monthly_window) if data.monthly_window else None,
+        )
+
+
+class CreditStatusDetails(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    has_credits: bool
+    unlimited: bool
+    balance: str | None = None
+    approx_local_messages: list[JsonValue] | None = None
+    approx_cloud_messages: list[JsonValue] | None = None
+
+    @classmethod
+    def from_data(cls, data: CreditStatusDetailsData) -> "CreditStatusDetails":
+        return cls(
+            has_credits=data.has_credits,
+            unlimited=data.unlimited,
+            balance=data.balance,
+            approx_local_messages=data.approx_local_messages,
+            approx_cloud_messages=data.approx_cloud_messages,
+        )
+
+
+class AdditionalRateLimitStatus(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    quota_key: str | None = None
+    limit_name: str
+    display_label: str | None = None
+    metered_feature: str
+    rate_limit: RateLimitStatusDetails | None = None
+
+    @classmethod
+    def from_data(cls, data: AdditionalRateLimitData) -> "AdditionalRateLimitStatus":
+        return cls(
+            quota_key=data.quota_key,
+            limit_name=data.limit_name,
+            display_label=data.display_label,
+            metered_feature=data.metered_feature,
+            rate_limit=RateLimitStatusDetails.from_data(data.rate_limit) if data.rate_limit else None,
+        )
+
+
+class RateLimitResetCredits(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    available_count: int
+
+    @classmethod
+    def from_data(cls, data: RateLimitResetCreditsData) -> "RateLimitResetCredits":
+        return cls(available_count=data.available_count)
+
+
+class RateLimitStatusPayload(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    plan_type: str
+    rate_limit: RateLimitStatusDetails | None = None
+    credits: CreditStatusDetails | None = None
+    rate_limit_reset_credits: RateLimitResetCredits | None = None
+    additional_rate_limits: list[AdditionalRateLimitStatus] = []
+
+    @classmethod
+    def from_data(cls, data: RateLimitStatusPayloadData) -> "RateLimitStatusPayload":
+        return cls(
+            plan_type=data.plan_type,
+            rate_limit=RateLimitStatusDetails.from_data(data.rate_limit) if data.rate_limit else None,
+            credits=CreditStatusDetails.from_data(data.credits) if data.credits else None,
+            rate_limit_reset_credits=(
+                RateLimitResetCredits.from_data(data.rate_limit_reset_credits)
+                if data.rate_limit_reset_credits
+                else None
+            ),
+            additional_rate_limits=[AdditionalRateLimitStatus.from_data(arl) for arl in data.additional_rate_limits],
+        )
+
+
+class ConsumeRateLimitResetCreditRequest(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    redeem_request_id: str
+
+
+class ConsumeRateLimitResetCreditResponse(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    code: str
+    windows_reset: int = 0
+
+
+class ReasoningLevelSchema(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    effort: str
+    description: str
+
+
+class CodexTruncationPolicy(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    mode: Literal["bytes", "tokens"]
+    limit: int = Field(strict=True, ge=-(2**63), le=2**63 - 1)
+
+
+class CodexModelEntry(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    slug: str
+    display_name: str
+    description: str
+    base_instructions: str = ""
+    default_reasoning_level: str | None = None
+    supported_reasoning_levels: list[ReasoningLevelSchema] = []
+    supported_in_api: bool = True
+    priority: int = 0
+    minimal_client_version: str | None = None
+    supports_reasoning_summaries: bool = False
+    support_verbosity: bool = False
+    default_verbosity: str | None = None
+    supports_parallel_tool_calls: bool = False
+    context_window: int = 0
+    input_modalities: list[str] = []
+    available_in_plans: list[str] = []
+    prefer_websockets: bool = False
+    visibility: str = "list"
+    truncation_policy: CodexTruncationPolicy
+    experimental_supported_tools: list[str]
+
+
+class ModelMetadata(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    display_name: str
+    description: str
+    context_window: int
+    input_context_window: int | None = None
+    max_output_tokens: int | None = None
+    input_modalities: list[str]
+    supported_reasoning_levels: list[ReasoningLevelSchema]
+    default_reasoning_level: str | None = None
+    supports_reasoning_summaries: bool = False
+    support_verbosity: bool = False
+    default_verbosity: str | None = None
+    prefer_websockets: bool = False
+    supports_parallel_tool_calls: bool = False
+    supported_in_api: bool = True
+    minimal_client_version: str | None = None
+    priority: int = 0
+    additional_speed_tiers: list[str] | None = None
+    service_tiers: list[dict[str, JsonValue]] | None = None
+    default_service_tier: str | None = None
+
+
+class ModelListItem(BaseModel):
+    # Cursor's local-provider discovery reads OpenAI-compatible /v1/models
+    # entries and preserves provider-specific model capability fields. Keep
+    # allowing those extras so clients can learn the model context window and
+    # trigger their own compaction instead of relying on provider-side failures.
+    model_config = ConfigDict(extra="allow")
+
+    id: str
+    object: str = "model"
+    created: int
+    owned_by: str
+    metadata: ModelMetadata | None = None
+
+
+class ModelListResponse(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    object: str = "list"
+    data: list[ModelListItem]
+
+
+class CodexModelsResponse(BaseModel):
+    models: list[CodexModelEntry]
+    object: str = "list"
+    data: list[ModelListItem] = []
+
+
+class V1UsageLimitResponse(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    limit_type: str
+    limit_window: str
+    max_value: int
+    current_value: int
+    remaining_value: int
+    model_filter: str | None = None
+    reset_at: str
+    source: str = "api_key_limit"
+
+
+class AccountPoolUsageResponse(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    primary: float | None = None
+    secondary: float | None = None
+
+
+class V1UsageResponse(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    request_count: int
+    total_tokens: int
+    cached_input_tokens: int
+    total_cost_usd: float
+    limits: list[V1UsageLimitResponse]
+    upstream_limits: list[V1UsageLimitResponse] = []
+    account_pool_usage: AccountPoolUsageResponse | None = None
+
+
+class V1ResetCreditEntry(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    account_id: str
+    email: str
+    redeem_id: str
+    expired_at: datetime | None = Field(serialization_alias="expiredAt")
+
+
+class V1ResetCreditRedeemRequest(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    account_id: str
+    redeem_id: str
+
+
+class V1ResetCreditRedeemResponse(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    code: str
+    windows_reset: int
+    redeemed_at: datetime | None = None
+
+
+class WarmupRequest(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    mode: str = "normal"
+
+
+class WarmupSubmittedAccount(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    account_id: str
+    request_id: str
+    model: str
+
+
+class WarmupSkippedAccount(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    account_id: str
+    reason: str
+
+
+class WarmupFailedAccount(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    account_id: str
+    error_code: str
+    error_message: str
+
+
+class WarmupResponse(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    mode: str
+    total_accounts: int
+    submitted: list[WarmupSubmittedAccount]
+    skipped: list[WarmupSkippedAccount]
+    failed: list[WarmupFailedAccount]

app/modules/proxy/selection_errors.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from typing import Protocol
+
+from app.core.errors import OpenAIErrorEnvelope, openai_error
+from app.core.resilience.overload import is_local_overload_error_code
+
+USAGE_LIMIT_REACHED = "usage_limit_reached"
+
+
+class SelectionFailure(Protocol):
+    error_message: str | None
+    error_code: str | None
+    resets_at: int | None
+
+
+def selection_failure_response(selection: SelectionFailure) -> tuple[int, OpenAIErrorEnvelope]:
+    """Map an account-selection failure to its externally visible HTTP response.
+
+    The ``usage_limit_reached`` mapping is strictly for upstream usage/quota
+    exhaustion of the whole eligible pool. Local capacity codes (account caps,
+    admission gates, fair-share throttles) resolve against the canonical
+    ``LOCAL_OVERLOAD_CODES`` registry so they keep their stable 429
+    ``rate_limit_error`` contract and are never reclassified as upstream usage
+    exhaustion or collapsed into a generic 503.
+    """
+    code = selection.error_code or "no_accounts"
+    message = selection.error_message or "No active accounts available"
+    if code == USAGE_LIMIT_REACHED:
+        return (
+            429,
+            openai_error(
+                code,
+                message,
+                error_type=USAGE_LIMIT_REACHED,
+                resets_at=selection.resets_at,
+            ),
+        )
+    if is_local_overload_error_code(code):
+        return 429, openai_error(code, message, error_type="rate_limit_error")
+    return 503, openai_error(code, message)