codexLB-up 1.25.0b10__py3-none-any.whl

app/core/utils/proxy_env.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import os
+import urllib.request
+from collections.abc import Mapping
+from urllib.parse import urlparse, urlunparse
+
+_WEBSOCKET_PROXY_ENV_PRIORITY: dict[str, tuple[str, ...]] = {
+    "ws": (
+        "ws",
+        "socks",
+        "https",
+        "http",
+        "all",
+    ),
+    "wss": (
+        "wss",
+        "socks",
+        "https",
+        "all",
+    ),
+}
+_HTTP_PROXY_ENV_PRIORITY: dict[str, tuple[str, ...]] = {
+    "http": ("socks", "http", "all"),
+    "https": ("socks", "https", "all"),
+}
+
+STANDARD_OUTBOUND_PROXY_ENV_NAMES: tuple[str, ...] = tuple(
+    dict.fromkeys(f"{name}_proxy" for names in _WEBSOCKET_PROXY_ENV_PRIORITY.values() for name in names)
+)
+
+
+def outbound_proxy_env_configured(environ: Mapping[str, str | None] = os.environ) -> bool:
+    return any(
+        name.lower() in STANDARD_OUTBOUND_PROXY_ENV_NAMES and value is not None and value.strip()
+        for name, value in environ.items()
+    )
+
+
+def resolve_websocket_proxy_from_env(url: str, environ: Mapping[str, str | None] = os.environ) -> str | None:
+    return _resolve_proxy_from_env(url, _WEBSOCKET_PROXY_ENV_PRIORITY, environ)
+
+
+def resolve_http_proxy_from_env(url: str, environ: Mapping[str, str | None] = os.environ) -> str | None:
+    return _resolve_proxy_from_env(url, _HTTP_PROXY_ENV_PRIORITY, environ)
+
+
+def _resolve_proxy_from_env(
+    url: str,
+    priority: Mapping[str, tuple[str, ...]],
+    environ: Mapping[str, str | None],
+) -> str | None:
+    parsed = urlparse(url)
+    scheme = parsed.scheme.lower()
+    env_names = priority.get(scheme)
+    if env_names is None:
+        return None
+
+    hostname = parsed.hostname
+    port = parsed.port or (443 if scheme in {"https", "wss"} else 80)
+    proxies = _sanitized_proxy_env(environ)
+    proxy_bypass_environment = getattr(urllib.request, "proxy_bypass_environment")
+    if hostname and proxy_bypass_environment(f"{hostname}:{port}", proxies):
+        return None
+
+    for name in env_names:
+        proxy_url = proxies.get(name)
+        if proxy_url is not None:
+            return proxy_url
+    return None
+
+
+def _sanitized_proxy_env(environ: Mapping[str, str | None] = os.environ) -> dict[str, str]:
+    proxies: dict[str, str] = {}
+    for name, value in environ.items():
+        if name == "HTTP_PROXY" and environ.get("REQUEST_METHOD"):
+            continue
+        _add_proxy_env_value(proxies, name, value)
+    for name, value in environ.items():
+        if name == "HTTP_PROXY" and environ.get("REQUEST_METHOD"):
+            continue
+        if name == name.lower():
+            _add_proxy_env_value(proxies, name, value)
+    return proxies
+
+
+def _add_proxy_env_value(proxies: dict[str, str], name: str, value: str | None) -> None:
+    normalized_name = name.lower()
+    if not normalized_name.endswith("_proxy"):
+        return
+    proxy_name = normalized_name.removesuffix("_proxy")
+    if not value:
+        proxies.pop(proxy_name, None)
+        return
+    proxies[proxy_name] = _normalize_proxy_url(proxy_name, value.strip())
+
+
+def _normalize_proxy_url(proxy_name: str, proxy_url: str) -> str:
+    if proxy_name != "socks":
+        return proxy_url
+    parsed = urlparse(proxy_url)
+    if parsed.scheme not in {"http", "https"}:
+        return proxy_url
+    return urlunparse(parsed._replace(scheme="socks5h"))

