codex-lb 0.3.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

app/modules/proxy/load_balancer.py

@@ -16,8 +16,8 @@ from app.core.balancer.types import UpstreamError
 from app.core.usage.quota import apply_usage_quota
 from app.db.models import Account, UsageHistory
 from app.modules.accounts.repository import AccountsRepository
+from app.modules.proxy.repo_bundle import ProxyRepoFactory
 from app.modules.proxy.sticky_repository import StickySessionsRepository
-from app.modules.usage.repository import UsageRepository
 from app.modules.usage.updater import UsageUpdater
 
 
@@ -37,16 +37,8 @@ class AccountSelection:
 
 
 class LoadBalancer:
-    def __init__(
-        self,
-        accounts_repo: AccountsRepository,
-        usage_repo: UsageRepository,
-        sticky_repo: StickySessionsRepository | None = None,
-    ) -> None:
-        self._accounts_repo = accounts_repo
-        self._usage_repo = usage_repo
-        self._usage_updater = UsageUpdater(usage_repo, accounts_repo)
-        self._sticky_repo = sticky_repo
+    def __init__(self, repo_factory: ProxyRepoFactory) -> None:
+        self._repo_factory = repo_factory
         self._runtime: dict[str, RuntimeState] = {}
 
     async def select_account(
@@ -56,43 +48,53 @@ class LoadBalancer:
         reallocate_sticky: bool = False,
         prefer_earlier_reset_accounts: bool = False,
     ) -> AccountSelection:
-        accounts = await self._accounts_repo.list_accounts()
-        latest_primary = await self._usage_repo.latest_by_account()
-        await self._usage_updater.refresh_accounts(accounts, latest_primary)
-        latest_primary = await self._usage_repo.latest_by_account()
-        latest_secondary = await self._usage_repo.latest_by_account(window="secondary")
-
-        states, account_map = _build_states(
-            accounts=accounts,
-            latest_primary=latest_primary,
-            latest_secondary=latest_secondary,
-            runtime=self._runtime,
-        )
+        selected_snapshot: Account | None = None
+        error_message: str | None = None
+        async with self._repo_factory() as repos:
+            accounts = await repos.accounts.list_accounts()
+            latest_primary = await repos.usage.latest_by_account()
+            updater = UsageUpdater(repos.usage, repos.accounts)
+            await updater.refresh_accounts(accounts, latest_primary)
+            latest_primary = await repos.usage.latest_by_account()
+            latest_secondary = await repos.usage.latest_by_account(window="secondary")
+
+            states, account_map = _build_states(
+                accounts=accounts,
+                latest_primary=latest_primary,
+                latest_secondary=latest_secondary,
+                runtime=self._runtime,
+            )
 
-        result = await self._select_with_stickiness(
-            states=states,
-            account_map=account_map,
-            sticky_key=sticky_key,
-            reallocate_sticky=reallocate_sticky,
-            prefer_earlier_reset_accounts=prefer_earlier_reset_accounts,
-        )
-        for state in states:
-            account = account_map.get(state.account_id)
-            if account:
-                await self._sync_state(account, state)
-
-        if result.account is None:
-            return AccountSelection(account=None, error_message=result.error_message)
-
-        selected = account_map.get(result.account.account_id)
-        if selected:
-            selected.status = result.account.status
-            selected.deactivation_reason = result.account.deactivation_reason
-            runtime = self._runtime.setdefault(selected.id, RuntimeState())
-            runtime.last_selected_at = time.time()
-        if selected is None:
-            return AccountSelection(account=None, error_message=result.error_message)
-        return AccountSelection(account=selected, error_message=None)
+            result = await self._select_with_stickiness(
+                states=states,
+                account_map=account_map,
+                sticky_key=sticky_key,
+                reallocate_sticky=reallocate_sticky,
+                prefer_earlier_reset_accounts=prefer_earlier_reset_accounts,
+                sticky_repo=repos.sticky_sessions,
+            )
+            for state in states:
+                account = account_map.get(state.account_id)
+                if account:
+                    await self._sync_state(repos.accounts, account, state)
+
+            if result.account is None:
+                error_message = result.error_message
+            else:
+                selected = account_map.get(result.account.account_id)
+                if selected is None:
+                    error_message = result.error_message
+                else:
+                    selected.status = result.account.status
+                    selected.deactivation_reason = result.account.deactivation_reason
+                    selected_snapshot = _clone_account(selected)
+
+        if selected_snapshot is None:
+            return AccountSelection(account=None, error_message=error_message)
+
+        runtime = self._runtime.setdefault(selected_snapshot.id, RuntimeState())
+        runtime.last_selected_at = time.time()
+        return AccountSelection(account=selected_snapshot, error_message=None)
 
     async def _select_with_stickiness(
         self,
@@ -102,21 +104,22 @@ class LoadBalancer:
         sticky_key: str | None,
         reallocate_sticky: bool,
         prefer_earlier_reset_accounts: bool,
+        sticky_repo: StickySessionsRepository | None,
     ) -> SelectionResult:
