codex-lb 0.1.2__py3-none-any.whl

app/modules/accounts/service.py

@@ -0,0 +1,211 @@
+from __future__ import annotations
+
+from datetime import datetime, timedelta, timezone
+
+from app.core import usage as usage_core
+from app.core.auth import (
+    DEFAULT_EMAIL,
+    DEFAULT_PLAN,
+    claims_from_auth,
+    extract_id_token_claims,
+    fallback_account_id,
+    parse_auth_json,
+)
+from app.core.crypto import TokenEncryptor
+from app.core.usage.logs import cost_from_log
+from app.core.utils.time import from_epoch_seconds, to_utc_naive, utcnow
+from app.db.models import Account, AccountStatus, UsageHistory
+from app.modules.accounts.repository import AccountsRepository
+from app.modules.accounts.schemas import (
+    AccountAuthStatus,
+    AccountImportResponse,
+    AccountSummary,
+    AccountTokenStatus,
+    AccountUsage,
+)
+from app.modules.proxy.usage_updater import UsageUpdater
+from app.modules.request_logs.repository import RequestLogsRepository
+from app.modules.usage.repository import UsageRepository
+
+
+class AccountsService:
+    def __init__(
+        self,
+        repo: AccountsRepository,
+        usage_repo: UsageRepository | None = None,
+        logs_repo: RequestLogsRepository | None = None,
+    ) -> None:
+        self._repo = repo
+        self._usage_repo = usage_repo
+        self._logs_repo = logs_repo
+        self._usage_updater = UsageUpdater(usage_repo, repo) if usage_repo else None
+        self._encryptor = TokenEncryptor()
+
+    async def list_accounts(self) -> list[AccountSummary]:
+        accounts = await self._repo.list_accounts()
+        if not accounts:
+            return []
+        await self._refresh_usage(accounts)
+        primary_usage = await self._usage_repo.latest_by_account(window="primary") if self._usage_repo else {}
+        secondary_usage = await self._usage_repo.latest_by_account(window="secondary") if self._usage_repo else {}
+        cost_by_account = await self._costs_last_24h()
+        return [
+            self._account_to_summary(
+                account,
+                primary_usage.get(account.id),
+                secondary_usage.get(account.id),
+                cost_by_account.get(account.id),
+            )
+            for account in accounts
+        ]
+
+    async def import_account(self, raw: bytes) -> AccountImportResponse:
+        auth = parse_auth_json(raw)
+        claims = claims_from_auth(auth)
+
+        email = claims.email or DEFAULT_EMAIL
+        plan_type = claims.plan_type or DEFAULT_PLAN
+        account_id = claims.account_id or fallback_account_id(email)
+        last_refresh = to_utc_naive(auth.last_refresh_at) if auth.last_refresh_at else utcnow()
+
+        account = Account(
+            id=account_id,
+            email=email,
+            plan_type=plan_type,
+            access_token_encrypted=self._encryptor.encrypt(auth.tokens.access_token),
+            refresh_token_encrypted=self._encryptor.encrypt(auth.tokens.refresh_token),
+            id_token_encrypted=self._encryptor.encrypt(auth.tokens.id_token),
+            last_refresh=last_refresh,
+            status=AccountStatus.ACTIVE,
+            deactivation_reason=None,
+        )
+
+        saved = await self._repo.upsert(account)
+        if self._usage_repo and self._usage_updater:
+            latest_usage = await self._usage_repo.latest_by_account(window="primary")
+            await self._usage_updater.refresh_accounts([saved], latest_usage)
+        return AccountImportResponse(
+            account_id=saved.id,
+            email=saved.email,
+            plan_type=saved.plan_type,
+            status=saved.status,
+        )
+
+    async def reactivate_account(self, account_id: str) -> bool:
+        return await self._repo.update_status(account_id, AccountStatus.ACTIVE, None)
+
+    async def pause_account(self, account_id: str) -> bool:
+        return await self._repo.update_status(account_id, AccountStatus.PAUSED, None)
+
+    async def delete_account(self, account_id: str) -> bool:
+        return await self._repo.delete(account_id)
+
+    def _account_to_summary(
+        self,
+        account: Account,
+        primary_usage: UsageHistory | None,
+        secondary_usage: UsageHistory | None,
+        cost_usd_24h: float | None,
+    ) -> AccountSummary:
+        auth_status = self._build_auth_status(account)
+        primary_used_percent = _normalize_used_percent(primary_usage) or 0.0
+        secondary_used_percent = _normalize_used_percent(secondary_usage) or 0.0
+        primary_remaining_percent = usage_core.remaining_percent_from_used(primary_used_percent) or 0.0
+        secondary_remaining_percent = usage_core.remaining_percent_from_used(secondary_used_percent) or 0.0
+        reset_at_primary = from_epoch_seconds(primary_usage.reset_at) if primary_usage is not None else None
+        reset_at_secondary = from_epoch_seconds(secondary_usage.reset_at) if secondary_usage is not None else None
+        capacity_primary = usage_core.capacity_for_plan(account.plan_type, "primary")
+        capacity_secondary = usage_core.capacity_for_plan(account.plan_type, "secondary")
+        remaining_credits_primary = usage_core.remaining_credits_from_percent(
+            primary_used_percent,
+            capacity_primary,
+        )
+        remaining_credits_secondary = usage_core.remaining_credits_from_percent(
+            secondary_used_percent,
+            capacity_secondary,
+        )
+        return AccountSummary(
+            account_id=account.id,
+            email=account.email,
+            display_name=account.email,
+            plan_type=account.plan_type,
+            status=account.status.value,
+            usage=AccountUsage(
+                primary_remaining_percent=primary_remaining_percent,
+                secondary_remaining_percent=secondary_remaining_percent,
+            ),
+            reset_at_primary=reset_at_primary,
+            reset_at_secondary=reset_at_secondary,
+            last_refresh_at=account.last_refresh,
+            capacity_credits_primary=capacity_primary,
+            remaining_credits_primary=remaining_credits_primary,
+            capacity_credits_secondary=capacity_secondary,
+            remaining_credits_secondary=remaining_credits_secondary,
+            cost_usd_24h=cost_usd_24h if cost_usd_24h is not None else 0.0,
+            deactivation_reason=account.deactivation_reason,
+            auth=auth_status,
+        )
+
+    def _build_auth_status(self, account: Account) -> AccountAuthStatus:
+        access_token = self._decrypt_token(account.access_token_encrypted)
+        refresh_token = self._decrypt_token(account.refresh_token_encrypted)
+        id_token = self._decrypt_token(account.id_token_encrypted)
+
+        access_expires = _token_expiry(access_token)
+        refresh_state = "stored" if refresh_token else "missing"
+        id_state = "unknown"
+        if id_token:
+            claims = extract_id_token_claims(id_token)
+            if claims.model_dump(exclude_none=True):
+                id_state = "parsed"
+
+        return AccountAuthStatus(
+            access=AccountTokenStatus(expires_at=access_expires),
+            refresh=AccountTokenStatus(state=refresh_state),
+            id_token=AccountTokenStatus(state=id_state),
+        )
+
+    def _decrypt_token(self, encrypted: bytes | None) -> str | None:
+        if not encrypted:
+            return None
+        try:
+            return self._encryptor.decrypt(encrypted)
+        except Exception:
+            return None
+
+    async def _refresh_usage(self, accounts: list[Account]) -> None:
+        if not self._usage_repo or not self._usage_updater:
+            return
+        latest_usage = await self._usage_repo.latest_by_account(window="primary")
+        await self._usage_updater.refresh_accounts(accounts, latest_usage)
+
+    async def _costs_last_24h(self) -> dict[str, float]:
+        if not self._logs_repo:
+            return {}
+        since = utcnow() - timedelta(hours=24)
+        logs = await self._logs_repo.list_since(since)
+        totals: dict[str, float] = {}
+        for log in logs:
+            cost = cost_from_log(log)
+            if cost is None:
+                continue
+            totals[log.account_id] = totals.get(log.account_id, 0.0) + cost
+        return {account_id: round(total, 6) for account_id, total in totals.items()}
+
+
+def _token_expiry(token: str | None) -> datetime | None:
+    if not token:
+        return None
+    claims = extract_id_token_claims(token)
+    exp = claims.exp
+    if isinstance(exp, (int, float)):
+        return datetime.fromtimestamp(exp, tz=timezone.utc)
+    if isinstance(exp, str) and exp.isdigit():
+        return datetime.fromtimestamp(int(exp), tz=timezone.utc)
+    return None
+
+
+def _normalize_used_percent(entry: UsageHistory | None) -> float | None:
+    if not entry:
+        return None
+    return entry.used_percent

