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

app/core/openai/v1_requests.py

@@ -1,9 +1,8 @@
 from __future__ import annotations
 
-from typing import cast
-
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
+from app.core.openai.message_coercion import coerce_messages
 from app.core.openai.requests import (
     ResponsesCompactRequest,
     ResponsesReasoning,
@@ -21,7 +20,7 @@ class V1ResponsesRequest(BaseModel):
     input: list[JsonValue] | None = None
     instructions: str | None = None
     tools: list[JsonValue] = Field(default_factory=list)
-    tool_choice: str | None = None
+    tool_choice: str | dict[str, JsonValue] | None = None
     parallel_tool_calls: bool | None = None
     reasoning: ResponsesReasoning | None = None
     store: bool | None = None
@@ -54,7 +53,7 @@ class V1ResponsesRequest(BaseModel):
         input_items: list[JsonValue] = input_value if isinstance(input_value, list) else []
 
         if messages is not None:
-            instruction_text, input_items = _coerce_messages(instruction_text, messages)
+            instruction_text, input_items = coerce_messages(instruction_text, messages)
 
         data["instructions"] = instruction_text
         data["input"] = input_items
@@ -86,63 +85,8 @@ class V1ResponsesCompactRequest(BaseModel):
         input_items: list[JsonValue] = input_value if isinstance(input_value, list) else []
 
         if messages is not None:
-            instruction_text, input_items = _coerce_messages(instruction_text, messages)
+            instruction_text, input_items = coerce_messages(instruction_text, messages)
 
         data["instructions"] = instruction_text
         data["input"] = input_items
         return ResponsesCompactRequest.model_validate(data)
-
-
-def _coerce_messages(existing_instructions: str, messages: list[JsonValue]) -> tuple[str, list[JsonValue]]:
-    instruction_parts: list[str] = []
-    input_messages: list[JsonValue] = []
-    for message in messages:
-        if not isinstance(message, dict):
-            raise ValueError("Each message must be an object.")
-        message_dict = cast(dict[str, JsonValue], message)
-        role_value = message_dict.get("role")
-        role = role_value if isinstance(role_value, str) else None
-        if role in ("system", "developer"):
-            content_text = _content_to_text(message_dict.get("content"))
-            if content_text:
-                instruction_parts.append(content_text)
-            continue
-        input_messages.append(cast(JsonValue, message_dict))
-    merged = _merge_instructions(existing_instructions, instruction_parts)
-    return merged, input_messages
-
-
-def _merge_instructions(existing: str, extra_parts: list[str]) -> str:
-    if not extra_parts:
-        return existing
-    extra = "\n".join([part for part in extra_parts if part])
-    if not extra:
-        return existing
-    if existing:
-        return f"{existing}\n{extra}"
-    return extra
-
-
-def _content_to_text(content: object) -> str | None:
-    if content is None:
-        return None
-    if isinstance(content, str):
-        return content
-    if isinstance(content, list):
-        parts: list[str] = []
-        for part in content:
-            if isinstance(part, str):
-                parts.append(part)
-            elif isinstance(part, dict):
-                part_dict = cast(dict[str, JsonValue], part)
-                text = part_dict.get("text")
-                if isinstance(text, str):
-                    parts.append(text)
-        return "\n".join([part for part in parts if part])
-    if isinstance(content, dict):
-        content_dict = cast(dict[str, JsonValue], content)
-        text = content_dict.get("text")
-        if isinstance(text, str):
-            return text
-        return None
-    return None

app/db/session.py

@@ -13,7 +13,7 @@ from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_asyn
 from app.core.config.settings import get_settings
 from app.db.migrations import run_migrations
 
-DATABASE_URL = get_settings().database_url
+_settings = get_settings()
 
 logger = logging.getLogger(__name__)
 
@@ -43,15 +43,32 @@ def _configure_sqlite_engine(engine: Engine, *, enable_wal: bool) -> None:
             cursor.close()
 
 
-if _is_sqlite_url(DATABASE_URL):
+if _is_sqlite_url(_settings.database_url):
+    is_sqlite_memory = _is_sqlite_memory_url(_settings.database_url)
+    if is_sqlite_memory:
+        engine = create_async_engine(
+            _settings.database_url,
+            echo=False,
+            connect_args={"timeout": _SQLITE_BUSY_TIMEOUT_SECONDS},
+        )
+    else:
+        engine = create_async_engine(
+            _settings.database_url,
+            echo=False,
+            pool_size=_settings.database_pool_size,
+            max_overflow=_settings.database_max_overflow,
+            pool_timeout=_settings.database_pool_timeout_seconds,
+            connect_args={"timeout": _SQLITE_BUSY_TIMEOUT_SECONDS},
+        )
+    _configure_sqlite_engine(engine.sync_engine, enable_wal=not is_sqlite_memory)
+else:
     engine = create_async_engine(
-        DATABASE_URL,
+        _settings.database_url,
         echo=False,
-        connect_args={"timeout": _SQLITE_BUSY_TIMEOUT_SECONDS},
+        pool_size=_settings.database_pool_size,
+        max_overflow=_settings.database_max_overflow,
+        pool_timeout=_settings.database_pool_timeout_seconds,
     )
-    _configure_sqlite_engine(engine.sync_engine, enable_wal=not _is_sqlite_memory_url(DATABASE_URL))
-else:
-    engine = create_async_engine(DATABASE_URL, echo=False)
 
 SessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)
 
