codexLB-up 1.25.0b10__py3-none-any.whl

app/modules/proxy/durable_bridge_runtime.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from uuid import uuid4
+
+_HTTP_BRIDGE_OWNER_PROCESS_EPOCH = uuid4().hex
+
+
+def http_bridge_owner_process_epoch() -> str:
+    return _HTTP_BRIDGE_OWNER_PROCESS_EPOCH

app/modules/proxy/durable_bridge_transcript_codec.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import hmac
+import struct
+import zlib
+from collections.abc import Sequence
+from dataclasses import dataclass
+from hashlib import sha256
+
+DURABLE_BRIDGE_CHUNK_CODEC = "zlib-length-prefixed-utf8-v1"
+DURABLE_BRIDGE_CHUNK_MAX_EVENTS = 256
+DURABLE_BRIDGE_TRANSCRIPT_MAX_EVENTS = 65_536
+_FRAME_LENGTH = struct.Struct(">I")
+
+
+class DurableBridgeTranscriptDecodeError(ValueError):
+    """Raised when a durable chunk cannot be decoded without ambiguity."""
+
+
+@dataclass(frozen=True, slots=True)
+class EncodedDurableBridgeTranscriptChunk:
+    codec: str
+    payload: bytes
+    event_count: int
+    uncompressed_bytes: int
+    payload_sha256: str
+
+
+def encode_durable_bridge_transcript_chunk(
+    events: Sequence[str],
+) -> EncodedDurableBridgeTranscriptChunk:
+    if not events:
+        raise ValueError("durable transcript chunks require at least one event")
+    if len(events) > DURABLE_BRIDGE_CHUNK_MAX_EVENTS:
+        raise ValueError("durable transcript chunk exceeds event-count limit")
+    framed = bytearray()
+    for event in events:
+        encoded = event.encode("utf-8")
+        if len(encoded) > 0xFFFFFFFF:
+            raise ValueError("durable transcript event exceeds framing limit")
+        framed.extend(_FRAME_LENGTH.pack(len(encoded)))
+        framed.extend(encoded)
+    canonical = bytes(framed)
+    return EncodedDurableBridgeTranscriptChunk(
+        codec=DURABLE_BRIDGE_CHUNK_CODEC,
+        payload=zlib.compress(canonical),
+        event_count=len(events),
+        uncompressed_bytes=len(canonical),
+        payload_sha256=sha256(canonical).hexdigest(),
+    )
+
+
+def decode_durable_bridge_transcript_chunk(
+    *,
+    codec: str,
+    payload: bytes,
+    event_count: int,
+    uncompressed_bytes: int,
+    payload_sha256: str,
+    max_uncompressed_bytes: int,
+) -> tuple[str, ...]:
+    if codec != DURABLE_BRIDGE_CHUNK_CODEC:
+        raise DurableBridgeTranscriptDecodeError("unknown durable transcript chunk codec")
+    if type(event_count) is not int or type(uncompressed_bytes) is not int:
+        raise DurableBridgeTranscriptDecodeError("durable transcript chunk numeric metadata must be integral")
+    if event_count <= 0:
+        raise DurableBridgeTranscriptDecodeError("durable transcript chunk event count must be positive")
+    if event_count > DURABLE_BRIDGE_CHUNK_MAX_EVENTS:
+        raise DurableBridgeTranscriptDecodeError("durable transcript chunk exceeds event-count limit")
+    if uncompressed_bytes < 0 or uncompressed_bytes > max_uncompressed_bytes:
+        raise DurableBridgeTranscriptDecodeError("durable transcript chunk exceeds decode byte limit")
+
+    decompressor = zlib.decompressobj()
+    try:
+        canonical = decompressor.decompress(payload, uncompressed_bytes + 1)
+    except zlib.error as exc:
+        raise DurableBridgeTranscriptDecodeError("invalid durable transcript compressed payload") from exc
+    if (
+        len(canonical) != uncompressed_bytes
+        or decompressor.unconsumed_tail
+        or decompressor.unused_data
+        or not decompressor.eof
+    ):
+        raise DurableBridgeTranscriptDecodeError("durable transcript chunk byte count mismatch")
+    try:
+        expected_digest = bytes.fromhex(payload_sha256)
+    except (TypeError, ValueError) as exc:
+        raise DurableBridgeTranscriptDecodeError("durable transcript chunk hash is not hexadecimal") from exc
+    if len(expected_digest) != sha256().digest_size:
+        raise DurableBridgeTranscriptDecodeError("durable transcript chunk hash has invalid length")
+    if not hmac.compare_digest(sha256(canonical).digest(), expected_digest):
+        raise DurableBridgeTranscriptDecodeError("durable transcript chunk hash mismatch")
+
+    events: list[str] = []
+    offset = 0
+    for _ in range(event_count):
+        header_end = offset + _FRAME_LENGTH.size
+        if header_end > len(canonical):
+            raise DurableBridgeTranscriptDecodeError("durable transcript chunk has truncated frame header")
+        (event_size,) = _FRAME_LENGTH.unpack(canonical[offset:header_end])
+        event_end = header_end + event_size
+        if event_end > len(canonical):
+            raise DurableBridgeTranscriptDecodeError("durable transcript chunk has truncated event")
+        try:
+            events.append(canonical[header_end:event_end].decode("utf-8"))
+        except UnicodeDecodeError as exc:
+            raise DurableBridgeTranscriptDecodeError("durable transcript chunk contains invalid UTF-8") from exc
+        offset = event_end
+    if offset != len(canonical):
+        raise DurableBridgeTranscriptDecodeError("durable transcript chunk contains trailing bytes")
+    return tuple(events)
+
+
+__all__ = [
+    "DURABLE_BRIDGE_CHUNK_CODEC",
+    "DURABLE_BRIDGE_CHUNK_MAX_EVENTS",
+    "DURABLE_BRIDGE_TRANSCRIPT_MAX_EVENTS",
+    "DurableBridgeTranscriptDecodeError",
+    "EncodedDurableBridgeTranscriptChunk",
+    "decode_durable_bridge_transcript_chunk",
+    "encode_durable_bridge_transcript_chunk",
+]

