codexLB-up 1.25.0b10__py3-none-any.whl

app/core/utils/sse.py

@@ -0,0 +1,286 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import re
+from collections.abc import AsyncIterator, Callable, Mapping
+from typing import cast
+
+from app.core.errors import ResponseFailedEvent
+from app.core.types import JsonValue
+from app.core.utils.json_guards import is_json_dict
+from app.core.utils.shared_future import _await_task_deferring_cancellation, wait_on_shared_future
+
+type JsonPayload = Mapping[str, JsonValue] | ResponseFailedEvent
+
+# The SSE spec delimits lines only by CR, LF, or CRLF. str.splitlines() also
+# breaks on other Unicode boundaries (VT, FF, FS/GS/RS, NEL, U+2028, U+2029),
+# and U+2028/U+2029 are valid *unescaped* inside JSON strings, so splitting on
+# them would corrupt a data: payload that legitimately contains one.
+_SSE_LINE_BOUNDARY = re.compile(r"\r\n|\r|\n")
+
+SSE_KEEPALIVE_FRAME = ": keepalive\n\n"
+CODEX_KEEPALIVE_FRAME = 'event: codex.keepalive\ndata: {"type":"codex.keepalive"}\n\n'
+
+# The exact single-event shape ``format_sse_event`` emits (and the upstream
+# Codex backend sends): a leading ``event: <type>`` line, one JSON-object
+# ``data:`` line, LF-only framing, and a blank-line terminator. Blocks that
+# match can expose their event type without a JSON parse and are safe to
+# relay downstream byte-for-byte.
+_CANONICAL_SSE_BLOCK = re.compile(r"\Aevent: ([^\r\n]+)\ndata: \{[^\r\n]*\n\n\Z")
+
+
+def sse_event_type_from_block(event_block: str) -> str | None:
+    """Cheaply extract the event type from a canonically framed SSE block.
+
+    Returns the ``event:`` line's value only when the block matches the exact
+    shape ``format_sse_event`` produces (see ``_CANONICAL_SSE_BLOCK``).
+    Anything else — data-only blocks, multi-line data, CR/CRLF framing,
+    comment or ``id:`` lines, non-object data payloads, or an ``event:`` field
+    that appears after ``data:`` (legal SSE, but not canonical here) — returns
+    ``None`` so callers fall back to a full parse.
+    """
+    match = _CANONICAL_SSE_BLOCK.match(event_block)
+    if match is None:
+        return None
+    return match.group(1)
+
+
+async def inject_sse_keepalives(
+    source: AsyncIterator[str],
+    interval_seconds: float,
+    *,
+    keepalive_frame: str = SSE_KEEPALIVE_FRAME,
+    on_keepalive: Callable[[], None] | None = None,
+) -> AsyncIterator[str]:
+    """Wrap an SSE event iterator and emit comment heartbeats on idle gaps.
+
+    Comment frames (lines starting with ``:``) are mandated by the SSE spec to
+    be ignored by parsers, so they are safe to inject between event blocks.
+    They keep the TCP path warm so half-open sockets surface as write errors
+    instead of hanging forever, and let aggressive intermediaries see traffic.
+
+    A non-positive ``interval_seconds`` disables injection entirely.
+    """
+    iterator = source.__aiter__()
+    try:
+        if interval_seconds <= 0:
+            async for chunk in iterator:
+                yield chunk
+            return
+
+        async def _next_chunk(it: AsyncIterator[str]) -> str:
+            return await it.__anext__()
+
+        pending: asyncio.Task[str] | None = None
+        try:
+            while True:
+                if pending is None:
+                    pending = asyncio.create_task(_next_chunk(iterator))
+                try:
+                    # Not ``wait_for(shield(pending))``: Python 3.14's shield
+                    # leaks a done-callback onto ``pending`` at every keepalive
+                    # timeout, so a quiet upstream accumulates callbacks (and
+                    # O(n) remove scans) until the next chunk arrives.
+                    chunk = await wait_on_shared_future(
+                        pending,
+                        timeout=interval_seconds,
+                    )
+                except asyncio.TimeoutError:
+                    if on_keepalive is not None:
+                        on_keepalive()
+                    yield keepalive_frame
+                    continue
+                except StopAsyncIteration:
+                    pending = None
+                    break
+                pending = None
+                yield chunk
+        finally:
+            if pending is not None and not pending.done():
+                pending.cancel()
+                try:
+                    # Not ``await pending``: a level-cancelled consumer scope
+                    # re-cancels this task on every loop iteration, and a
+                    # direct task await cascades each cancel into ``pending``
+                    # and any cancellation-deferring wait beneath it (the
+                    # startup-probe task), spinning the loop until that
+                    # cleanup completes. The canonical helper shields this
+                    # task and waits through a proxy future, so the chunk task
+                    # sees only the explicit cancel above; it still settles
+                    # before the outer finalizers close the iterator chain the
+                    # chunk task is driving, and its eventual exception is
+                    # retrieved here rather than logged as never retrieved.
+                    await _await_task_deferring_cancellation(pending)
+                except BaseException:
+                    pass
+    finally:
+        aclose = getattr(iterator, "aclose", None)
+        if aclose is not None:
+            await aclose()
+        elif iterator is not source:
+            source_aclose = getattr(source, "aclose", None)
+            if source_aclose is not None:
+                await source_aclose()
+
+
+def format_sse_event(payload: JsonPayload) -> str:
+    data = json.dumps(payload, ensure_ascii=True, separators=(",", ":"))
+    event_type = payload.get("type")
+    if isinstance(event_type, str) and event_type:
+        return f"event: {event_type}\ndata: {data}\n\n"
+    return f"data: {data}\n\n"
+
+
+def _is_single_line_json_object_text(text: str) -> bool:
+    return text.startswith("{") and "\n" not in text and "\r" not in text
+
+
+def format_sse_event_from_text(payload: Mapping[str, JsonValue], text: str) -> str:
+    """Frame an already-serialized JSON object without re-dumping it.
+
+    ``text`` MUST be a JSON serialization of ``payload`` (the upstream frame the
+    payload was parsed from, or a compact re-dump of it). The result carries the
+    same ``event:`` line ``format_sse_event(payload)`` would emit, but its
+    ``data:`` line is ``text`` verbatim, so non-ASCII stays UTF-8 instead of
+    being ``\\uXXXX``-escaped and the JSON is never re-serialized. Anything
+    that would not produce a canonical single-line block (multi-line or
+    whitespace-prefixed text) falls back to ``format_sse_event``.
+    """
+    if not _is_single_line_json_object_text(text):
+        return format_sse_event(payload)
+    event_type = payload.get("type")
+    if isinstance(event_type, str) and event_type:
+        return f"event: {event_type}\ndata: {text}\n\n"
+    return f"data: {text}\n\n"
+
+
+def format_sse_data(payload: Mapping[str, JsonValue]) -> str:
+    data = json.dumps(payload, ensure_ascii=True, separators=(",", ":"))
+    return f"data: {data}\n\n"
+
+
+class ParsedSseBlock(str):
+    """An SSE event block that carries its already-parsed ``data:`` payload.
+
+    Behaves exactly like ``str`` (framing, encoding, comparison, ``startswith``);
+    ``payload`` is what ``parse_sse_data_json`` returns for the same text, so a
+    chain of stream stages parses each block once. The payload object is shared
+    by every stage that receives the block: consumers MUST treat it as read-only
+    and copy before mutating. Derived strings (``strip()``, concatenation) are
+    plain ``str`` and take the full parse path again. ``is_local`` records a
+    producer-generated block without changing its wire payload or retry markers.
+    ``response_id_is_local`` also covers a local ID added to a real upstream event.
+    """
+
+    # ``__slots__`` is not supported on ``str`` subclasses; the instance dict
+    # costs ~100-200 B per block, far less than a redundant JSON parse.
+    payload: dict[str, JsonValue] | None
+    is_local: bool
+    response_id_is_local: bool
+
+    def __new__(
+        cls,
+        block: str,
+        payload: dict[str, JsonValue] | None,
+        *,
+        is_local: bool = False,
+        response_id_is_local: bool = False,
+    ) -> ParsedSseBlock:
+        instance = super().__new__(cls, block)
+        instance.payload = payload
+        instance.is_local = is_local
+        instance.response_id_is_local = response_id_is_local
+        return instance
+
+
+def format_local_sse_event(payload: dict[str, JsonValue] | ResponseFailedEvent) -> ParsedSseBlock:
+    """Frame a locally generated event with provenance confined to the carrier."""
+    return ParsedSseBlock(
+        format_sse_event(payload), cast(dict[str, JsonValue], payload), is_local=True, response_id_is_local=True
+    )
+
+
+def sse_block_with_payload(event_block: str, payload: dict[str, JsonValue] | None) -> ParsedSseBlock:
+    """Attach ``payload`` (the result of ``parse_sse_data_json(event_block)``) to the block."""
+    if isinstance(event_block, ParsedSseBlock) and event_block.payload is payload:
+        return event_block
+    return ParsedSseBlock(
+        event_block,
+        payload,
+        is_local=isinstance(event_block, ParsedSseBlock) and event_block.is_local,
+        response_id_is_local=isinstance(event_block, ParsedSseBlock) and event_block.response_id_is_local,
+    )
+
+
+def parse_sse_data_json_text(text: str) -> dict[str, JsonValue] | None:
+    """Parse one ``data:`` value as a JSON object.
+
+    Equivalent to ``parse_sse_data_json(f"data: {text}\\n\\n")`` but skips SSE
+    line parsing when ``text`` is a single-line JSON object — the shape every
+    upstream websocket frame has. Other inputs (blank, ``[DONE]``, multi-line or
+    whitespace-prefixed text) take the exact SSE field path so the ``None`` and
+    line-joining semantics stay identical.
+    """
+    if _is_single_line_json_object_text(text):
+        try:
+            payload = json.loads(text)
+        except json.JSONDecodeError:
+            return None
+        if is_json_dict(payload):
+            return payload
+        return None
+    return parse_sse_data_json(f"data: {text}\n\n")
+
+
+def parse_sse_data_json(event_block: str) -> dict[str, JsonValue] | None:
+    if isinstance(event_block, ParsedSseBlock):
+        return event_block.payload
+    data = extract_sse_data(event_block)
+    if data is None:
+        return None
+    try:
+        payload = json.loads(data)
+    except json.JSONDecodeError:
+        return None
+    if is_json_dict(payload):
+        return payload
+    return None
+
+
+def extract_sse_data(event_block: str) -> str | None:
+    data_lines = _extract_sse_data_lines(event_block)
+    if data_lines is None:
+        return None
+    data = "\n".join(data_lines)
+    if not data.strip():
+        return None
+    if data.strip() == "[DONE]":
+        return None
+    return data
+
+
+def _extract_sse_data_lines(event_block: str) -> list[str] | None:
+    data_lines: list[str] = []
+    for raw_line in _SSE_LINE_BOUNDARY.split(event_block):
+        if not raw_line:
+            continue
+        if raw_line.startswith(":"):
+            continue
+
+        field, value = _parse_sse_field(raw_line)
+        if field == "data":
+            data_lines.append(value)
+
+    if not data_lines:
+        return None
+    return data_lines
+
+
+def _parse_sse_field(line: str) -> tuple[str, str]:
+    if ":" not in line:
+        return line, ""
+    field, value = line.split(":", 1)
+    if value.startswith(" "):
+        value = value[1:]
+    return field, value