app/core/utils/request_id.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from contextvars import ContextVar, Token
+from uuid import uuid4
+
+_REQUEST_ID: ContextVar[str | None] = ContextVar("request_id", default=None)
+_REQUEST_SCOPE_ID: ContextVar[str | None] = ContextVar("request_scope_id", default=None)
+
+
+def get_request_id() -> str | None:
+    return _REQUEST_ID.get()
+
+
+def set_request_id(value: str | None) -> Token[str | None]:
+    return _REQUEST_ID.set(value)
+
+
+def reset_request_id(token: Token[str | None]) -> None:
+    _REQUEST_ID.reset(token)
+
+
+def clear_request_id() -> None:
+    _REQUEST_ID.set(None)
+
+
+def ensure_request_id(value: str | None = None) -> str:
+    """Return the active request ID, storing an explicit or generated value when needed."""
+    if value:
+        _REQUEST_ID.set(value)
+        return value
+    current = _REQUEST_ID.get()
+    if current:
+        return current
+    generated = str(uuid4())
+    _REQUEST_ID.set(generated)
+    return generated
+
+
+def get_request_scope_id() -> str | None:
+    return _REQUEST_SCOPE_ID.get()
+
+
+def set_request_scope_id(value: str | None) -> Token[str | None]:
+    return _REQUEST_SCOPE_ID.set(value)
+
+
+def reset_request_scope_id(token: Token[str | None]) -> None:
+    _REQUEST_SCOPE_ID.reset(token)
+
+
+def clear_request_scope_id() -> None:
+    _REQUEST_SCOPE_ID.set(None)
+
+
+def ensure_request_scope_id(value: str | None = None) -> str:
+    """Return a server-owned ID that is unique to one inbound request scope."""
+    if value:
+        _REQUEST_SCOPE_ID.set(value)
+        return value
+    current = _REQUEST_SCOPE_ID.get()
+    if current:
+        return current
+    generated = str(uuid4())
+    _REQUEST_SCOPE_ID.set(generated)
+    return generated

app/core/utils/retry.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import random
+import re
+
+# Units the retry hint may use, longest-first within each family so that, for
+# example, ``ms`` is preferred over ``m`` and ``minutes`` over ``m``.
+_UNIT_ALTERNATION = r"ms|milliseconds?|hours?|hrs?|h|minutes?|mins?|m|seconds?|secs?|s"
+
+# A unit literal is only valid when it is not immediately followed by another
+# letter, otherwise the single-letter alternatives swallow the prefix of an
+# unsupported longer word (``m`` from ``month``, ``h`` from ``half``) and the
+# hint is silently mis-scaled. Digits still follow units in compound hints
+# (``6m0s``), so the boundary forbids letters only, not digits or whitespace.
+_UNIT_BOUNDARY = r"(?![A-Za-z])"
+
+# Seconds-per-unit for every literal ``_UNIT_ALTERNATION`` can capture.
+_UNIT_SECONDS: dict[str, float] = {
+    "ms": 0.001,
+    "millisecond": 0.001,
+    "milliseconds": 0.001,
+    "h": 3600.0,
+    "hr": 3600.0,
+    "hrs": 3600.0,
+    "hour": 3600.0,
+    "hours": 3600.0,
+    "m": 60.0,
+    "min": 60.0,
+    "mins": 60.0,
+    "minute": 60.0,
+    "minutes": 60.0,
+    "s": 1.0,
+    "sec": 1.0,
+    "secs": 1.0,
+    "second": 1.0,
+    "seconds": 1.0,
+}
+
+# Capture the contiguous run of ``<number><unit>`` components that immediately
+# follows "try again in" so compound hints such as ``6m0s`` or ``1h2m3s`` are
+# read in full instead of stopping at the first unit.
+_RETRY_PATTERN = re.compile(rf"(?i)try again in\s*((?:\d+(?:\.\d+)?\s*(?:{_UNIT_ALTERNATION}){_UNIT_BOUNDARY}\s*)+)")
+_DURATION_TOKEN = re.compile(rf"(?i)(\d+(?:\.\d+)?)\s*({_UNIT_ALTERNATION}){_UNIT_BOUNDARY}")
+
+_BACKOFF_INITIAL_DELAY_MS = 200
+_BACKOFF_FACTOR = 2.0
+_BACKOFF_JITTER_MIN = 0.9
+_BACKOFF_JITTER_MAX = 1.1
+
+
+def parse_retry_after(message: str) -> float | None:
+    match = _RETRY_PATTERN.search(message or "")
+    if not match:
+        return None
+    total = 0.0
+    matched = False
+    for value, unit in _DURATION_TOKEN.findall(match.group(1)):
+        multiplier = _UNIT_SECONDS.get(unit.lower())
+        if multiplier is None:
+            continue
+        if unit.lower() == "ms":
+            total += float(value) / 1000
+        else:
+            total += float(value) * multiplier
+        matched = True
+    if not matched:
+        return None
+    return total
+
+
+def backoff_seconds(attempt: int) -> float:
+    if attempt < 1:
+        attempt = 1
+    exponent = _BACKOFF_FACTOR ** (attempt - 1)
+    base_ms = _BACKOFF_INITIAL_DELAY_MS * exponent
+    jitter = random.uniform(_BACKOFF_JITTER_MIN, _BACKOFF_JITTER_MAX)
+    return (base_ms * jitter) / 1000.0