-        if not sticky_key or not self._sticky_repo:
+        if not sticky_key or not sticky_repo:
             return select_account(states, prefer_earlier_reset=prefer_earlier_reset_accounts)
 
         if reallocate_sticky:
             chosen = select_account(states, prefer_earlier_reset=prefer_earlier_reset_accounts)
             if chosen.account is not None and chosen.account.account_id in account_map:
-                await self._sticky_repo.upsert(sticky_key, chosen.account.account_id)
+                await sticky_repo.upsert(sticky_key, chosen.account.account_id)
             return chosen
 
-        existing = await self._sticky_repo.get_account_id(sticky_key)
+        existing = await sticky_repo.get_account_id(sticky_key)
         if existing:
             pinned = next((state for state in states if state.account_id == existing), None)
             if pinned is None:
-                await self._sticky_repo.delete(sticky_key)
+                await sticky_repo.delete(sticky_key)
             else:
                 pinned_result = select_account([pinned], prefer_earlier_reset=prefer_earlier_reset_accounts)
                 if pinned_result.account is not None:
@@ -124,29 +127,33 @@ class LoadBalancer:
 
         chosen = select_account(states, prefer_earlier_reset=prefer_earlier_reset_accounts)
         if chosen.account is not None and chosen.account.account_id in account_map:
-            await self._sticky_repo.upsert(sticky_key, chosen.account.account_id)
+            await sticky_repo.upsert(sticky_key, chosen.account.account_id)
         return chosen
 
     async def mark_rate_limit(self, account: Account, error: UpstreamError) -> None:
         state = self._state_for(account)
         handle_rate_limit(state, error)
-        await self._sync_state(account, state)
+        async with self._repo_factory() as repos:
+            await self._sync_state(repos.accounts, account, state)
 
     async def mark_quota_exceeded(self, account: Account, error: UpstreamError) -> None:
         state = self._state_for(account)
         handle_quota_exceeded(state, error)
-        await self._sync_state(account, state)
+        async with self._repo_factory() as repos:
+            await self._sync_state(repos.accounts, account, state)
 
     async def mark_permanent_failure(self, account: Account, error_code: str) -> None:
         state = self._state_for(account)
         handle_permanent_failure(state, error_code)
-        await self._sync_state(account, state)
+        async with self._repo_factory() as repos:
+            await self._sync_state(repos.accounts, account, state)
 
     async def record_error(self, account: Account) -> None:
         state = self._state_for(account)
         state.error_count += 1
         state.last_error_at = time.time()
-        await self._sync_state(account, state)
+        async with self._repo_factory() as repos:
+            await self._sync_state(repos.accounts, account, state)
 
     def _state_for(self, account: Account) -> AccountState:
         runtime = self._runtime.setdefault(account.id, RuntimeState())
@@ -164,7 +171,12 @@ class LoadBalancer:
             deactivation_reason=account.deactivation_reason,
         )
 
