codexLB-up 1.25.0b10__py3-none-any.whl

app/modules/proxy/_load_balancer/throughput_cohort.py

@@ -0,0 +1,233 @@
+"""Replica-local, fleet-relative output-throughput weight per (account, model).
+
+First-token latency (``ttft_cohort.py``) misses the dominant per-account speed
+difference seen in production: on ``gpt-5.6-sol`` the fleet splits into two
+output-throughput cohorts (nine accounts near 40 tok/s, ten near 66-75 tok/s)
+whose first-token latency is the same, while on ``gpt-6-astra`` the same
+accounts are uniform. Throughput is therefore a property of the (account,
+model) pair, not of the account, and a slow-on-one-model account must not be
+discounted on a model where it is not slow.
+
+This module keeps, per account and per model, a bounded list of recent output
+throughputs -- total ``output_tokens`` over the generation span
+``latency_ms - latency_first_token_ms`` -- for successful ``normal`` turns
+with at least 200 output tokens, a single upstream attempt and no local
+queueing (the same funnel filters as the TTFT sampler), on any input size and
+reasoning effort. ``output_tokens`` is the total the upstream reports,
+reasoning tokens included. Where the reasoning lands depends on whether the
+client asked for reasoning summaries: with summaries the first summary delta
+is the first token, the reasoning runs inside the span and the total is what
+the span produced; without summaries the whole reasoning prefix sits inside
+the TTFT and the sample overstates the span's throughput by ``total /
+visible`` -- several-fold on a turn that reasons for longer than it answers,
+not a small bias. No numerator is exact for both kinds of turn (visible tokens
+alone would under-read every summary-streaming turn by the same factor). The
+total is kept because the comparison is fleet-relative: fresh draws spread
+clients across accounts, so the summary/no-summary mix is shared and the
+per-account medians move together; a per-account median only shifts when more
+than half of that account's in-window samples are inflated, and the fleet
+reference only when more than half of the accounts' medians are, so a spurious
+discount needs an uneven mix that persists over most of an hour on most of the
+fleet. Inflation reads *fast*, so it never discounts the account it lands on.
+
+Estimator: the per-(account, model) estimate is the median of the sample
+throughputs. Unlike first-token latency, throughput is not bimodal -- it is
+the serving path's speed, largely independent of the prompt -- so a median
+neither hides a mode share nor follows a single stalled stream (a network
+pause or a slow client) or a short burst; the 200-token floor keeps the
+quantization of ``tokens / duration`` small. The fleet reference is the median
+of the per-account estimates for the requested model; an account below the
+reference by more than the deadband is discounted to ``max(floor, account /
+fleet)``. ``latency_cohort.apply_latency_cohort_weights`` combines this with the
+TTFT multiplier as the minimum of the two (never their product) and only for
+fresh weighted draws.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from statistics import median
+from typing import Any
+
+from app.modules.proxy._load_balancer.ttft_cohort import (
+    TTFT_MAX_SAMPLES,
+    TTFT_MIN_ACCOUNTS,
+    TTFT_SAMPLE_REQUEST_KIND,
+    TTFT_SAMPLE_WINDOW_SECONDS,
+)
+from app.modules.proxy._load_balancer.types import RuntimeState
+
+# Same window/cap pattern as the TTFT sampler, per (account, model).
+TPS_SAMPLE_WINDOW_SECONDS = TTFT_SAMPLE_WINDOW_SECONDS
+TPS_MAX_SAMPLES = TTFT_MAX_SAMPLES
+# Below this many samples for a model an account's estimate is noise and stays neutral.
+TPS_MIN_SAMPLES = 8
+# Short answers quantize ``tokens / duration`` too coarsely to compare.
+TPS_SAMPLE_MIN_OUTPUT_TOKENS = 200
+# Rows written without a model (direct WebSocket rows carry ``""``) or with the
+# placeholder cannot be attributed to a per-model cohort.
+_UNKNOWN_MODEL = "unknown"
+# An account within 15% below the fleet reference is not discounted at all.
+TPS_WEIGHT_DEADBAND = 0.15
+# A 10x slower account still keeps half of its weight: the window keeps
+# sampling it and the discount lifts as soon as its throughput recovers.
+TPS_WEIGHT_FLOOR = 0.5
+
+
+def _sampleable_model(model: str | None) -> str | None:
+    if not model or model == _UNKNOWN_MODEL:
+        return None
+    return model
+
+
+def _prune(samples_by_model: dict[str, list[tuple[float, float]]], now: float) -> None:
+    oldest_kept = now - TPS_SAMPLE_WINDOW_SECONDS
+    for model in list(samples_by_model):
+        samples = samples_by_model[model]
+        stale = 0
+        for recorded_at, _ in samples:
+            if recorded_at >= oldest_kept:
+                break
+            stale += 1
+        if stale:
+            del samples[:stale]
+        excess = len(samples) - TPS_MAX_SAMPLES
+        if excess > 0:
+            del samples[:excess]
+        if not samples:
+            del samples_by_model[model]
+
+
+def _throughput_sample(
+    *,
+    account_id: str | None,
+    status: str,
+    request_kind: str,
+    model: str | None,
+    latency_ms: int | None,
+    latency_first_token_ms: int | None,
+    output_tokens: int | None,
+    queued_wait_ms: int,
+    retried: bool,
+) -> tuple[str, float] | None:
+    """``(model, tokens_per_second)`` for an eligible row, else ``None``."""
+    sample_model = _sampleable_model(model)
+    if (
+        not account_id
+        or status != "success"
+        or request_kind != TTFT_SAMPLE_REQUEST_KIND
+        or retried
+        or sample_model is None
+        or latency_first_token_ms is None
+        or latency_first_token_ms < 0
+        or latency_ms is None
+        or latency_ms <= latency_first_token_ms
+        or output_tokens is None
+        or output_tokens < TPS_SAMPLE_MIN_OUTPUT_TOKENS
+        or queued_wait_ms > 0
+    ):
+        return None
+    return sample_model, output_tokens / ((latency_ms - latency_first_token_ms) / 1000.0)
+
+
+def record_tps_sample(
+    balancer: Any,
+    *,
+    account_id: str | None,
+    status: str,
+    request_kind: str,
+    model: str | None,
+    latency_ms: int | None,
+    latency_first_token_ms: int | None,
+    output_tokens: int | None,
+    queued_wait_ms: int,
+    retried: bool = False,
+) -> None:
+    """Record one eligible output throughput for ``(account_id, model)`` at the balancer clock.
+
+    ``retried`` and ``queued_wait_ms`` have the meaning documented on
+    ``ttft_cohort.record_ttft_sample``: a retried send or a capacity wait
+    leaves a failed attempt inside ``latency_ms``, and a queued row's span is
+    not the upstream's. ``latency_ms`` MUST end at the upstream terminal event,
+    not at the end of downstream delivery or local settlement/cleanup: the
+    request-log funnel passes the ``latency_upstream_terminal_ms`` each path
+    stamps when it parses the terminal frame (HTTP stream settlement, bridge
+    reader, WebSocket reader) in place of the row's latency. The throughput is
+    ``output_tokens`` over the generation span
+    ``latency_ms - latency_first_token_ms``; rows without a
+    model, with the ``unknown`` placeholder, without a first token, with a
+    non-positive span or with fewer than ``TPS_SAMPLE_MIN_OUTPUT_TOKENS``
+    output tokens add nothing.
+
+    Same concurrency contract as the TTFT sampler: synchronous, lock-free, no
+    await between reading and writing the per-account map; no-op when
+    ``balancer`` does not expose the runtime map or clock (test doubles).
+    """
+    runtime_map = getattr(balancer, "_runtime", None)
+    clock = getattr(balancer, "_clock", None)
+    if not isinstance(runtime_map, dict) or clock is None or not account_id:
+        return
+    sample = _throughput_sample(
+        account_id=account_id,
+        status=status,
+        request_kind=request_kind,
+        model=model,
+        latency_ms=latency_ms,
+        latency_first_token_ms=latency_first_token_ms,
+        output_tokens=output_tokens,
+        queued_wait_ms=queued_wait_ms,
+        retried=retried,
+    )
+    if sample is None:
+        return
+    sample_model, tokens_per_second = sample
+    now = float(clock.time())
+    runtime = runtime_map.setdefault(account_id, RuntimeState())
+    samples_by_model = runtime.tps_samples
+    if samples_by_model is None:
+        samples_by_model = runtime.tps_samples = {}
+    samples_by_model.setdefault(sample_model, []).append((now, tokens_per_second))
+    _prune(samples_by_model, now)
+
+
+def account_tps_estimate(runtime: RuntimeState | None, model: str | None, now: float) -> float | None:
+    """Median in-window throughput (tok/s) of ``runtime`` on ``model``; ``None`` below ``TPS_MIN_SAMPLES``."""
+    if runtime is None or not runtime.tps_samples or not model:
+        return None
+    samples = runtime.tps_samples.get(model)
+    if not samples:
+        return None
+    oldest_kept = now - TPS_SAMPLE_WINDOW_SECONDS
+    values = [tps for recorded_at, tps in samples if recorded_at >= oldest_kept]
+    if len(values) < TPS_MIN_SAMPLES:
+        return None
+    return float(median(values))
+
+
+def tps_weight_multiplier(account_tps: float, fleet_tps: float) -> float:
+    """Draw-weight multiplier in ``[floor, 1.0]`` for an account throughput relative to the fleet."""
+    if fleet_tps <= 0.0 or account_tps <= 0.0 or account_tps >= fleet_tps * (1.0 - TPS_WEIGHT_DEADBAND):
+        return 1.0
+    return max(TPS_WEIGHT_FLOOR, account_tps / fleet_tps)
+
+
+def fleet_tps_estimates(
+    runtime_by_account_id: Mapping[str, RuntimeState],
+    model: str | None,
+    now: float,
+) -> tuple[dict[str, float], float | None]:
+    """Per-account estimates on ``model`` and the fleet reference (their median).
+
+    The reference is ``None`` -- every multiplier neutral -- when no model is
+    requested or fewer than ``TTFT_MIN_ACCOUNTS`` accounts hold enough samples
+    for that model; evidence on other models never contributes.
+    """
+    if not model:
+        return {}, None
+    estimates = {
+        account_id: estimate
+        for account_id, runtime in runtime_by_account_id.items()
+        if (estimate := account_tps_estimate(runtime, model, now)) is not None
+    }
+    fleet_tps = median(estimates.values()) if len(estimates) >= TTFT_MIN_ACCOUNTS else None
+    return estimates, fleet_tps

app/modules/proxy/_load_balancer/ttft_cohort.py

@@ -0,0 +1,214 @@
+"""Replica-local, fleet-relative first-token latency weight per account.
+
+The balancer has no latency signal: two accounts with equal remaining credits
+receive the same share of fresh weighted draws even when one of them answers
+the same small prompts noticeably later (a cohort whose upstream reasons on a
+large share of turns, or one on a slower upstream path). Every conversation
+that lands there pays that cohort for its whole sticky lifetime.
+
+This module keeps a bounded list of recent first-token latencies per account
+for one narrow slice of traffic -- successful ``normal`` turns with a small
+uncached input, no or low reasoning effort, a single upstream attempt and no
+local queueing (the request kinds ``warmup``, ``compaction`` and
+``realtime_live``, any row that waited on the response-create gate or bridge
+queue, and any WebSocket/bridge row whose first-token latency spans a retried
+send or an account-capacity wait are excluded) -- and derives a
+draw-weight multiplier relative to the fleet: an account whose estimate sits
+above the fleet median by more than the deadband is discounted to
+``max(floor, fleet / account)``. ``latency_cohort.apply_latency_cohort_weights``
+combines it with the per-model output-throughput multiplier
+(``throughput_cohort.py``) as the minimum of the two and applies it to fresh
+weighted draws; sticky owners, deterministic probes and deterministic
+strategies never consult it. The signal is replica-local and never persisted.
+
+Estimator: the per-account estimate is the mean of the samples below the top
+decile (an upper-trimmed mean). The slow cohort observed in production is
+bimodal -- a large share of its turns reason before the first token -- so a
+plain median sits on the no-reasoning mode and misses the share, while a plain
+mean is dominated by a single 30 s stall; trimming the tail keeps the mode
+share visible and drops the stall. The fleet reference is the median of the
+per-account estimates, so a uniformly slow fleet (upstream-wide degradation)
+stays neutral. Estimator changes are one-function edits here, not settings.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping
+from statistics import median
+from typing import Any
+
+from app.modules.proxy._load_balancer.types import RuntimeState
+
+# One hour of eligible turns: long enough for a quiet account to keep evidence,
+# short enough that an account changing cohort is re-weighted within the hour.
+TTFT_SAMPLE_WINDOW_SECONDS = 3600.0
+# Memory bound per account when an account exceeds ~one eligible turn a
+# minute; the window bounds staleness, the cap bounds the list.
+TTFT_MAX_SAMPLES = 64
+# Below this many samples an account's estimate is noise and stays neutral.
+TTFT_MIN_SAMPLES = 8
+# Below this many accounts with evidence there is no fleet to be relative to.
+TTFT_MIN_ACCOUNTS = 3
+# Only the small-input slice is comparable across accounts; the cap applies
+# to uncached input tokens (input minus cached prefix) so long-lived sticky
+# sessions with a large cached prefix keep sampling.
+TTFT_SAMPLE_MAX_INPUT_TOKENS = 20_000
+TTFT_SAMPLE_EFFORTS: frozenset[str | None] = frozenset({None, "minimal", "low"})
+TTFT_SAMPLE_REQUEST_KIND = "normal"
+# Share of the slowest samples dropped from each account's estimate.
+TTFT_TRIM_FRACTION = 0.1
+# An account within 15% of the fleet reference is not discounted at all.
+TTFT_WEIGHT_DEADBAND = 0.15
+# A 10x outlier still keeps half of its weight: the window keeps sampling it
+# and the discount lifts as soon as its samples recover.
+TTFT_WEIGHT_FLOOR = 0.5
+
+
+def _prune(samples: list[tuple[float, int]], now: float) -> None:
+    oldest_kept = now - TTFT_SAMPLE_WINDOW_SECONDS
+    stale = 0
+    for recorded_at, _ in samples:
+        if recorded_at >= oldest_kept:
+            break
+        stale += 1
+    if stale:
+        del samples[:stale]
+    excess = len(samples) - TTFT_MAX_SAMPLES
+    if excess > 0:
+        del samples[:excess]
+
+
+def _eligible(
+    *,
+    account_id: str | None,
+    status: str,
+    request_kind: str,
+    latency_first_token_ms: int | None,
+    input_tokens: int | None,
+    cached_input_tokens: int | None,
+    reasoning_effort: str | None,
+    latency_upstream_send_ms: int | None,
+    queued_wait_ms: int,
+    retried: bool,
+) -> bool:
+    return (
+        bool(account_id)
+        and status == "success"
+        and request_kind == TTFT_SAMPLE_REQUEST_KIND
+        and not retried
+        and latency_first_token_ms is not None
+        and latency_upstream_send_ms is not None
+        and 0 <= latency_upstream_send_ms <= latency_first_token_ms
+        and input_tokens is not None
+        and input_tokens - (cached_input_tokens or 0) < TTFT_SAMPLE_MAX_INPUT_TOKENS
+        and reasoning_effort in TTFT_SAMPLE_EFFORTS
+        and queued_wait_ms <= 0
+    )
+
+
+def record_ttft_sample(
+    balancer: Any,
+    *,
+    account_id: str | None,
+    status: str,
+    request_kind: str,
+    latency_first_token_ms: int | None,
+    input_tokens: int | None,
+    cached_input_tokens: int | None = None,
+    reasoning_effort: str | None,
+    latency_upstream_send_ms: int | None,
+    queued_wait_ms: int,
+    retried: bool = False,
+) -> None:
+    """Record one eligible first-token latency for ``account_id`` at the balancer clock.
+
+    The sample is ``latency_first_token_ms - latency_upstream_send_ms``: the
+    row's first-token latency runs from the request start, and on the bridge
+    path the start precedes session lookup, reconnect, prewarm, image inlining
+    and payload slimming, so the account is only charged from its
+    ``response.create`` send. ``latency_upstream_send_ms`` is the start-to-send
+    offset (``0`` when the caller's clock already starts at the send, as the
+    HTTP stream's attempt clock does); a row without one (``None``) or with a
+    send after its first token has no usable anchor and is never sampled.
+    ``retried`` marks a row whose first-token latency spans more than one
+    upstream send or an account-capacity wait (bridge retries, transparent
+    direct-WebSocket replays after a dropped upstream, retries after a
+    capacity wait or an owner-pinned quota error); such a latency includes
+    the failed attempt and the reconnect and, after an account switch, would
+    be charged to the destination account, so it is never sampled.
+    ``queued_wait_ms`` is the response-create gate wait (which includes the
+    global response-create admission wait) plus the bridge-queue wait.
+
+    Called from the request-log funnel on stream close. Synchronous and
+    lock-free on purpose: there is no await between reading and writing the
+    per-account list, so on the single-threaded loop it cannot interleave
+    with selection, and the stream close never waits on the per-account lock
+    that ``record_errors`` holds across a database write. No-op when
+    ``balancer`` does not expose the runtime map or clock (test doubles).
+    """
+    runtime_map = getattr(balancer, "_runtime", None)
+    clock = getattr(balancer, "_clock", None)
+    if not isinstance(runtime_map, dict) or clock is None:
+        return
+    if not _eligible(
+        account_id=account_id,
+        status=status,
+        request_kind=request_kind,
+        latency_first_token_ms=latency_first_token_ms,
+        input_tokens=input_tokens,
+        cached_input_tokens=cached_input_tokens,
+        reasoning_effort=reasoning_effort,
+        latency_upstream_send_ms=latency_upstream_send_ms,
+        queued_wait_ms=queued_wait_ms,
+        retried=retried,
+    ):
+        return
+    if account_id is None or latency_first_token_ms is None or latency_upstream_send_ms is None:
+        return  # narrowed by ``_eligible``
+    now = float(clock.time())
+    runtime = runtime_map.setdefault(account_id, RuntimeState())
+    samples = runtime.ttft_samples
+    if samples is None:
+        samples = runtime.ttft_samples = []
+    samples.append((now, int(latency_first_token_ms - latency_upstream_send_ms)))
+    _prune(samples, now)
+
+
+def account_ttft_estimate_ms(runtime: RuntimeState | None, now: float) -> float | None:
+    """Upper-trimmed mean of the in-window samples; ``None`` below ``TTFT_MIN_SAMPLES``."""
+    if runtime is None or not runtime.ttft_samples:
+        return None
+    oldest_kept = now - TTFT_SAMPLE_WINDOW_SECONDS
+    values = sorted(ttft_ms for recorded_at, ttft_ms in runtime.ttft_samples if recorded_at >= oldest_kept)
+    if len(values) < TTFT_MIN_SAMPLES:
+        return None
+    kept = values[: len(values) - math.ceil(len(values) * TTFT_TRIM_FRACTION)]
+    return sum(kept) / len(kept)
+
+
+def ttft_weight_multiplier(account_ms: float, fleet_ms: float) -> float:
+    """Draw-weight multiplier in ``[floor, 1.0]`` for an account estimate relative to the fleet."""
+    if fleet_ms <= 0.0 or account_ms <= fleet_ms * (1.0 + TTFT_WEIGHT_DEADBAND):
+        return 1.0
+    return max(TTFT_WEIGHT_FLOOR, fleet_ms / account_ms)
+
+
+def fleet_ttft_estimates(
+    runtime_by_account_id: Mapping[str, RuntimeState],
+    now: float,
+) -> tuple[dict[str, float], float | None]:
+    """Per-account estimates and the fleet reference (their median).
+
+    Computed over the whole runtime map so a selection narrowed to a subset is
+    still weighed against the full pool; the reference is ``None`` -- every
+    multiplier neutral -- while fewer than ``TTFT_MIN_ACCOUNTS`` accounts hold
+    enough samples.
+    """
+    estimates = {
+        account_id: estimate
+        for account_id, runtime in runtime_by_account_id.items()
+        if (estimate := account_ttft_estimate_ms(runtime, now)) is not None
+    }
+    fleet_ms = median(estimates.values()) if len(estimates) >= TTFT_MIN_ACCOUNTS else None
+    return estimates, fleet_ms