app/core/utils/shared_future.py

@@ -0,0 +1,172 @@
+"""Await shared futures without per-waiter callbacks on the shared object.
+
+``asyncio.wait_for(asyncio.shield(shared), timeout)`` attaches done callbacks
+to ``shared`` for every waiter and removes them with O(n) list scans when a
+waiter is cancelled or times out. With many waiters piled onto one long-lived
+future (the http-bridge inflight/capacity registries, refresh singleflight),
+a mass timeout turns the event loop into an O(N^2) callback grinder. Python
+3.14's ``shield`` additionally leaks one ``_clear_awaited_by_callback`` per
+attempt onto the still-pending future, so each retry cycle makes every later
+scan more expensive. In the 2026-08-20 production incident this starved the
+event loop for hours (98% of GIL samples inside ``Future.remove_done_callback``)
+with zero client sessions attached.
+
+``wait_on_shared_future`` keeps exactly one fan-out callback on the shared
+future regardless of waiter count. Each waiter awaits its own single-use proxy
+future, so waiter timeout and cancellation are O(1) set operations that never
+touch the shared future's callback list.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Coroutine
+from typing import Any, TypeVar, cast
+
+import anyio
+from anyio.lowlevel import checkpoint_if_cancelled
+
+from app.core.clock import REAL_SCHEDULER, Scheduler
+
+_T = TypeVar("_T")
+_TaskResultT = TypeVar("_TaskResultT")
+
+_WAITERS_ATTR = "_shared_future_fanout_waiters"
+
+
+def _fan_out(shared: "asyncio.Future[_T]", waiters: "set[asyncio.Future[_T]]") -> None:
+    cancelled = shared.cancelled()
+    # Read the exception here, before the loop, so it is consumed whether or
+    # not a waiter is still registered. A waiter that timed out discards its
+    # proxy, so a shared future completing in that gap reaches this callback
+    # with an empty set -- and an exception nobody retrieved makes asyncio log
+    # "Task exception was never retrieved" when the task is collected. In
+    # production that surfaced as an ERROR + traceback per SSE stream whose
+    # source ended (``StopAsyncIteration``) while the keepalive injector was
+    # between waits, on requests that had otherwise succeeded.
+    exc = None if cancelled else shared.exception()
+    for waiter in waiters:
+        if waiter.done():
+            continue
+        if cancelled:
+            waiter.cancel()
+            continue
+        if exc is not None:
+            waiter.set_exception(exc)
+            # Consume eagerly: a waiter whose task was cancelled between this
+            # fan-out and its resumption would otherwise log
+            # "exception was never retrieved" from the proxy destructor.
+            waiter.exception()
+        else:
+            waiter.set_result(shared.result())
+    waiters.clear()
+
+
+async def wait_on_shared_future(
+    shared: "asyncio.Future[_T]",
+    *,
+    timeout: float | None = None,
+    scheduler: Scheduler = REAL_SCHEDULER,
+) -> _T:
+    """Drop-in equivalent of ``wait_for(shield(shared), timeout)`` for futures
+    awaited by many concurrent waiters.
+
+    - ``shared``'s result, exception, or cancellation propagates to every
+      waiter exactly as with ``shield``.
+    - ``timeout`` raises ``TimeoutError``; ``shared`` is never cancelled or
+      otherwise mutated by a waiter timing out or being cancelled.
+    - Cancelling the awaiting task detaches its proxy in O(1) and leaves
+      ``shared`` (and the work it represents) running.
+    - ``scheduler`` only matters for tests: a timed wait runs through its
+      ``wait_for`` so virtual time can expire it. The production default is
+      the real scheduler, i.e. ``asyncio.wait_for`` verbatim.
+    """
+    if shared.done():
+        return shared.result()
+    waiters: set[asyncio.Future[_T]] | None = getattr(shared, _WAITERS_ATTR, None)
+    if waiters is None:
+        # No await between the ``done()`` check above and this registration,
+        # so the fan-out callback cannot have fired with an empty set.
+        waiters = set()
+        setattr(shared, _WAITERS_ATTR, waiters)
+        shared.add_done_callback(lambda done, _waiters=waiters: _fan_out(done, _waiters))
+    proxy: asyncio.Future[_T] = asyncio.get_running_loop().create_future()
+    waiters.add(proxy)
+    try:
+        if timeout is None:
+            return await proxy
+        return await scheduler.wait_for(proxy, timeout)
+    finally:
+        waiters.discard(proxy)
+
+
+async def _await_task_deferring_cancellation(
+    task: asyncio.Task[_TaskResultT],
+) -> tuple[_TaskResultT, asyncio.CancelledError | None]:
+    """Finish critical cleanup while preserving the caller's cancellation."""
+
+    cancellation: asyncio.CancelledError | None = None
+    result: _TaskResultT | None = None
+    # The anyio shield keeps a level-cancelled Starlette scope from re-raising
+    # into every ``await``, which would otherwise busy-spin this loop until the
+    # owned task completes. The shield only covers the task anyio tracks: when
+    # this helper runs inside a plain ``asyncio.create_task`` task that an
+    # anyio-tracked task awaits directly (``await task``), each level re-cancel
+    # of the tracked task cascades down the ``_fut_waiter`` chain and re-enters
+    # this loop anyway. Callers that hand such a task across a task boundary
+    # must await it through ``wait_on_shared_future`` so the proxy future, not
+    # the owned task, absorbs the repeated cancels (2026-09-07 production
+    # busy spin). ``wait_on_shared_future`` keeps the loop's waits
+    # off the task's done-callback list: Python 3.14's ``asyncio.shield``
+    # leaks a callback per cancelled wait, so re-shielding a task wedged on a
+    # lock grew 100k+ callbacks and O(n^2) remove scans in the 2026-08-30
+    # production event-loop livelock.
+    with anyio.CancelScope(shield=True):
+        while True:
+            try:
+                result = await wait_on_shared_future(task)
+                break
+            except asyncio.CancelledError as exc:
+                if task.cancelled():
+                    raise
+                cancellation = cancellation or exc
+    if cancellation is None:
+        # The shield also blocks the level cancellation this helper promises
+        # to surface. Probe for it without suspending so callers still get
+        # their cancellation marker after the owned task finished.
+        try:
+            await checkpoint_if_cancelled()
+        except asyncio.CancelledError as exc:
+            cancellation = exc
+    return cast(_TaskResultT, result), cancellation
+
+
+async def _await_result_deferring_cancellation(
+    awaitable: "Awaitable[_TaskResultT]",
+    *,
+    scheduler: Scheduler = REAL_SCHEDULER,
+) -> tuple[_TaskResultT, asyncio.CancelledError | None]:
+    """``_await_task_deferring_cancellation`` for a bare awaitable.
+
+    A coroutine is spawned through ``scheduler`` (``asyncio.create_task``
+    under the real default, an owned task under a simulation); any other
+    awaitable is wrapped by ``asyncio.ensure_future`` as before.
+    """
+
+    task: asyncio.Task[_TaskResultT]
+    if asyncio.iscoroutine(awaitable):
+        task = scheduler.create_task(cast(Coroutine[Any, Any, _TaskResultT], awaitable))
+    else:
+        task = asyncio.ensure_future(awaitable)
+    return await _await_task_deferring_cancellation(task)
+
+
+async def _await_cleanup_deferring_cancellation(
+    awaitable: "Awaitable[object]",
+    *,
+    scheduler: Scheduler = REAL_SCHEDULER,
+) -> asyncio.CancelledError | None:
+    """Finish required cleanup, returning the deferred cancellation marker."""
+
+    _, cancellation = await _await_result_deferring_cancellation(awaitable, scheduler=scheduler)
+    return cancellation