app/core/utils/time.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+
+
+def utcnow() -> datetime:
+    return datetime.now(timezone.utc).replace(tzinfo=None)
+
+
+def to_utc_naive(value: datetime) -> datetime:
+    if value.tzinfo is None:
+        return value
+    return value.astimezone(timezone.utc).replace(tzinfo=None)
+
+
+def naive_utc_to_epoch(dt: datetime) -> int:
+    """Convert a naive-UTC datetime to Unix epoch seconds."""
+    return int(dt.replace(tzinfo=timezone.utc).timestamp())
+
+
+def from_epoch_seconds(value: int | float | None) -> datetime | None:
+    if value is None:
+        return None
+    return datetime.fromtimestamp(value, tz=timezone.utc)

app/db/account_identity_lock.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from collections.abc import Collection
+from hashlib import sha256
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+
+_POSTGRES_ACCOUNT_IDENTITY_LOCK_TIMEOUT_MS = 30_000
+
+
+def advisory_lock_key(scope: str, value: str) -> int:
+    digest = sha256(f"{scope}:{value}".encode("utf-8")).digest()
+    return int.from_bytes(digest[:8], byteorder="big", signed=True)
+
+
+def account_identity_lock_key(chatgpt_account_id: str) -> int:
+    """Return the existing PostgreSQL lock namespace for one upstream identity."""
+    return advisory_lock_key("account-id", f"chatgpt:{chatgpt_account_id}")
+
+
+async def lock_postgresql_account_identities(
+    session: AsyncSession,
+    chatgpt_account_ids: Collection[str | None],
+) -> tuple[int, ...]:
+    """Lock upstream identity membership in canonical transaction-scoped order."""
+    bind = session.get_bind()
+    if bind is None or bind.dialect.name != "postgresql":
+        return ()
+    lock_keys = tuple(
+        sorted(
+            {
+                account_identity_lock_key(chatgpt_account_id)
+                for chatgpt_account_id in chatgpt_account_ids
+                if chatgpt_account_id
+            }
+        )
+    )
+    try:
+        if lock_keys:
+            # Match the database layer's existing 30-second contention budget.
+            # Transaction-local scope bounds every subsequent lock in this unit
+            # of work and PostgreSQL restores it at transaction end.
+            await session.execute(
+                text("SELECT set_config('lock_timeout', :timeout, true)"),
+                {"timeout": f"{_POSTGRES_ACCOUNT_IDENTITY_LOCK_TIMEOUT_MS}ms"},
+            )
+        for lock_key in lock_keys:
+            await session.execute(
+                text("SELECT pg_advisory_xact_lock(:lock_key)"),
+                {"lock_key": lock_key},
+            )
+    except BaseException:
+        await session.rollback()
+        raise
+    return lock_keys