app/modules/proxy/_load_balancer/tunables.py

@@ -0,0 +1,97 @@
+"""Dashboard-managed routing weights and overload isolation (C2-2 routing/overload).
+
+Five balancer knobs -- the overload isolation window, the error-rate weighting
+kill switch, the in-flight pressure penalty, the leased-token weight and the
+account lease TTL -- live in ``dashboard_settings`` with the process
+environment as the fallback the dashboard inherits while its column is NULL
+(``configuration-tiers``: code default < environment < dashboard).
+
+Selection is a hot path that must not touch the database and must not read
+settings while holding a runtime lock, so the values are resolved *once per
+selection or lease operation* from the ``SettingsCache`` snapshot the proxy
+service already holds for that stage (the same snapshot the concurrency caps
+come from), frozen into a :class:`RoutingTunables`, and threaded through the
+selection entry points next to ``routing_strategy`` and the caps; nothing is
+re-read inside a lock section. Stages of one request share the cache's TTL
+window. A caller without a snapshot (tests, tools) resolves from the
+environment alone.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from app.core.config.inheritable import resolve_inheritable
+from app.modules.proxy._load_balancer.types import AccountLeaseKind
+
+# Code defaults, kept in step with the ``Settings`` fields of the same name so
+# a caller without any settings object still gets the shipped behaviour.
+DEFAULT_OVERLOAD_ISOLATION_SECONDS = 1800
+DEFAULT_ERROR_RATE_WEIGHTING_ENABLED = True
+DEFAULT_INFLIGHT_PENALTY_PCT = 2.5
+DEFAULT_LEASE_TOKEN_WEIGHT = 1.0
+DEFAULT_LEASE_TTL_SECONDS = 900.0
+
+
+@dataclass(frozen=True, slots=True)
+class RoutingTunables:
+    """Effective routing/overload knobs for one selection or lease operation."""
+
+    overload_isolation_seconds: float = float(DEFAULT_OVERLOAD_ISOLATION_SECONDS)
+    error_rate_weighting_enabled: bool = DEFAULT_ERROR_RATE_WEIGHTING_ENABLED
+    inflight_penalty_pct: float = DEFAULT_INFLIGHT_PENALTY_PCT
+    lease_token_weight: float = DEFAULT_LEASE_TOKEN_WEIGHT
+    lease_ttl_seconds: float = DEFAULT_LEASE_TTL_SECONDS
+
+
+def resolve_routing_tunables(
+    dashboard_settings: object | None,
+    *,
+    startup_settings: object,
+) -> RoutingTunables:
+    """Resolve the five knobs as code default < environment < dashboard.
+
+    ``dashboard_settings`` is the cached ``dashboard_settings`` row (or ``None``
+    when the caller has no snapshot); ``startup_settings`` is the process
+    ``Settings``. Both are read with ``getattr`` so partial stand-ins used by
+    tests resolve to the code default for the fields they omit.
+    """
+
+    def pick[T: (int, float, bool)](name: str, default: T) -> T:
+        column = getattr(dashboard_settings, name, None)
+        environment = getattr(startup_settings, name, None)
+        return resolve_inheritable(column, environment, default).value
+
+    return RoutingTunables(
+        overload_isolation_seconds=float(pick("proxy_overload_isolation_seconds", DEFAULT_OVERLOAD_ISOLATION_SECONDS)),
+        error_rate_weighting_enabled=bool(
+            pick("proxy_account_error_rate_weighting_enabled", DEFAULT_ERROR_RATE_WEIGHTING_ENABLED)
+        ),
+        inflight_penalty_pct=float(pick("proxy_account_inflight_penalty_pct", DEFAULT_INFLIGHT_PENALTY_PCT)),
+        lease_token_weight=float(pick("proxy_account_lease_token_weight", DEFAULT_LEASE_TOKEN_WEIGHT)),
+        lease_ttl_seconds=float(pick("proxy_account_lease_ttl_seconds", DEFAULT_LEASE_TTL_SECONDS)),
+    )
+
+
+_ACCOUNT_STREAM_LEASE_STALE_GRACE_SECONDS = 60.0
+
+
+def account_lease_stale_ttl_seconds(
+    kind: AccountLeaseKind, settings: object, *, routing_tunables: RoutingTunables
+) -> float:
+    """Age after which a lease of ``kind`` is reclaimed as stale.
+
+    The lease TTL comes from the dashboard snapshot; a stream lease also
+    outlives the effective request budgets (``settings`` carries the C2-1
+    dashboard overlay) plus a grace period.
+    """
+    ttl_seconds = float(routing_tunables.lease_ttl_seconds)
+    if kind != "stream":
+        return ttl_seconds
+    valid_stream_budget_seconds = max(
+        ttl_seconds,
+        float(getattr(settings, "proxy_request_budget_seconds", ttl_seconds)),
+        float(getattr(settings, "http_responses_stream_request_budget_seconds", ttl_seconds)),
+        float(getattr(settings, "http_responses_session_bridge_request_budget_seconds", ttl_seconds)),
+    )
+    return max(ttl_seconds, valid_stream_budget_seconds + _ACCOUNT_STREAM_LEASE_STALE_GRACE_SECONDS)

app/modules/proxy/_load_balancer/types.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+AccountLeaseKind = Literal["response_create", "stream"]
+
+MAX_SELECTION_ATTEMPTS = 4
+
+
+@dataclass
+class RuntimeState:
+    reset_at: float | None = None
+    cooldown_until: float | None = None
+    last_error_at: float | None = None
+    last_selected_at: float | None = None
+    error_count: int = 0
+    version: int = 0
+    health_version: int = 0
+    blocked_at: float | None = None
+    health_tier: int = 0
+    drain_entered_at: float | None = None
+    probe_success_streak: int = 0
+    inflight_response_creates: int = 0
+    inflight_streams: int = 0
+    leased_tokens: float = 0.0
+    leases: dict[str, AccountLease] | None = None
+    stream_key_inflight: dict[str, int] | None = None
+    # Upstream-overload soft backoff (see ``_load_balancer/overload_backoff.py``).
+    # Recent ``server_is_overloaded`` rejection times inside the trip window,
+    # the deprioritization deadline once tripped, the exponential level, and
+    # when the level last tripped (for decay). Replica-local, never persisted.
+    overload_rejections: list[float] | None = None
+    # Recent *soft* overload observations: bare ``server_error`` terminals that
+    # upstream returns for the same admission-rejection condition but without
+    # the explicit overload code. Counted at a fractional weight so a genuine
+    # one-off fault never trips the window on its own, while a sustained
+    # ``server_error`` refusal still deprioritizes the account.
+    soft_overload_rejections: list[float] | None = None
+    overload_backoff_until: float | None = None
+    overload_backoff_level: int = 0
+    overload_last_trip_at: float | None = None
+    # Isolation stage deadline (subset of the backoff deadline): while set and
+    # in the future, soft sticky owners are rerouted too.
+    overload_isolated_until: float | None = None
+    # Short burst cooldown set by a code-less upstream HTTP 429 (per-account
+    # burst/concurrency rejection, no ``rate_limit_exceeded`` / usage code):
+    # fresh (unbound) selection and fresh sticky bindings avoid the account
+    # until this deadline while another candidate exists. Established sticky
+    # owners and hard continuity owners are untouched. Replica-local, never
+    # persisted, independent of ``cooldown_until`` / persisted RATE_LIMITED.
+    burst_backoff_until: float | None = None
+    # Recent upstream outcome window (see ``_load_balancer/error_rate.py``):
+    # minute bucket -> [successes, failures], pruned to the configured window.
+    # Feeds the selection weight multiplier; replica-local, never persisted.
+    outcome_buckets: dict[int, list[int]] | None = None
+    # Recent eligible first-token latencies (see ``_load_balancer/ttft_cohort.py``):
+    # ``(recorded_at, ttft_ms)`` pruned to the window/cap. Replica-local, never persisted.
+    ttft_samples: list[tuple[float, int]] | None = None
+    # Recent eligible output throughputs per model (see
+    # ``_load_balancer/throughput_cohort.py``): model -> ``(recorded_at, tokens_per_second)``
+    # pruned to the window/cap; models without in-window samples are dropped.
+    tps_samples: dict[str, list[tuple[float, float]]] | None = None
+    # Transition-log gates of ``_load_balancer/latency_cohort.py``: the last
+    # applied first-token multiplier (model-agnostic) and the last applied
+    # non-neutral throughput multiplier per model (neutral models are dropped,
+    # so the map only grows with models the account is discounted on).
+    ttft_weight: float = 1.0
+    tps_weights: dict[str, float] | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class ProbeReservation:
+    account_id: str
+    previous_last_selected_at: float | None
+    reserved_at: float
+    expected_runtime_version: int
+
+
+@dataclass(frozen=True, slots=True)
+class AccountLease:
+    lease_id: str
+    account_id: str
+    kind: AccountLeaseKind
+    acquired_at: float
+    estimated_tokens: float = 0.0
+    api_key_id: str | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class AccountConcurrencyCaps:
+    response_create_limit: int
+    stream_limit: int
+    configured_response_create_limit: int | None = None
+    configured_stream_limit: int | None = None
+    replica_count: int = 1