app/modules/proxy/fair_share.py

@@ -0,0 +1,131 @@
+"""Congestion-aware per-API-key fair-share stream admission.
+
+Stream slots are a pool-wide scarce resource: the per-account stream cap
+protects individual accounts from upstream rate limits, but slot allocation
+across API keys is otherwise first-come-first-served, so one bursty key can
+hold every slot and starve every other key. This module provides the pure
+max-min fair-share decision used by the load balancer's stream admission:
+
+- Work-conserving: while pool utilization is below the configured congestion
+  threshold, every key admits unconditionally — idle capacity is never wasted
+  on a static per-key limit.
+- Under congestion, a key may take a new slot only while it holds fewer than
+  ``max(MIN_GUARANTEE_STREAMS, pool_capacity // active_keys)`` streams. Keys
+  under their fair share always admit (a key holding fewer than the minimum
+  guarantee can never be denied, making light interactive keys
+  starvation-proof), while over-share keys cannot reacquire freed slots until
+  they drop back under — freed capacity flows to the keys the congestion was
+  starving.
+
+All quantities are computed over the selection's candidate account set, so
+account-scoped keys are measured against the pool they can actually use. All
+arithmetic is integer-only so decisions are deterministic and exactly
+testable. The module performs no I/O and holds no state; the impure edge
+(runtime counters, locking, metrics) lives in the load balancer.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+API_KEY_STREAM_FAIR_SHARE_ERROR_CODE = "api_key_stream_fair_share"
+
+MIN_GUARANTEE_STREAMS = 2
+
+
+@dataclass(frozen=True, slots=True)
+class FairShareDecision:
+    admitted: bool
+    congested: bool
+    fair_share: int
+    pool_capacity: int
+    pool_inflight: int
+    active_key_count: int
+    requester_inflight: int
+
+
+def effective_stream_pool_capacity(
+    *,
+    candidate_account_count: int,
+    stream_limit: int,
+    stream_reserve_slots: int,
+) -> int:
+    """Return the candidate pool's stream capacity, or 0 when ungateable.
+
+    Mirrors the per-account admission arithmetic: each account contributes
+    ``max(1, stream_limit - stream_reserve_slots)`` slots. A nonpositive
+    stream limit means unlimited caps and a zero return disables the gate.
+    """
+    if stream_limit <= 0 or candidate_account_count <= 0:
+        return 0
+    per_account = max(1, stream_limit - max(0, stream_reserve_slots))
+    return candidate_account_count * per_account
+
+
+def evaluate_stream_fair_share(
+    *,
+    pool_capacity: int,
+    pool_inflight: int,
+    requester_inflight: int,
+    other_active_key_count: int,
+    threshold_pct: int,
+    min_guarantee: int = MIN_GUARANTEE_STREAMS,
+) -> FairShareDecision:
+    """Return the fair-share admission decision for one stream request.
+
+    ``other_active_key_count`` counts keys other than the requester holding at
+    least one in-flight stream on the candidate accounts; the requester is
+    always counted exactly once on top of it, whether or not it is already
+    active. Congestion is ``pool_inflight * 100 >= pool_capacity *
+    threshold_pct`` (integer math); a ``threshold_pct`` or ``pool_capacity``
+    of zero disables the gate and admits.
+    """
+    active_key_count = other_active_key_count + 1
+    if threshold_pct <= 0 or pool_capacity <= 0:
+        return FairShareDecision(
+            admitted=True,
+            congested=False,
+            fair_share=0,
+            pool_capacity=pool_capacity,
+            pool_inflight=pool_inflight,
+            active_key_count=active_key_count,
+            requester_inflight=requester_inflight,
+        )
+    congested = pool_inflight * 100 >= pool_capacity * threshold_pct
+    fair_share = max(min_guarantee, pool_capacity // active_key_count)
+    admitted = not congested or requester_inflight + 1 <= fair_share
+    return FairShareDecision(
+        admitted=admitted,
+        congested=congested,
+        fair_share=fair_share,
+        pool_capacity=pool_capacity,
+        pool_inflight=pool_inflight,
+        active_key_count=active_key_count,
+        requester_inflight=requester_inflight,
+    )
+
+
+class ApiKeyFairShareDenialError(Exception):
+    """Fair-share gate denial from a direct (selection-less) lease acquire.
+
+    ``LoadBalancer.select_account`` reports denials through the selection
+    error surface; ``LoadBalancer.acquire_account_lease`` returns ``None``
+    only for plain per-account cap denials, so the fair-share denial raises
+    this instead of overloading that return value. Carries the full
+    ``FairShareDecision`` for callers that surface the denial numbers.
+    """
+
+    def __init__(self, decision: FairShareDecision) -> None:
+        super().__init__(fair_share_denial_message(decision))
+        self.decision = decision
+
+
+def fair_share_denial_message(decision: FairShareDecision) -> str:
+    """Human-readable denial message carrying the decision's numbers."""
+    return (
+        "API key stream fair share exceeded under pool congestion; "
+        f"key holds {decision.requester_inflight} of a fair share of {decision.fair_share} streams "
+        f"({decision.pool_inflight}/{decision.pool_capacity} pool streams in flight across "
+        f"{decision.active_key_count} active keys). "
+        "Waiting for the key's streams to finish or the pool to decongest."
+    )