app/db/alembic/__init__.py

@@ -0,0 +1 @@
+from __future__ import annotations

app/db/alembic/env.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from logging.config import fileConfig
+
+from alembic import context
+from sqlalchemy import engine_from_config, pool
+
+from app.core.config.settings import get_settings
+from app.db.migration_url import to_sync_database_url
+from app.db.models import Base
+
+config = context.config
+
+if config.config_file_name is not None:
+    fileConfig(config.config_file_name)
+
+target_metadata = Base.metadata
+
+
+def _sync_database_url() -> str:
+    configured = config.get_main_option("sqlalchemy.url")
+    if configured:
+        return configured
+    return to_sync_database_url(get_settings().database_url)
+
+
+def run_migrations_offline() -> None:
+    url = _sync_database_url()
+    config.set_main_option("sqlalchemy.url", url)
+
+    context.configure(
+        url=url,
+        target_metadata=target_metadata,
+        literal_binds=True,
+        dialect_opts={"paramstyle": "named"},
+        compare_type=True,
+        render_as_batch=url.startswith("sqlite"),
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def run_migrations_online() -> None:
+    section = config.get_section(config.config_ini_section, {})
+    section["sqlalchemy.url"] = _sync_database_url()
+
+    connectable = engine_from_config(
+        section,
+        prefix="sqlalchemy.",
+        poolclass=pool.NullPool,
+        future=True,
+    )
+
+    with connectable.connect() as connection:
+        context.configure(
+            connection=connection,
+            target_metadata=target_metadata,
+            compare_type=True,
+            render_as_batch=connection.dialect.name == "sqlite",
+        )
+
+        with context.begin_transaction():
+            context.run_migrations()
+
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()

app/db/alembic/revision_ids.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import re
+
+REVISION_ID_PATTERN = re.compile(r"^\d{8}_\d{6}_[a-z0-9_]+$")
+
+OLD_TO_NEW_REVISION_MAP: dict[str, str] = {
+    "000_base_schema": "20260213_000000_base_schema",
+    "001_normalize_account_plan_types": "20260213_000100_normalize_account_plan_types",
+    "002_add_request_logs_reasoning_effort": "20260213_000200_add_request_logs_reasoning_effort",
+    "003_add_accounts_reset_at": "20260213_000300_add_accounts_reset_at",
+    "004_add_accounts_chatgpt_account_id": "20260213_000400_add_accounts_chatgpt_account_id",
+    "005_add_dashboard_settings": "20260213_000500_add_dashboard_settings",
+    "006_add_dashboard_settings_totp": "20260213_000600_add_dashboard_settings_totp",
+    "007_add_dashboard_settings_password": "20260213_000700_add_dashboard_settings_password",
+    "008_add_api_keys": "20260213_000800_add_api_keys",
+    "009_add_api_key_limits": "20260214_000000_add_api_key_limits",
+    "010_add_idx_logs_requested_at": "20260215_000000_add_idx_logs_requested_at",
+    "011_add_api_key_usage_reservations": "20260218_000000_add_api_key_usage_reservations",
+    "012_add_import_without_overwrite_and_drop_accounts_email_unique": (
+        "20260218_000100_add_import_without_overwrite_and_drop_accounts_email_unique"
+    ),
+    "013_add_dashboard_settings_routing_strategy": "20260225_000000_add_dashboard_settings_routing_strategy",
+    "014_add_api_firewall_allowlist": "20260228_030000_add_api_firewall_allowlist",
+    "013_add_api_key_enforcement_fields": (
+        "20260218_000100_add_import_without_overwrite_and_drop_accounts_email_unique"
+    ),
+    "20260410_020000_restore_import_without_overwrite_default_false": "20260409_020000_fix_http_bridge_last_seen_index",
+    "20260525_000000_merge_routing_settings_security_heads": "20260513_000000_add_accounts_alias",
+}
+
+LEGACY_MIGRATION_TO_NEW_REVISION: dict[str, str] = {
+    "001_normalize_account_plan_types": OLD_TO_NEW_REVISION_MAP["001_normalize_account_plan_types"],
+    "002_add_request_logs_reasoning_effort": OLD_TO_NEW_REVISION_MAP["002_add_request_logs_reasoning_effort"],
+    "003_add_accounts_reset_at": OLD_TO_NEW_REVISION_MAP["003_add_accounts_reset_at"],
+    "004_add_accounts_chatgpt_account_id": OLD_TO_NEW_REVISION_MAP["004_add_accounts_chatgpt_account_id"],
+    "005_add_dashboard_settings": OLD_TO_NEW_REVISION_MAP["005_add_dashboard_settings"],
+    "006_add_dashboard_settings_totp": OLD_TO_NEW_REVISION_MAP["006_add_dashboard_settings_totp"],
+    "007_add_dashboard_settings_password": OLD_TO_NEW_REVISION_MAP["007_add_dashboard_settings_password"],
+    "008_add_api_keys": OLD_TO_NEW_REVISION_MAP["008_add_api_keys"],
+    "009_add_api_key_limits": OLD_TO_NEW_REVISION_MAP["009_add_api_key_limits"],
+    "010_add_idx_logs_requested_at": OLD_TO_NEW_REVISION_MAP["010_add_idx_logs_requested_at"],
+}