-    async def _sync_state(self, account: Account, state: AccountState) -> None:
+    async def _sync_state(
+        self,
+        accounts_repo: AccountsRepository,
+        account: Account,
+        state: AccountState,
+    ) -> None:
         runtime = self._runtime.setdefault(account.id, RuntimeState())
         runtime.reset_at = state.reset_at
         runtime.cooldown_until = state.cooldown_until
@@ -177,7 +189,7 @@ class LoadBalancer:
         reset_changed = account.reset_at != reset_at_int
 
         if status_changed or reason_changed or reset_changed:
-            await self._accounts_repo.update_status(
+            await accounts_repo.update_status(
                 account.id,
                 state.status,
                 state.deactivation_reason,
@@ -251,3 +263,8 @@ def _state_from_account(
         error_count=runtime.error_count,
         deactivation_reason=account.deactivation_reason,
     )
+
+
+def _clone_account(account: Account) -> Account:
+    data = {column.name: getattr(account, column.name) for column in Account.__table__.columns}
+    return Account(**data)

app/modules/proxy/repo_bundle.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import AsyncContextManager
+
+from app.modules.accounts.repository import AccountsRepository
+from app.modules.proxy.sticky_repository import StickySessionsRepository
+from app.modules.request_logs.repository import RequestLogsRepository
+from app.modules.settings.repository import SettingsRepository
+from app.modules.usage.repository import UsageRepository
+
+
+@dataclass(slots=True)
+class ProxyRepositories:
+    accounts: AccountsRepository
+    usage: UsageRepository
+    request_logs: RequestLogsRepository
+    sticky_sessions: StickySessionsRepository
+    settings: SettingsRepository
+
+
+ProxyRepoFactory = Callable[[], AsyncContextManager[ProxyRepositories]]

app/modules/proxy/service.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 import logging
 import time
 from collections.abc import Sequence
@@ -28,7 +29,6 @@ from app.core.utils.sse import format_sse_event
 from app.core.utils.time import utcnow
 from app.db.models import Account, UsageHistory
 from app.modules.accounts.auth_manager import AuthManager
-from app.modules.accounts.repository import AccountsRepository
 from app.modules.proxy.helpers import (
     _apply_error_metadata,
     _credits_headers,
@@ -45,33 +45,18 @@ from app.modules.proxy.helpers import (
     _window_snapshot,
 )
 from app.modules.proxy.load_balancer import LoadBalancer
-from app.modules.proxy.sticky_repository import StickySessionsRepository
+from app.modules.proxy.repo_bundle import ProxyRepoFactory, ProxyRepositories
 from app.modules.proxy.types import RateLimitStatusPayloadData
-from app.modules.request_logs.repository import RequestLogsRepository
-from app.modules.settings.repository import SettingsRepository
-from app.modules.usage.repository import UsageRepository
 from app.modules.usage.updater import UsageUpdater
 
 logger = logging.getLogger(__name__)
 
 
 class ProxyService:
-    def __init__(
-        self,
-        accounts_repo: AccountsRepository,
-        usage_repo: UsageRepository,
-        logs_repo: RequestLogsRepository,
-        sticky_repo: StickySessionsRepository,
-        settings_repo: SettingsRepository,
-    ) -> None:
-        self._accounts_repo = accounts_repo
-        self._usage_repo = usage_repo
-        self._logs_repo = logs_repo
-        self._settings_repo = settings_repo
+    def __init__(self, repo_factory: ProxyRepoFactory) -> None:
+        self._repo_factory = repo_factory
         self._encryptor = TokenEncryptor()
-        self._auth_manager = AuthManager(accounts_repo)
-        self._load_balancer = LoadBalancer(accounts_repo, usage_repo, sticky_repo)
-        self._usage_updater = UsageUpdater(usage_repo, accounts_repo)
+        self._load_balancer = LoadBalancer(repo_factory)
 
     def stream_responses(
         self,
@@ -80,6 +65,7 @@ class ProxyService:
         *,
         propagate_http_errors: bool = False,
     ) -> AsyncIterator[str]:
+        _maybe_log_proxy_request_payload("stream", payload, headers)
         _maybe_log_proxy_request_shape("stream", payload, headers)
         filtered = filter_inbound_headers(headers)
         return self._stream_with_retry(
@@ -93,11 +79,14 @@ class ProxyService:
         payload: ResponsesCompactRequest,
         headers: Mapping[str, str],
     ) -> OpenAIResponsePayload:
+        _maybe_log_proxy_request_payload("compact", payload, headers)
         _maybe_log_proxy_request_shape("compact", payload, headers)
         filtered = filter_inbound_headers(headers)
-        settings = await self._settings_repo.get_or_create()
-        prefer_earlier_reset = settings.prefer_earlier_reset_accounts
-        sticky_key = _sticky_key_from_compact_payload(payload) if settings.sticky_threads_enabled else None
+        async with self._repo_factory() as repos:
+            settings = await repos.settings.get_or_create()
+            prefer_earlier_reset = settings.prefer_earlier_reset_accounts
+            sticky_threads_enabled = settings.sticky_threads_enabled
+        sticky_key = _sticky_key_from_compact_payload(payload) if sticky_threads_enabled else None
         selection = await self._load_balancer.select_account(
             sticky_key=sticky_key,
             reallocate_sticky=sticky_key is not None,
@@ -136,69 +125,71 @@ class ProxyService:
 
     async def rate_limit_headers(self) -> dict[str, str]:
         now = utcnow()
-        accounts = await self._accounts_repo.list_accounts()
-        account_map = {account.id: account for account in accounts}
-
         headers: dict[str, str] = {}
-        primary_minutes = await self._usage_repo.latest_window_minutes("primary")
-        if primary_minutes is None:
-            primary_minutes = usage_core.default_window_minutes("primary")
-        if primary_minutes:
-            primary_rows = await self._usage_repo.aggregate_since(
-                now - timedelta(minutes=primary_minutes),
-                window="primary",
-            )
-            if primary_rows:
-                summary = usage_core.summarize_usage_window(
-                    [row.to_window_row() for row in primary_rows],
-                    account_map,
-                    "primary",
+        async with self._repo_factory() as repos:
+            accounts = await repos.accounts.list_accounts()
+            account_map = {account.id: account for account in accounts}
+
+            primary_minutes = await repos.usage.latest_window_minutes("primary")
+            if primary_minutes is None:
+                primary_minutes = usage_core.default_window_minutes("primary")
+            if primary_minutes:
+                primary_rows = await repos.usage.aggregate_since(
+                    now - timedelta(minutes=primary_minutes),
+                    window="primary",
                 )
-                headers.update(_rate_limit_headers("primary", summary))
-
-        secondary_minutes = await self._usage_repo.latest_window_minutes("secondary")
-        if secondary_minutes is None:
-            secondary_minutes = usage_core.default_window_minutes("secondary")
-        if secondary_minutes:
-            secondary_rows = await self._usage_repo.aggregate_since(
-                now - timedelta(minutes=secondary_minutes),
-                window="secondary",
-            )
-            if secondary_rows:
-                summary = usage_core.summarize_usage_window(
-                    [row.to_window_row() for row in secondary_rows],
-                    account_map,
-                    "secondary",
+                if primary_rows:
+                    summary = usage_core.summarize_usage_window(
+                        [row.to_window_row() for row in primary_rows],
+                        account_map,
+                        "primary",
+                    )
+                    headers.update(_rate_limit_headers("primary", summary))
+
+            secondary_minutes = await repos.usage.latest_window_minutes("secondary")
+            if secondary_minutes is None:
+                secondary_minutes = usage_core.default_window_minutes("secondary")
+            if secondary_minutes:
+                secondary_rows = await repos.usage.aggregate_since(
+                    now - timedelta(minutes=secondary_minutes),
+                    window="secondary",
                 )
-                headers.update(_rate_limit_headers("secondary", summary))
+                if secondary_rows:
+                    summary = usage_core.summarize_usage_window(
+                        [row.to_window_row() for row in secondary_rows],
+                        account_map,
+                        "secondary",
+                    )
+                    headers.update(_rate_limit_headers("secondary", summary))
 
-        latest_usage = await self._usage_repo.latest_by_account()
-        headers.update(_credits_headers(latest_usage.values()))
+            latest_usage = await repos.usage.latest_by_account()
+            headers.update(_credits_headers(latest_usage.values()))
         return headers
 
     async def get_rate_limit_payload(self) -> RateLimitStatusPayloadData:
-        accounts = await self._accounts_repo.list_accounts()
-        await self._refresh_usage(accounts)
-        selected_accounts = _select_accounts_for_limits(accounts)
-        if not selected_accounts:
-            return RateLimitStatusPayloadData(plan_type="guest")
-
-        account_map = {account.id: account for account in selected_accounts}
-        primary_rows = await self._latest_usage_rows(account_map, "primary")
-        secondary_rows = await self._latest_usage_rows(account_map, "secondary")
-
-        primary_summary = _summarize_window(primary_rows, account_map, "primary")
-        secondary_summary = _summarize_window(secondary_rows, account_map, "secondary")
-
-        now_epoch = int(time.time())
-        primary_window = _window_snapshot(primary_summary, primary_rows, "primary", now_epoch)
-        secondary_window = _window_snapshot(secondary_summary, secondary_rows, "secondary", now_epoch)
-
-        return RateLimitStatusPayloadData(
-            plan_type=_plan_type_for_accounts(selected_accounts),
-            rate_limit=_rate_limit_details(primary_window, secondary_window),
-            credits=_credits_snapshot(await self._latest_usage_entries(account_map)),
-        )
+        async with self._repo_factory() as repos:
+            accounts = await repos.accounts.list_accounts()
+            await self._refresh_usage(repos, accounts)
+            selected_accounts = _select_accounts_for_limits(accounts)
+            if not selected_accounts:
+                return RateLimitStatusPayloadData(plan_type="guest")
+
+            account_map = {account.id: account for account in selected_accounts}
+            primary_rows = await self._latest_usage_rows(repos, account_map, "primary")
+            secondary_rows = await self._latest_usage_rows(repos, account_map, "secondary")
+
+            primary_summary = _summarize_window(primary_rows, account_map, "primary")
+            secondary_summary = _summarize_window(secondary_rows, account_map, "secondary")
+
+            now_epoch = int(time.time())
+            primary_window = _window_snapshot(primary_summary, primary_rows, "primary", now_epoch)
+            secondary_window = _window_snapshot(secondary_summary, secondary_rows, "secondary", now_epoch)
+
+            return RateLimitStatusPayloadData(
+                plan_type=_plan_type_for_accounts(selected_accounts),
+                rate_limit=_rate_limit_details(primary_window, secondary_window),
+                credits=_credits_snapshot(await self._latest_usage_entries(repos, account_map)),
+            )
 
     async def _stream_with_retry(
         self,
@@ -208,9 +199,11 @@ class ProxyService:
         propagate_http_errors: bool,
     ) -> AsyncIterator[str]:
         request_id = ensure_request_id()
-        settings = await self._settings_repo.get_or_create()
-        prefer_earlier_reset = settings.prefer_earlier_reset_accounts
-        sticky_key = _sticky_key_from_payload(payload) if settings.sticky_threads_enabled else None
+        async with self._repo_factory() as repos:
+            settings = await repos.settings.get_or_create()
+            prefer_earlier_reset = settings.prefer_earlier_reset_accounts
+            sticky_threads_enabled = settings.sticky_threads_enabled
+        sticky_key = _sticky_key_from_payload(payload) if sticky_threads_enabled else None
         max_attempts = 3
         for attempt in range(max_attempts):
             selection = await self._load_balancer.select_account(
@@ -398,20 +391,21 @@ class ProxyService:
             )
             with anyio.CancelScope(shield=True):
                 try:
-                    await self._logs_repo.add_log(
-                        account_id=account_id_value,
-                        request_id=request_id,
-                        model=model,
-                        input_tokens=input_tokens,
-                        output_tokens=output_tokens,
-                        cached_input_tokens=cached_input_tokens,
-                        reasoning_tokens=reasoning_tokens,
-                        reasoning_effort=reasoning_effort,
-                        latency_ms=latency_ms,
-                        status=status,
-                        error_code=error_code,
-                        error_message=error_message,
-                    )
+                    async with self._repo_factory() as repos:
+                        await repos.request_logs.add_log(
+                            account_id=account_id_value,
+                            request_id=request_id,
+                            model=model,
+                            input_tokens=input_tokens,
+                            output_tokens=output_tokens,
+                            cached_input_tokens=cached_input_tokens,
+                            reasoning_tokens=reasoning_tokens,
+                            reasoning_effort=reasoning_effort,
+                            latency_ms=latency_ms,
+                            status=status,
+                            error_code=error_code,
+                            error_message=error_message,
+                        )
                 except Exception:
                     logger.warning(
                         "Failed to persist request log account_id=%s request_id=%s",
@@ -420,18 +414,20 @@ class ProxyService:
                         exc_info=True,
                     )
 
-    async def _refresh_usage(self, accounts: list[Account]) -> None:
-        latest_usage = await self._usage_repo.latest_by_account(window="primary")
-        await self._usage_updater.refresh_accounts(accounts, latest_usage)
+    async def _refresh_usage(self, repos: ProxyRepositories, accounts: list[Account]) -> None:
+        latest_usage = await repos.usage.latest_by_account(window="primary")
+        updater = UsageUpdater(repos.usage, repos.accounts)
+        await updater.refresh_accounts(accounts, latest_usage)
 
     async def _latest_usage_rows(
         self,
+        repos: ProxyRepositories,
         account_map: dict[str, Account],
         window: str,
     ) -> list[UsageWindowRow]:
         if not account_map:
             return []
-        latest = await self._usage_repo.latest_by_account(window=window)
+        latest = await repos.usage.latest_by_account(window=window)
         return [
             UsageWindowRow(
                 account_id=entry.account_id,
@@ -445,15 +441,18 @@ class ProxyService:
 
     async def _latest_usage_entries(
         self,
+        repos: ProxyRepositories,
         account_map: dict[str, Account],
     ) -> list[UsageHistory]:
         if not account_map:
             return []
-        latest = await self._usage_repo.latest_by_account()
+        latest = await repos.usage.latest_by_account()
         return [entry for entry in latest.values() if entry.account_id in account_map]
 
     async def _ensure_fresh(self, account: Account, *, force: bool = False) -> Account:
-        return await self._auth_manager.ensure_fresh(account, force=force)
+        async with self._repo_factory() as repos:
+            auth_manager = AuthManager(repos.accounts)
+            return await auth_manager.ensure_fresh(account, force=force)
 
     async def _handle_proxy_error(self, account: Account, exc: ProxyResponseError) -> None:
         error = _parse_openai_error(exc.payload)
@@ -526,6 +525,32 @@ def _maybe_log_proxy_request_shape(
     )
 
 
+def _maybe_log_proxy_request_payload(
+    kind: str,
+    payload: ResponsesRequest | ResponsesCompactRequest,
+    headers: Mapping[str, str],
+) -> None:
+    settings = get_settings()
+    if not settings.log_proxy_request_payload:
+        return
+
+    request_id = get_request_id()
+    payload_dict = payload.model_dump(mode="json", exclude_none=True)
+    extra = payload.model_extra or {}
+    if extra:
+        payload_dict = {**payload_dict, "_extra": extra}
+    header_keys = _interesting_header_keys(headers)
+    payload_json = json.dumps(payload_dict, ensure_ascii=True, separators=(",", ":"))
+
+    logger.warning(
+        "proxy_request_payload request_id=%s kind=%s payload=%s headers=%s",
+        request_id,
+        kind,
+        payload_json,
+        header_keys,
+    )
+
+
 def _hash_identifier(value: str) -> str:
     digest = sha256(value.encode("utf-8")).hexdigest()
     return f"sha256:{digest[:12]}"