app/modules/proxy/file_pin_repository.py

@@ -0,0 +1,210 @@
+from __future__ import annotations
+
+from collections.abc import Collection
+
+from sqlalchemy import Integer, bindparam, text
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.sql.elements import TextClause
+
+from app.db.session import sqlite_writer_section
+
+_TABLE = "file_account_pins"
+
+# Ownership TTLs stay entirely in the database clock domain. PostgreSQL's
+# clock_timestamp() is evaluated when each clause executes. A successful claim
+# is followed by a guarded refresh in the same transaction so even an INSERT
+# that waited behind an ultimately rolled-back unique contender receives its
+# full TTL after the wait. SQLite's padded strftime form matches SQLAlchemy
+# DateTime's six-digit fractional width, preserving exact lexicographic expiry
+# checks.
+_POSTGRES_NOW = "clock_timestamp()"
+_POSTGRES_NOW_PLUS_TTL = "clock_timestamp() + make_interval(secs => :ttl)"
+_POSTGRES_STATEMENT_NOW = "statement_timestamp()"
+_SQLITE_NOW = "(strftime('%Y-%m-%d %H:%M:%f', 'now') || '000')"
+_SQLITE_NOW_PLUS_TTL = "(strftime('%Y-%m-%d %H:%M:%f', 'now', '+' || :ttl || ' seconds') || '000')"
+
+_POSTGRES_CLAIM = text(
+    f"""
+    INSERT INTO {_TABLE} (file_id, account_id, expires_at)
+    VALUES (:file_id, :account_id, {_POSTGRES_NOW_PLUS_TTL})
+    ON CONFLICT (file_id) DO UPDATE SET
+        account_id = excluded.account_id,
+        expires_at = {_POSTGRES_NOW_PLUS_TTL}
+    WHERE {_TABLE}.account_id = :account_id
+       OR {_TABLE}.expires_at <= {_POSTGRES_NOW}
+    RETURNING account_id
+    """
+).bindparams(bindparam("ttl", type_=Integer))
+
+_SQLITE_CLAIM = text(
+    f"""
+    INSERT INTO {_TABLE} (file_id, account_id, expires_at)
+    VALUES (:file_id, :account_id, {_SQLITE_NOW_PLUS_TTL})
+    ON CONFLICT (file_id) DO UPDATE SET
+        account_id = excluded.account_id,
+        expires_at = {_SQLITE_NOW_PLUS_TTL}
+    WHERE {_TABLE}.account_id = :account_id
+       OR {_TABLE}.expires_at <= {_SQLITE_NOW}
+    RETURNING account_id
+    """
+).bindparams(bindparam("ttl", type_=Integer))
+
+_POSTGRES_CLEANUP = text(f"DELETE FROM {_TABLE} WHERE expires_at <= {_POSTGRES_STATEMENT_NOW}")
+_SQLITE_CLEANUP = text(f"DELETE FROM {_TABLE} WHERE expires_at <= {_SQLITE_NOW}")
+
+_POSTGRES_REFRESH = text(
+    f"""
+    UPDATE {_TABLE}
+    SET expires_at = {_POSTGRES_NOW_PLUS_TTL}
+    WHERE file_id = :file_id
+      AND account_id = :account_id
+    RETURNING account_id
+    """
+).bindparams(bindparam("ttl", type_=Integer))
+_SQLITE_REFRESH = text(
+    f"""
+    UPDATE {_TABLE}
+    SET expires_at = {_SQLITE_NOW_PLUS_TTL}
+    WHERE file_id = :file_id
+      AND account_id = :account_id
+    RETURNING account_id
+    """
+).bindparams(bindparam("ttl", type_=Integer))
+
+_POSTGRES_GET_LIVE = text(f"SELECT account_id FROM {_TABLE} WHERE file_id = :file_id AND expires_at > {_POSTGRES_NOW}")
+_SQLITE_GET_LIVE = text(f"SELECT account_id FROM {_TABLE} WHERE file_id = :file_id AND expires_at > {_SQLITE_NOW}")
+
+_POSTGRES_GET_LIVE_MANY = text(
+    f"""
+    SELECT file_id, account_id
+    FROM {_TABLE}
+    WHERE file_id IN :file_ids
+      AND expires_at > {_POSTGRES_NOW}
+    """
+).bindparams(bindparam("file_ids", expanding=True))
+_SQLITE_GET_LIVE_MANY = text(
+    f"""
+    SELECT file_id, account_id
+    FROM {_TABLE}
+    WHERE file_id IN :file_ids
+      AND expires_at > {_SQLITE_NOW}
+    """
+).bindparams(bindparam("file_ids", expanding=True))
+
+_GET_ACCOUNT = text(f"SELECT account_id FROM {_TABLE} WHERE file_id = :file_id")
+
+
+class FileAccountPinOwnershipConflict(RuntimeError):
+    def __init__(self, file_id: str, persisted_account_id: str, requested_account_id: str) -> None:
+        super().__init__(
+            f"Live file ownership conflict for {file_id!r}: "
+            f"persisted={persisted_account_id!r} requested={requested_account_id!r}"
+        )
+        self.file_id = file_id
+        self.persisted_account_id = persisted_account_id
+        self.requested_account_id = requested_account_id
+
+
+def build_file_account_pin_claim(*, dialect_name: str) -> TextClause:
+    if dialect_name == "postgresql":
+        return _POSTGRES_CLAIM
+    if dialect_name == "sqlite":
+        return _SQLITE_CLAIM
+    raise RuntimeError(f"Unsupported database dialect for file account pins: {dialect_name}")
+
+
+def build_file_account_pin_cleanup(*, dialect_name: str) -> TextClause:
+    if dialect_name == "postgresql":
+        return _POSTGRES_CLEANUP
+    if dialect_name == "sqlite":
+        return _SQLITE_CLEANUP
+    raise RuntimeError(f"Unsupported database dialect for file account pins: {dialect_name}")
+
+
+def build_file_account_pin_refresh(*, dialect_name: str) -> TextClause:
+    if dialect_name == "postgresql":
+        return _POSTGRES_REFRESH
+    if dialect_name == "sqlite":
+        return _SQLITE_REFRESH
+    raise RuntimeError(f"Unsupported database dialect for file account pins: {dialect_name}")
+
+
+def build_file_account_pin_live_lookup(*, dialect_name: str, many: bool = False) -> TextClause:
+    if dialect_name == "postgresql":
+        return _POSTGRES_GET_LIVE_MANY if many else _POSTGRES_GET_LIVE
+    if dialect_name == "sqlite":
+        return _SQLITE_GET_LIVE_MANY if many else _SQLITE_GET_LIVE
+    raise RuntimeError(f"Unsupported database dialect for file account pins: {dialect_name}")
+
+
+class FileAccountPinRepository:
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+
+    async def claim(self, file_id: str, account_id: str, *, ttl_seconds: int) -> None:
+        if ttl_seconds <= 0:
+            raise ValueError("File account pin TTL must be positive")
+        dialect_name = self._dialect_name()
+        params = {
+            "file_id": file_id,
+            "account_id": account_id,
+            "ttl": ttl_seconds,
+        }
+        async with sqlite_writer_section():
+            await self._session.execute(build_file_account_pin_cleanup(dialect_name=dialect_name))
+            persisted_account_id = (
+                await self._session.execute(
+                    build_file_account_pin_claim(dialect_name=dialect_name),
+                    params,
+                )
+            ).scalar_one_or_none()
+            if persisted_account_id is None:
+                persisted_account_id = await self._session.scalar(
+                    _GET_ACCOUNT,
+                    {"file_id": file_id},
+                )
+            if persisted_account_id != account_id:
+                await self._session.rollback()
+                raise FileAccountPinOwnershipConflict(
+                    file_id,
+                    persisted_account_id or "<missing>",
+                    account_id,
+                )
+            refreshed_account_id = (
+                await self._session.execute(
+                    build_file_account_pin_refresh(dialect_name=dialect_name),
+                    params,
+                )
+            ).scalar_one_or_none()
+            if refreshed_account_id != account_id:
+                await self._session.rollback()
+                raise RuntimeError(f"Failed to refresh file account pin after claim: {file_id!r}")
+            await self._session.commit()
+
+    async def get_live_account_id(self, file_id: str) -> str | None:
+        return await self._session.scalar(
+            build_file_account_pin_live_lookup(dialect_name=self._dialect_name()),
+            {"file_id": file_id},
+        )
+
+    async def get_live_account_ids(self, file_ids: Collection[str]) -> dict[str, str]:
+        unique_file_ids = tuple(dict.fromkeys(file_ids))
+        if not unique_file_ids:
+            return {}
+        rows = (
+            (
+                await self._session.execute(
+                    build_file_account_pin_live_lookup(
+                        dialect_name=self._dialect_name(),
+                        many=True,
+                    ),
+                    {"file_ids": unique_file_ids},
+                )
+            )
+            .tuples()
+            .all()
+        )
+        return dict(rows)
+
+    def _dialect_name(self) -> str:
+        return self._session.get_bind().dialect.name