@@ -116,7 +133,7 @@ async def get_session() -> AsyncIterator[AsyncSession]:
 async def init_db() -> None:
     from app.db.models import Base
 
-    _ensure_sqlite_dir(DATABASE_URL)
+    _ensure_sqlite_dir(_settings.database_url)
 
     async with engine.begin() as conn:
         await conn.run_sync(Base.metadata.create_all)

app/dependencies.py

@@ -11,6 +11,7 @@ from app.db.session import SessionLocal, _safe_close, _safe_rollback, get_sessio
 from app.modules.accounts.repository import AccountsRepository
 from app.modules.accounts.service import AccountsService
 from app.modules.oauth.service import OauthService
+from app.modules.proxy.repo_bundle import ProxyRepositories
 from app.modules.proxy.service import ProxyService
 from app.modules.proxy.sticky_repository import StickySessionsRepository
 from app.modules.request_logs.repository import RequestLogsRepository
@@ -79,7 +80,12 @@ def get_usage_context(
     usage_repository = UsageRepository(session)
     request_logs_repository = RequestLogsRepository(session)
     accounts_repository = AccountsRepository(session)
-    service = UsageService(usage_repository, request_logs_repository, accounts_repository)
+    service = UsageService(
+        usage_repository,
+        request_logs_repository,
+        accounts_repository,
+        refresh_repo_factory=_usage_refresh_context,
+    )
     return UsageContext(
         session=session,
         usage_repository=usage_repository,
@@ -101,6 +107,40 @@ async def _accounts_repo_context() -> AsyncIterator[AccountsRepository]:
         await _safe_close(session)
 
 
+@asynccontextmanager
+async def _usage_refresh_context() -> AsyncIterator[tuple[UsageRepository, AccountsRepository]]:
+    session = SessionLocal()
+    try:
+        yield UsageRepository(session), AccountsRepository(session)
+    except BaseException:
+        await _safe_rollback(session)
+        raise
+    finally:
+        if session.in_transaction():
+            await _safe_rollback(session)
+        await _safe_close(session)
+
+
+@asynccontextmanager
+async def _proxy_repo_context() -> AsyncIterator[ProxyRepositories]:
+    session = SessionLocal()
+    try:
+        yield ProxyRepositories(
+            accounts=AccountsRepository(session),
+            usage=UsageRepository(session),
+            request_logs=RequestLogsRepository(session),
+            sticky_sessions=StickySessionsRepository(session),
+            settings=SettingsRepository(session),
+        )
+    except BaseException:
+        await _safe_rollback(session)
+        raise
+    finally:
+        if session.in_transaction():
+            await _safe_rollback(session)
+        await _safe_close(session)
+
+
 def get_oauth_context(
     session: AsyncSession = Depends(get_session),
 ) -> OauthContext:
@@ -108,21 +148,8 @@ def get_oauth_context(
     return OauthContext(service=OauthService(accounts_repository, repo_factory=_accounts_repo_context))
 
 
-def get_proxy_context(
-    session: AsyncSession = Depends(get_session),
-) -> ProxyContext:
-    accounts_repository = AccountsRepository(session)
-    usage_repository = UsageRepository(session)
-    request_logs_repository = RequestLogsRepository(session)
-    sticky_repository = StickySessionsRepository(session)
-    settings_repository = SettingsRepository(session)
-    service = ProxyService(
-        accounts_repository,
-        usage_repository,
-        request_logs_repository,
-        sticky_repository,
-        settings_repository,
-    )
+def get_proxy_context() -> ProxyContext:
+    service = ProxyService(repo_factory=_proxy_repo_context)
     return ProxyContext(service=service)
 
 

app/main.py

@@ -1,23 +1,19 @@
 from __future__ import annotations
 
-import logging
 from contextlib import asynccontextmanager
 from pathlib import Path
-from uuid import uuid4
 
-from fastapi import FastAPI, Request
-from fastapi.exception_handlers import (
-    http_exception_handler,
-    request_validation_exception_handler,
-)
-from fastapi.exceptions import RequestValidationError
-from fastapi.responses import FileResponse, JSONResponse, RedirectResponse, Response
+from fastapi import FastAPI
+from fastapi.responses import FileResponse, RedirectResponse
 from fastapi.staticfiles import StaticFiles
-from starlette.exceptions import HTTPException as StarletteHTTPException
 
 from app.core.clients.http import close_http_client, init_http_client
-from app.core.errors import dashboard_error
-from app.core.utils.request_id import get_request_id, reset_request_id, set_request_id
+from app.core.handlers import add_exception_handlers
+from app.core.middleware import (
+    add_api_unhandled_error_middleware,
+    add_request_decompression_middleware,
+    add_request_id_middleware,
+)
 from app.db.session import close_db, init_db
 from app.modules.accounts import api as accounts_api
 from app.modules.health import api as health_api
@@ -27,8 +23,6 @@ from app.modules.request_logs import api as request_logs_api
 from app.modules.settings import api as settings_api
 from app.modules.usage import api as usage_api
 
-logger = logging.getLogger(__name__)
-
 
 @asynccontextmanager
 async def lifespan(_: FastAPI):
@@ -47,59 +41,10 @@ async def lifespan(_: FastAPI):
 def create_app() -> FastAPI:
     app = FastAPI(title="codex-lb", version="0.1.0", lifespan=lifespan)
 
-    @app.middleware("http")
-    async def request_id_middleware(request: Request, call_next) -> JSONResponse:
-        inbound_request_id = request.headers.get("x-request-id") or request.headers.get("request-id")
-        request_id = inbound_request_id or str(uuid4())
-        token = set_request_id(request_id)
-        try:
-            response = await call_next(request)
-        except Exception:
-            reset_request_id(token)
-            raise
-        response.headers.setdefault("x-request-id", request_id)
-        return response
-
-    @app.middleware("http")
-    async def api_unhandled_error_middleware(request: Request, call_next) -> Response:
-        try:
-            return await call_next(request)
-        except Exception:
-            if request.url.path.startswith("/api/"):
-                logger.exception(
-                    "Unhandled API error request_id=%s",
-                    get_request_id(),
-                )
-                return JSONResponse(
-                    status_code=500,
-                    content=dashboard_error("internal_error", "Unexpected error"),
-                )
-            raise
-
-    @app.exception_handler(RequestValidationError)
-    async def _validation_error_handler(
-        request: Request,
-        exc: RequestValidationError,
-    ) -> Response:
-        if request.url.path.startswith("/api/"):
-            return JSONResponse(
-                status_code=422,
-                content=dashboard_error("validation_error", "Invalid request payload"),
-            )
-        return await request_validation_exception_handler(request, exc)
-
-    @app.exception_handler(StarletteHTTPException)
-    async def _http_error_handler(
-        request: Request,
-        exc: StarletteHTTPException,
-    ) -> Response:
-        if request.url.path.startswith("/api/"):
-            detail = exc.detail if isinstance(exc.detail, str) else "Request failed"
-            return JSONResponse(
-                status_code=exc.status_code,
-                content=dashboard_error(f"http_{exc.status_code}", detail),
-            )
-        return await http_exception_handler(request, exc)
+    add_request_decompression_middleware(app)
+    add_request_id_middleware(app)
+    add_api_unhandled_error_middleware(app)
+    add_exception_handlers(app)
 
     app.include_router(proxy_api.router)
     app.include_router(proxy_api.v1_router)

app/modules/accounts/repository.py

@@ -19,19 +19,19 @@ class AccountsRepository:
     async def upsert(self, account: Account) -> Account:
         existing = await self._session.get(Account, account.id)
         if existing:
-            existing.chatgpt_account_id = account.chatgpt_account_id
-            existing.email = account.email
-            existing.plan_type = account.plan_type
-            existing.access_token_encrypted = account.access_token_encrypted
-            existing.refresh_token_encrypted = account.refresh_token_encrypted
-            existing.id_token_encrypted = account.id_token_encrypted
-            existing.last_refresh = account.last_refresh
-            existing.status = account.status
-            existing.deactivation_reason = account.deactivation_reason
+            _apply_account_updates(existing, account)
             await self._session.commit()
             await self._session.refresh(existing)
             return existing
 
+        result = await self._session.execute(select(Account).where(Account.email == account.email))
+        existing_by_email = result.scalar_one_or_none()
+        if existing_by_email:
+            _apply_account_updates(existing_by_email, account)
+            await self._session.commit()
+            await self._session.refresh(existing_by_email)
+            return existing_by_email
+
         self._session.add(account)
         await self._session.commit()
         await self._session.refresh(account)
@@ -89,3 +89,15 @@ class AccountsRepository:
         )
         await self._session.commit()
         return result.scalar_one_or_none() is not None
+
+
+def _apply_account_updates(target: Account, source: Account) -> None:
+    target.chatgpt_account_id = source.chatgpt_account_id
+    target.email = source.email
+    target.plan_type = source.plan_type
+    target.access_token_encrypted = source.access_token_encrypted
+    target.refresh_token_encrypted = source.refresh_token_encrypted
+    target.id_token_encrypted = source.id_token_encrypted
+    target.last_refresh = source.last_refresh
+    target.status = source.status
+    target.deactivation_reason = source.deactivation_reason

app/modules/proxy/api.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import time
 from collections.abc import AsyncIterator
 
 from fastapi import APIRouter, Body, Depends, Request, Response
@@ -7,6 +8,9 @@ from fastapi.responses import JSONResponse, StreamingResponse
 
 from app.core.clients.proxy import ProxyResponseError
 from app.core.errors import openai_error
+from app.core.openai.chat_requests import ChatCompletionsRequest
+from app.core.openai.chat_responses import collect_chat_completion, stream_chat_chunks
+from app.core.openai.models_catalog import MODEL_CATALOG
 from app.core.openai.requests import ResponsesCompactRequest, ResponsesRequest
 from app.core.openai.v1_requests import V1ResponsesCompactRequest, V1ResponsesRequest
 from app.dependencies import ProxyContext, get_proxy_context
@@ -35,6 +39,60 @@ async def v1_responses(
     return await _stream_responses(request, payload.to_responses_request(), context)
 
 
+@v1_router.get("/models")
+async def v1_models() -> JSONResponse:
+    created = int(time.time())
+    items = [
+        {
+            "id": model_id,
+            "object": "model",
+            "created": created,
+            "owned_by": "codex-lb",
+            "metadata": entry.model_dump(mode="json"),
+        }
+        for model_id, entry in MODEL_CATALOG.items()
+    ]
+    return JSONResponse({"object": "list", "data": items})
+
+
+@v1_router.post("/chat/completions")
+async def v1_chat_completions(
+    request: Request,
+    payload: ChatCompletionsRequest = Body(...),
+    context: ProxyContext = Depends(get_proxy_context),
+) -> Response:
+    rate_limit_headers = await context.service.rate_limit_headers()
+    responses_payload = payload.to_responses_request()
+    responses_payload.stream = True
+    stream = context.service.stream_responses(
+        responses_payload,
+        request.headers,
+        propagate_http_errors=True,
+    )
+    try:
+        first = await stream.__anext__()
+    except StopAsyncIteration:
+        first = None
+    except ProxyResponseError as exc:
+        return JSONResponse(status_code=exc.status_code, content=exc.payload, headers=rate_limit_headers)
+
+    stream_with_first = _prepend_first(first, stream)
+    if payload.stream:
+        return StreamingResponse(
+            stream_chat_chunks(stream_with_first, model=payload.model),
+            media_type="text/event-stream",
+            headers={"Cache-Control": "no-cache", **rate_limit_headers},
+        )
+
+    result = await collect_chat_completion(stream_with_first, model=payload.model)
+    status_code = 200
+    if isinstance(result, dict) and "error" in result:
+        error = result.get("error")
+        code = error.get("code") if isinstance(error, dict) else None
+        status_code = 503 if code == "no_accounts" else 502
+    return JSONResponse(content=result, status_code=status_code, headers=rate_limit_headers)
+
+
 async def _stream_responses(
     request: Request,
     payload: ResponsesRequest,

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)