import asyncio import json import math from datetime import datetime, timezone from decimal import Decimal, InvalidOperation from typing import Any from zoneinfo import ZoneInfo from app.config import Settings from app.models import ProviderUsage, UsageLimit from app.providers.base import GENERIC_ERROR, UsageProvider, utc_now def _number(value: Any, *, minimum: float = 0, maximum: float | None = None) -> float | None: if isinstance(value, bool) or not isinstance(value, (int, float)): return None try: number = float(value) except (TypeError, ValueError): return None if not math.isfinite(number) or number < minimum or (maximum is not None and number > maximum): return None return number def _decimal(value: Any, *, minimum: float = 0) -> float | None: if value is None or isinstance(value, bool): return None try: number = float(Decimal(str(value))) except (InvalidOperation, ValueError, TypeError): return None if not math.isfinite(number) or number < minimum: return None return number def _reset(value: Any, timezone_name: str) -> datetime | None: epoch = _number(value, minimum=0) if epoch is None: return None try: return datetime.fromtimestamp(epoch, timezone.utc).astimezone(ZoneInfo(timezone_name)) except (OverflowError, OSError, ValueError): return None def _duration_name(minutes: Any) -> str: duration = _number(minutes, minimum=0) if duration is None: return 'Usage Limit' if duration == 300: return '5 Hour Limit' if duration == 1440: return 'Daily Limit' if duration == 10080: return 'Weekly Limit' if duration < 60: label, amount = 'Minute', duration elif duration < 1440: label, amount = 'Hour', duration / 60 else: label, amount = 'Day', duration / 1440 rendered = str(int(amount)) if amount == int(amount) else f'{amount:g}' return f'{rendered} {label} Limit' def _window_limit(snapshot: dict[str, Any], key: str, window: Any, updated: datetime, timezone_name: str, suffix: str) -> UsageLimit | None: if not isinstance(window, dict): return None used = _number(window.get('usedPercent'), minimum=0, maximum=100) if used is None: return None duration = window.get('windowDurationMins') name = _duration_name(duration) limit_name = snapshot.get('limitName') if name == 'Usage Limit' and isinstance(limit_name, str) and limit_name.strip(): name = limit_name.strip() return UsageLimit(id=f'{key}-{suffix}', name=name, used_percent=used, remaining_percent=100 - used, reset_at=_reset(window.get('resetsAt'), timezone_name), window=name.removesuffix(' Limit').lower(), updated_at=updated) def parse_codex_result(payload: Any, timezone_name: str, now: datetime | None = None) -> list[UsageLimit]: if not isinstance(payload, dict) or not isinstance(payload.get('result'), dict): raise ValueError result = payload['result'] snapshots: list[tuple[str, dict[str, Any]]] = [] mapped = result.get('rateLimitsByLimitId') if isinstance(mapped, dict) and mapped: if any(not isinstance(value, dict) for value in mapped.values()): raise ValueError snapshots = [(str(key), value) for key, value in mapped.items()] elif isinstance(result.get('rateLimits'), dict): snapshots = [('rateLimits', result['rateLimits'])] else: raise ValueError updated = now or utc_now() if updated.tzinfo is None: updated = updated.replace(tzinfo=timezone.utc) limits: list[UsageLimit] = [] for key, snapshot in snapshots: for suffix, field in (('primary', 'primary'), ('secondary', 'secondary')): window = snapshot.get(field) if window is not None and not isinstance(window, dict): raise ValueError if isinstance(window, dict): parsed = _window_limit(snapshot, key, window, updated, timezone_name, suffix) if parsed is None: raise ValueError limits.append(parsed) credits = snapshot.get('credits') if credits is not None and not isinstance(credits, dict): raise ValueError if isinstance(credits, dict): if not isinstance(credits.get('hasCredits'), bool) or not isinstance(credits.get('unlimited'), bool): raise ValueError if isinstance(credits, dict) and credits.get('hasCredits') is True: balance = None if credits.get('unlimited') is True else _decimal(credits.get('balance')) if credits.get('unlimited') is not True and credits.get('balance') is not None and balance is None: raise ValueError limits.append(UsageLimit(id=f'{key}-credits', name='Additional Credits', remaining=balance, unit='credits', updated_at=updated)) individual = snapshot.get('individualLimit') if individual is not None and not isinstance(individual, dict): raise ValueError if isinstance(individual, dict): total = _decimal(individual.get('limit')) used = _decimal(individual.get('used')) remaining_percent = _number(individual.get('remainingPercent'), minimum=0, maximum=100) if total is None or used is None or remaining_percent is None or used > total: raise ValueError limits.append(UsageLimit(id=f'{key}-individual', name='Spend Limit', used_percent=100 - remaining_percent, remaining_percent=remaining_percent, used=used, remaining=total - used, limit=total, unit=None, reset_at=_reset(individual.get('resetsAt'), timezone_name), window='individual', updated_at=updated)) if not limits: raise ValueError return limits class CodexProvider(UsageProvider): def __init__(self, settings: Settings): super().__init__('codex', settings.provider_cache_ttl, 'Codex app-server') self.settings = settings async def _fetch_usage(self) -> ProviderUsage: process = None stderr_task = None try: process = await asyncio.create_subprocess_exec( self.settings.codex_command, 'app-server', '--listen', 'stdio://', stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE) stderr_task = asyncio.create_task(self._discard(process.stderr)) async with asyncio.timeout(self.settings.codex_timeout): await self._send(process, {'method': 'initialize', 'id': 1, 'params': { 'clientInfo': {'name': 'ai_usage_dashboard', 'title': 'AI Usage Dashboard', 'version': '1.0.0'}, 'capabilities': {'optOutNotificationMethods': []}}}) await self._read_response(process.stdout, 1) await self._send(process, {'method': 'initialized'}) await self._send(process, {'method': 'account/rateLimits/read', 'id': 2}) response = await self._read_response(process.stdout, 2) limits = parse_codex_result(response, self.settings.timezone) return ProviderUsage(provider='codex', status='ok', limits=limits, last_successful_update=utc_now(), source=self.source) except Exception: return ProviderUsage(provider='codex', status='error', error=GENERIC_ERROR, last_successful_update=self.last_successful_update, source=self.source) finally: if process is not None: try: await self._close_process(process) except Exception: pass if stderr_task is not None: if not stderr_task.done(): stderr_task.cancel() await asyncio.gather(stderr_task, return_exceptions=True) @staticmethod async def _discard(stream: asyncio.StreamReader) -> None: while await stream.read(8192): pass @staticmethod async def _send(process: asyncio.subprocess.Process, message: dict[str, Any]) -> None: process.stdin.write((json.dumps(message, separators=(',', ':')) + '\n').encode()) await process.stdin.drain() @staticmethod async def _read_response(stdout: asyncio.StreamReader, expected_id: int) -> dict[str, Any]: while True: line = await stdout.readline() if not line: raise ValueError try: message = json.loads(line) except (json.JSONDecodeError, TypeError): raise ValueError if isinstance(message, dict) and message.get('id') == expected_id: if 'result' not in message or not isinstance(message['result'], dict): raise ValueError return message @staticmethod async def _close_process(process: asyncio.subprocess.Process) -> None: if process.returncode is not None: return try: process.terminate() except ProcessLookupError: pass try: await asyncio.wait_for(process.wait(), timeout=0.5) except asyncio.TimeoutError: try: process.kill() except ProcessLookupError: return try: await process.wait() except ProcessLookupError: pass except ProcessLookupError: pass