app/modules/health/api.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from fastapi import APIRouter
+
+router = APIRouter(tags=["health"])
+
+
+@router.get("/health")
+async def health_check() -> dict:
+    return {"status": "ok"}

app/modules/oauth/api.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from fastapi import APIRouter, Body, Depends
+from fastapi.responses import JSONResponse
+
+from app.core.clients.oauth import OAuthError
+from app.core.errors import dashboard_error
+from app.dependencies import OauthContext, get_oauth_context
+from app.modules.oauth.schemas import (
+    OauthCompleteRequest,
+    OauthCompleteResponse,
+    OauthStartRequest,
+    OauthStartResponse,
+    OauthStatusResponse,
+)
+
+router = APIRouter(prefix="/api/oauth", tags=["dashboard"])
+
+
+@router.post("/start", response_model=OauthStartResponse)
+async def start_oauth(
+    request: OauthStartRequest,
+    context: OauthContext = Depends(get_oauth_context),
+) -> OauthStartResponse | JSONResponse:
+    try:
+        return await context.service.start_oauth(request)
+    except OAuthError as exc:
+        return JSONResponse(
+            status_code=502,
+            content=dashboard_error(exc.code, exc.message),
+        )
+    except NotImplementedError:
+        return JSONResponse(
+            status_code=501,
+            content=dashboard_error("not_implemented", "OAuth start is not implemented"),
+        )
+
+
+@router.get("/status", response_model=OauthStatusResponse)
+async def oauth_status(
+    context: OauthContext = Depends(get_oauth_context),
+) -> OauthStatusResponse | JSONResponse:
+    return await context.service.oauth_status()
+
+
+@router.post("/complete", response_model=OauthCompleteResponse)
+async def complete_oauth(
+    request: OauthCompleteRequest | None = Body(default=None),
+    context: OauthContext = Depends(get_oauth_context),
+) -> OauthCompleteResponse | JSONResponse:
+    try:
+        return await context.service.complete_oauth(request)
+    except NotImplementedError:
+        return JSONResponse(
+            status_code=501,
+            content=dashboard_error("not_implemented", "OAuth complete is not implemented"),
+        )

app/modules/oauth/schemas.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from app.modules.shared.schemas import DashboardModel
+
+
+class OauthStartRequest(DashboardModel):
+    force_method: str | None = None
+
+
+class OauthStartResponse(DashboardModel):
+    method: str
+    authorization_url: str | None = None
+    callback_url: str | None = None
+    verification_url: str | None = None
+    user_code: str | None = None
+    device_auth_id: str | None = None
+    interval_seconds: int | None = None
+    expires_in_seconds: int | None = None
+
+
+class OauthStatusResponse(DashboardModel):
+    status: str
+    error_message: str | None = None
+
+
+class OauthCompleteRequest(DashboardModel):
+    device_auth_id: str | None = None
+    user_code: str | None = None
+
+
+class OauthCompleteResponse(DashboardModel):
+    status: str

app/modules/oauth/service.py

@@ -0,0 +1,356 @@
+from __future__ import annotations
+
+import asyncio
+import secrets
+import time
+from contextlib import AbstractAsyncContextManager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Awaitable, Callable
+
+from aiohttp import web
+
+from app.core.auth import (
+    DEFAULT_EMAIL,
+    DEFAULT_PLAN,
+    OpenAIAuthClaims,
+    extract_id_token_claims,
+    fallback_account_id,
+)
+from app.core.clients.oauth import (
+    OAuthError,
+    OAuthTokens,
+    build_authorization_url,
+    exchange_authorization_code,
+    exchange_device_token,
+    generate_pkce_pair,
+    request_device_code,
+)
+from app.core.config.settings import get_settings
+from app.core.crypto import TokenEncryptor
+from app.core.utils.time import utcnow
+from app.db.models import Account, AccountStatus
+from app.modules.accounts.repository import AccountsRepository
+from app.modules.oauth.schemas import (
+    OauthCompleteRequest,
+    OauthCompleteResponse,
+    OauthStartRequest,
+    OauthStartResponse,
+    OauthStatusResponse,
+)
+
+_async_sleep = asyncio.sleep
+_SUCCESS_TEMPLATE = Path(__file__).resolve().parent / "templates" / "oauth_success.html"
+
+
+@dataclass
+class OAuthState:
+    status: str = "pending"
+    method: str | None = None
+    error_message: str | None = None
+    state_token: str | None = None
+    code_verifier: str | None = None
+    device_auth_id: str | None = None
+    user_code: str | None = None
+    interval_seconds: int | None = None
+    expires_at: float | None = None
+    callback_server: "OAuthCallbackServer | None" = None
+    poll_task: asyncio.Task[None] | None = None
+
+
+class OAuthStateStore:
+    def __init__(self) -> None:
+        self._lock = asyncio.Lock()
+        self._state = OAuthState(status="idle")
+
+    @property
+    def lock(self) -> asyncio.Lock:
+        return self._lock
+
+    @property
+    def state(self) -> OAuthState:
+        return self._state
+
+    async def reset(self) -> None:
+        async with self._lock:
+            await self._cleanup_locked()
+            self._state = OAuthState(status="idle")
+
+    async def _cleanup_locked(self) -> None:
+        task = self._state.poll_task
+        if task and not task.done():
+            task.cancel()
+        server = self._state.callback_server
+        if server:
+            await server.stop()
+
+
+class OAuthCallbackServer:
+    def __init__(
+        self,
+        handler: Callable[[web.Request], Awaitable[web.StreamResponse]],
+        host: str = "127.0.0.1",
+        port: int = 1455,
+    ) -> None:
+        self._handler = handler
+        self._host = host
+        self._port = port
+        self._runner: web.AppRunner | None = None
+        self._site: web.TCPSite | None = None
+
+    async def start(self) -> None:
+        app = web.Application()
+        app.router.add_get("/auth/callback", self._handler)
+        self._runner = web.AppRunner(app)
+        await self._runner.setup()
+        self._site = web.TCPSite(self._runner, self._host, self._port)
+        await self._site.start()
+
+    async def stop(self) -> None:
+        if self._runner:
+            await self._runner.cleanup()
+        self._runner = None
+        self._site = None
+
+
+_OAUTH_STORE = OAuthStateStore()
+
+
+class OauthService:
+    def __init__(
+        self,
+        accounts_repo: AccountsRepository,
+        repo_factory: Callable[[], AbstractAsyncContextManager[AccountsRepository]] | None = None,
+    ) -> None:
+        self._accounts_repo = accounts_repo
+        self._encryptor = TokenEncryptor()
+        self._store = _OAUTH_STORE
+        self._repo_factory = repo_factory
+
+    async def start_oauth(self, request: OauthStartRequest) -> OauthStartResponse:
+        force_method = (request.force_method or "").lower()
+        if not force_method:
+            accounts = await self._accounts_repo.list_accounts()
+            if accounts:
+                async with self._store.lock:
+                    await self._store._cleanup_locked()
+                    self._store._state = OAuthState(status="success")
+                return OauthStartResponse(method="browser")
+
+        if force_method == "device":
+            return await self._start_device_flow()
+
+        try:
+            return await self._start_browser_flow()
+        except OSError:
+            return await self._start_device_flow()
+
+    async def oauth_status(self) -> OauthStatusResponse:
+        async with self._store.lock:
+            state = self._store.state
+            status = state.status if state.status != "idle" else "pending"
+            return OauthStatusResponse(status=status, error_message=state.error_message)
+
+    async def complete_oauth(self, request: OauthCompleteRequest | None = None) -> OauthCompleteResponse:
+        payload = request or OauthCompleteRequest()
+        async with self._store.lock:
+            state = self._store.state
+            if payload.device_auth_id:
+                state.device_auth_id = payload.device_auth_id
+            if payload.user_code:
+                state.user_code = payload.user_code
+            if state.status == "success":
+                return OauthCompleteResponse(status="success")
+            if state.method != "device":
+                return OauthCompleteResponse(status="pending")
+            if state.poll_task and not state.poll_task.done():
+                return OauthCompleteResponse(status="pending")
+            if not state.device_auth_id or not state.user_code or not state.expires_at:
+                state.status = "error"
+                state.error_message = "Device code flow is not initialized."
+                return OauthCompleteResponse(status="error")
+
+            interval = state.interval_seconds if state.interval_seconds is not None else 0
+            interval = max(interval, 0)
+            poll_context = DevicePollContext(
+                device_auth_id=state.device_auth_id,
+                user_code=state.user_code,
+                interval_seconds=interval,
+                expires_at=state.expires_at,
+            )
+            state.poll_task = asyncio.create_task(self._poll_device_tokens(poll_context))
+            return OauthCompleteResponse(status="pending")
+
+    async def _start_browser_flow(self) -> OauthStartResponse:
+        await self._store.reset()
+        code_verifier, code_challenge = generate_pkce_pair()
+        state_token = secrets.token_urlsafe(16)
+        authorization_url = build_authorization_url(state=state_token, code_challenge=code_challenge)
+        settings = get_settings()
+
+        async with self._store.lock:
+            state = self._store.state
+            state.status = "pending"
+            state.method = "browser"
+            state.state_token = state_token
+            state.code_verifier = code_verifier
+            state.error_message = None
+
+        callback_server = OAuthCallbackServer(
+            self._handle_callback,
+            host=settings.oauth_callback_host,
+            port=settings.oauth_callback_port,
+        )
+        await callback_server.start()
+
+        async with self._store.lock:
+            self._store.state.callback_server = callback_server
+
+        return OauthStartResponse(
+            method="browser",
+            authorization_url=authorization_url,
+            callback_url=settings.oauth_redirect_uri,
+        )
+
+    async def _start_device_flow(self) -> OauthStartResponse:
+        await self._store.reset()
+        try:
+            device = await request_device_code()
+        except OAuthError as exc:
+            await self._set_error(exc.message)
+            raise
+
+        async with self._store.lock:
+            state = self._store.state
+            state.status = "pending"
+            state.method = "device"
+            state.device_auth_id = device.device_auth_id
+            state.user_code = device.user_code
+            state.interval_seconds = device.interval_seconds
+            state.expires_at = time.time() + device.expires_in_seconds
+            state.error_message = None
+
+        return OauthStartResponse(
+            method="device",
+            verification_url=device.verification_url,
+            user_code=device.user_code,
+            device_auth_id=device.device_auth_id,
+            interval_seconds=device.interval_seconds,
+            expires_in_seconds=device.expires_in_seconds,
+        )
+
+    async def _handle_callback(self, request: web.Request) -> web.Response:
+        params = request.rel_url.query
+        error = params.get("error")
+        code = params.get("code")
+        state = params.get("state")
+
+        if error:
+            await self._set_error(f"OAuth error: {error}")
+            return self._html_response(_error_html("Authorization failed."))
+
+        async with self._store.lock:
+            expected_state = self._store.state.state_token
+            verifier = self._store.state.code_verifier
+
+        if not code or not state or state != expected_state or not verifier:
+            await self._set_error("Invalid OAuth callback state.")
+            return self._html_response(_error_html("Invalid OAuth callback."))
+
+        try:
+            tokens = await exchange_authorization_code(code=code, code_verifier=verifier)
+            await self._persist_tokens(tokens)
+            await self._set_success()
+            html = _success_html()
+        except OAuthError as exc:
+            await self._set_error(exc.message)
+            html = _error_html(exc.message)
+
+        asyncio.create_task(self._stop_callback_server())
+        return self._html_response(html)
+
+    async def _poll_device_tokens(self, context: "DevicePollContext") -> None:
+        try:
+            while time.time() < context.expires_at:
+                tokens = await exchange_device_token(
+                    device_auth_id=context.device_auth_id,
+                    user_code=context.user_code,
+                )
+                if tokens:
+                    await self._persist_tokens(tokens)
+                    await self._set_success()
+                    return
+                await _async_sleep(context.interval_seconds)
+            await self._set_error("Device code expired.")
+        except OAuthError as exc:
+            await self._set_error(exc.message)
+        finally:
+            async with self._store.lock:
+                current = asyncio.current_task()
+                if self._store.state.poll_task is current:
+                    self._store.state.poll_task = None
+
+    async def _persist_tokens(self, tokens: OAuthTokens) -> None:
+        claims = extract_id_token_claims(tokens.id_token)
+        auth_claims = claims.auth or OpenAIAuthClaims()
+        account_id = auth_claims.chatgpt_account_id or claims.chatgpt_account_id
+        email = claims.email or DEFAULT_EMAIL
+        plan_type = auth_claims.chatgpt_plan_type or claims.chatgpt_plan_type or DEFAULT_PLAN
+        account_id = account_id or fallback_account_id(email)
+
+        account = Account(
+            id=account_id,
+            email=email,
+            plan_type=plan_type,
+            access_token_encrypted=self._encryptor.encrypt(tokens.access_token),
+            refresh_token_encrypted=self._encryptor.encrypt(tokens.refresh_token),
+            id_token_encrypted=self._encryptor.encrypt(tokens.id_token),
+            last_refresh=utcnow(),
+            status=AccountStatus.ACTIVE,
+            deactivation_reason=None,
+        )
+        if self._repo_factory:
+            async with self._repo_factory() as repo:
+                await repo.upsert(account)
+        else:
+            await self._accounts_repo.upsert(account)
+
+    async def _set_success(self) -> None:
+        async with self._store.lock:
+            self._store.state.status = "success"
+            self._store.state.error_message = None
+
+    async def _set_error(self, message: str) -> None:
+        async with self._store.lock:
+            self._store.state.status = "error"
+            self._store.state.error_message = message
+
+    async def _stop_callback_server(self) -> None:
+        async with self._store.lock:
+            server = self._store.state.callback_server
+            self._store.state.callback_server = None
+        if server:
+            await server.stop()
+
+    @staticmethod
+    def _html_response(html: str) -> web.Response:
+        return web.Response(text=html, content_type="text/html")
+
+
+@dataclass(frozen=True)
+class DevicePollContext:
+    device_auth_id: str
+    user_code: str
+    interval_seconds: int
+    expires_at: float
+
+
+def _success_html() -> str:
+    try:
+        return _SUCCESS_TEMPLATE.read_text(encoding="utf-8")
+    except OSError:
+        return "<html><body><h1>Login complete</h1><p>Return to the dashboard.</p></body></html>"
+
+
+def _error_html(message: str) -> str:
+    return f"<html><body><h1>Login failed</h1><p>{message}</p></body></html>"