codex-quota 1.0.0__py3-none-any.whl

codex_quota/__init__.py

@@ -0,0 +1,5 @@
+# src/codex_quota/__init__.py
+
+"""codex-quota: isolated Codex CLI profile management and quota monitoring."""
+
+__version__ = "1.0.0"

codex_quota/__main__.py

@@ -0,0 +1,3 @@
+from .cli import app
+
+app()

codex_quota/app_server.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+import json
+import subprocess
+import time
+from typing import Any
+
+from . import __version__
+from .codex import codex_env, find_codex
+from .errors import AppServerError
+from .filesystem import ensure_private_dir
+from .models import Profile, RateLimits, WindowLimit
+
+
+def _window(data: object) -> WindowLimit | None:
+    if not isinstance(data, dict):
+        return None
+    return WindowLimit(
+        used_percent=int(data.get("usedPercent", 0)),
+        window_duration_mins=int(data.get("windowDurationMins", 0)),
+        resets_at=data.get("resetsAt"),
+        raw=data,
+    )
+
+
+def decode_rate_limits(payload: dict[str, Any]) -> RateLimits:
+    rate_limits = payload.get("rateLimits") or payload.get("result", {}).get("rateLimits")
+    if not isinstance(rate_limits, dict):
+        raise AppServerError(f"Unexpected rate limit response: {payload}")
+
+    credits_data = rate_limits.get("credits")
+    return RateLimits(
+        primary=_window(rate_limits.get("primary")),
+        secondary=_window(rate_limits.get("secondary")),
+        plan_type=rate_limits.get("planType"),
+        rate_limit_reached_type=rate_limits.get("rateLimitReachedType"),
+        credits=credits_data if isinstance(credits_data, dict) else {},
+        raw=rate_limits,
+    )
+
+
+class AppServerClient:
+    def __init__(
+        self,
+        profile: Profile,
+        *,
+        codex_bin: str = "codex",
+        timeout: float = 30.0,
+    ) -> None:
+        self.profile = profile
+        self.codex_bin = codex_bin
+        self.timeout = timeout
+
+    def read_rate_limits(self) -> RateLimits:
+        binary = find_codex(self.codex_bin)
+        ensure_private_dir(self.profile.codex_home, parents=True, exist_ok=True)
+        with subprocess.Popen(
+            [binary, "app-server", "--stdio"],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            env=codex_env(self.profile),
+        ) as proc:
+            stdin = proc.stdin
+            stdout = proc.stdout
+            stderr = proc.stderr
+            if stdin is None or stdout is None or stderr is None:
+                raise AppServerError("codex app-server stdio pipes were not available")
+
+            def send(obj: dict[str, Any]) -> None:
+                stdin.write(json.dumps(obj) + "\n")
+                stdin.flush()
+
+            try:
+                request_id = 1
+                send(
+                    {
+                        "jsonrpc": "2.0",
+                        "id": request_id,
+                        "method": "initialize",
+                        "params": {
+                            "clientInfo": {
+                                "name": "codex-quota",
+                                "version": __version__,
+                            },
+                            "capabilities": {"experimentalApi": True},
+                        },
+                    }
+                )
+
+                deadline = time.monotonic() + self.timeout
+                initialized = False
+
+                while time.monotonic() < deadline:
+                    line = stdout.readline()
+
+                    if not line:
+                        if proc.poll() is not None:
+                            stderr_output = stderr.read()
+                            raise AppServerError(
+                                stderr_output.strip() or "codex app-server exited early"
+                            )
+                        time.sleep(0.02)
+                        continue
+
+                    try:
+                        msg = json.loads(line)
+                    except json.JSONDecodeError:
+                        continue
+
+                    if not isinstance(msg, dict):
+                        continue
+
+                    if msg.get("error"):
+                        raise AppServerError(json.dumps(msg["error"]))
+
+                    if msg.get("id") == 1 and not initialized:
+                        initialized = True
+                        send({"jsonrpc": "2.0", "method": "initialized", "params": {}})
+                        request_id = 2
+                        send(
+                            {
+                                "jsonrpc": "2.0",
+                                "id": request_id,
+                                "method": "account/rateLimits/read",
+                                "params": {},
+                            }
+                        )
+                        continue
+
+                    if msg.get("id") == request_id:
+                        result = msg.get("result")
+                        if not isinstance(result, dict):
+                            raise AppServerError(f"Malformed JSON-RPC response: {msg}")
+                        return decode_rate_limits(result)
+
+                raise AppServerError("Timed out waiting for Codex app-server rate limits")
+            finally:
+                proc.terminate()
+                try:
+                    proc.wait(timeout=2)
+                except subprocess.TimeoutExpired:
+                    proc.kill()
+
+
+def read_rate_limits(
+    profile: Profile,
+    codex_bin: str = "codex",
+    timeout: float = 30.0,
+) -> RateLimits:
+    return AppServerClient(profile, codex_bin=codex_bin, timeout=timeout).read_rate_limits()

codex_quota/cli.py

@@ -0,0 +1,308 @@
+from __future__ import annotations
+
+import json
+import sys
+from typing import Annotated, Any
+
+import typer
+from rich.console import Console
+
+from .codex import codex_version, simple_exec
+from .codex import login as codex_login
+from .config import ensure_config, load_config
+from .errors import CodexUsageError, ProfileNotFoundError
+from .models import AppConfig, ProfileStatus
+from .paths import default_paths
+from .profiles import add_profile, discover_profiles, ensure_profiles_dir, get_profile, profile_path
+from .profiles import remove_profile as remove_profile_dir
+from .render import render_profile_table, render_status_table
+from .service import check_failed, collect_statuses
+from .tui import run_fullscreen_tui
+
+app = typer.Typer(no_args_is_help=True, help="Manage Codex CLI profiles and quota usage.")
+profile_app = typer.Typer(no_args_is_help=True, help="Manage isolated Codex profiles.")
+app.add_typer(profile_app, name="profile")
+console = Console()
+
+
+def _config() -> AppConfig:
+    return load_config()
+
+
+def _print_error(exc: Exception) -> None:
+    console.print(f"[red]{exc}[/red]")
+
+
+def _window_payload(window: Any) -> dict[str, Any] | None:
+    if window is None:
+        return None
+    return {
+        "used_percent": window.used_percent,
+        "remaining_percent": window.remaining_percent,
+        "window_duration_mins": window.window_duration_mins,
+        "resets_at": window.resets_at,
+    }
+
+
+def _rate_limits_payload(
+    status: ProfileStatus,
+    *,
+    include_raw: bool = False,
+) -> dict[str, Any] | None:
+    limits = status.rate_limits
+    if limits is None:
+        return None
+
+    payload: dict[str, Any] = {
+        "blocked": limits.is_blocked,
+        "rate_limit_reached_type": limits.rate_limit_reached_type,
+        "primary": _window_payload(limits.primary),
+        "secondary": _window_payload(limits.secondary),
+    }
+    if include_raw:
+        payload["plan_type"] = limits.plan_type
+        payload["credits"] = limits.credits
+        payload["raw"] = limits.raw
+    return payload
+
+
+def _status_payload(
+    statuses: list[ProfileStatus],
+    *,
+    include_path: bool = False,
+    include_raw: bool = False,
+) -> list[dict[str, Any]]:
+    return [
+        {
+            "profile": item.profile.name,
+            "auth_ok": item.auth_ok,
+            "codex_ok": item.codex_ok,
+            "ok": item.ok,
+            "error": item.error,
+            **({"path": str(item.profile.codex_home)} if include_path else {}),
+            "rate_limits": _rate_limits_payload(item, include_raw=include_raw),
+        }
+        for item in statuses
+    ]
+
+
+@app.callback()
+def main() -> None:
+    """Manage isolated Codex profiles and inspect quota usage."""
+
+
+@app.command()
+def init(
+    overwrite_config: Annotated[bool, typer.Option("--overwrite-config")] = False,
+    verify: Annotated[bool, typer.Option("--verify", "--check")] = False,
+) -> None:
+    """Bootstrap directories and validate the local Codex environment."""
+    try:
+        config = ensure_config(overwrite=overwrite_config)
+        ensure_profiles_dir(config)
+        profiles = discover_profiles(config)
+
+        paths = default_paths()
+        console.print(f"Config file: {paths.config_file}")
+        console.print(f"Profiles dir: {config.profiles_dir}")
+        console.print(render_profile_table(profiles))
+
+        try:
+            console.print(f"Codex: {codex_version(config.codex_bin)}")
+        except CodexUsageError as exc:
+            console.print(f"[yellow]{exc}[/yellow]")
+            if verify:
+                raise typer.Exit(2) from exc
+
+        if verify:
+            statuses = collect_statuses(config)
+            console.print(render_status_table(statuses))
+            raise typer.Exit(2 if check_failed(statuses) else 0)
+    except CodexUsageError as exc:
+        _print_error(exc)
+        raise typer.Exit(2) from exc
+
+
+@profile_app.command("list")
+def profile_list() -> None:
+    """List profiles discovered under profiles_dir."""
+    try:
+        console.print(render_profile_table(discover_profiles(_config())))
+    except CodexUsageError as exc:
+        _print_error(exc)
+        raise typer.Exit(2) from exc
+
+
+@profile_app.command("add")
+def profile_add(
+    name: Annotated[str, typer.Argument(help="Profile name")],
+    exist_ok: Annotated[bool, typer.Option("--exist-ok", "--force")] = False,
+    run_login: Annotated[bool, typer.Option("--login")] = False,
+    device_auth: Annotated[
+        bool,
+        typer.Option("--device-auth", help="Pass --device-auth to codex login."),
+    ] = False,
+) -> None:
+    """Create a new isolated profile directory."""
+    try:
+        if device_auth and not run_login:
+            raise CodexUsageError("--device-auth requires --login")
+        config = _config()
+        profile = add_profile(config, name, exist_ok=exist_ok)
+        console.print(f"Created profile: {profile.name}")
+        console.print(str(profile.codex_home))
+
+        if run_login:
+            raise typer.Exit(codex_login(profile, config.codex_bin, device_auth=device_auth))
+    except CodexUsageError as exc:
+        _print_error(exc)
+        raise typer.Exit(2) from exc
+
+
+@profile_app.command("remove")
+def profile_remove(
+    name: Annotated[str, typer.Argument(help="Profile name")],
+    yes: Annotated[bool, typer.Option("--yes", "-y")] = False,
+) -> None:
+    """Remove a profile directory after safety checks."""
+    try:
+        config = _config()
+        profile = get_profile(config, name)
+
+        if profile.has_auth:
+            console.print("[yellow]Warning: auth.json exists in this profile.[/yellow]")
+
+        if not yes:
+            confirmed = typer.confirm(f"Delete profile '{name}' at {profile.codex_home}?")
+            if not confirmed:
+                raise typer.Exit(1)
+
+        removed = remove_profile_dir(config, name)
+        console.print(f"Removed profile: {removed.name}")
+    except CodexUsageError as exc:
+        _print_error(exc)
+        raise typer.Exit(2) from exc
+
+
+@profile_app.command("path")
+def profile_path_cmd(name: Annotated[str, typer.Argument(help="Profile name")]) -> None:
+    """Print the CODEX_HOME path for scripting."""
+    try:
+        typer.echo(str(profile_path(_config(), name)))
+    except CodexUsageError as exc:
+        _print_error(exc)
+        raise typer.Exit(2) from exc
+
+
+@app.command()
+def login(
+    profile_name: Annotated[str, typer.Argument(help="Profile name")],
+    create: Annotated[bool, typer.Option("--create")] = False,
+    device_auth: Annotated[
+        bool,
+        typer.Option("--device-auth", help="Pass --device-auth to codex login."),
+    ] = False,
+) -> None:
+    """Run the official Codex login flow for an existing profile."""
+    try:
+        config = _config()
+        profile = (
+            add_profile(config, profile_name, exist_ok=True)
+            if create
+            else get_profile(config, profile_name)
+        )
+        raise typer.Exit(codex_login(profile, config.codex_bin, device_auth=device_auth))
+    except CodexUsageError as exc:
+        _print_error(exc)
+        raise typer.Exit(2) from exc
+
+
+@app.command()
+def status(
+    profile_name: Annotated[str | None, typer.Argument(help="Optional profile name")] = None,
+    json_out: Annotated[bool, typer.Option("--json", help="Print machine-readable JSON.")] = False,
+    json_paths: Annotated[
+        bool,
+        typer.Option("--json-paths", help="Include absolute profile paths in JSON output."),
+    ] = False,
+    json_raw: Annotated[
+        bool,
+        typer.Option("--json-raw", help="Include raw rate-limit payload in JSON output."),
+    ] = False,
+) -> None:
+    """Read quota status for one or all discovered profiles."""
+    try:
+        statuses = collect_statuses(_config(), profile_name)
+        if json_out:
+            console.print_json(
+                json.dumps(_status_payload(statuses, include_path=json_paths, include_raw=json_raw))
+            )
+        else:
+            console.print(render_status_table(statuses))
+    except ProfileNotFoundError as exc:
+        _print_error(exc)
+        raise typer.Exit(2) from exc
+    except CodexUsageError as exc:
+        _print_error(exc)
+        raise typer.Exit(2) from exc
+
+
+@app.command()
+def check(
+    profile_name: Annotated[str | None, typer.Argument(help="Optional profile name")] = None,
+) -> None:
+    """Exit non-zero if any selected profile has auth, quota, or Codex problems."""
+    try:
+        statuses = collect_statuses(_config(), profile_name)
+        console.print(render_status_table(statuses))
+        raise typer.Exit(2 if check_failed(statuses) else 0)
+    except CodexUsageError as exc:
+        _print_error(exc)
+        raise typer.Exit(2) from exc
+
+
+@app.command("exec")
+def exec_prompt(
+    profile_name: Annotated[str, typer.Argument(help="Profile name")],
+    prompt: Annotated[str, typer.Argument(help="Prompt to send")],
+    raw: Annotated[bool, typer.Option("--raw", help="Do not force Codex JSON output.")] = False,
+) -> None:
+    """Run official codex exec under a selected CODEX_HOME."""
+    try:
+        config = _config()
+        proc = simple_exec(get_profile(config, profile_name), prompt, config, json_output=not raw)
+
+        if proc.stdout:
+            console.print(proc.stdout.rstrip())
+        if proc.stderr:
+            console.print(proc.stderr.rstrip(), style="yellow")
+
+        raise typer.Exit(proc.returncode)
+    except CodexUsageError as exc:
+        _print_error(exc)
+        raise typer.Exit(2) from exc
+
+
+@app.command()
+def tui(refresh: Annotated[float | None, typer.Option("--refresh", "-r")] = None) -> None:
+    """Open a Textual live quota view."""
+    try:
+        config = _config()
+        refresh_seconds = refresh or config.refresh_seconds
+
+        if refresh_seconds <= 0:
+            raise typer.BadParameter("--refresh must be greater than zero")
+
+        if not sys.stdin.isatty() or not console.is_terminal:
+            statuses = collect_statuses(config)
+            console.print(render_status_table(statuses))
+            raise typer.Exit(2 if check_failed(statuses) else 0)
+
+        run_fullscreen_tui(config, refresh_seconds)
+    except CodexUsageError as exc:
+        _print_error(exc)
+        raise typer.Exit(2) from exc
+
+
+if __name__ == "__main__":
+    app()

codex_quota/codex.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+
+from .errors import CodexCommandError, CodexNotFoundError
+from .filesystem import ensure_private_dir
+from .models import AppConfig, Profile
+
+
+def find_codex(codex_bin: str = "codex") -> str:
+    found = shutil.which(codex_bin)
+    if not found:
+        raise CodexNotFoundError(
+            "Codex CLI was not found. Install the official Codex CLI and ensure "
+            "it is available on PATH, or set codex_bin in config.json."
+        )
+    return found
+
+
+def codex_env(profile: Profile) -> dict[str, str]:
+    return {**os.environ, "CODEX_HOME": str(profile.codex_home)}
+
+
+def run_codex(
+    args: list[str],
+    *,
+    profile: Profile | None = None,
+    codex_bin: str = "codex",
+    timeout: float = 60.0,
+    check: bool = False,
+) -> subprocess.CompletedProcess[str]:
+    binary = find_codex(codex_bin)
+    env = codex_env(profile) if profile else os.environ.copy()
+    try:
+        proc = subprocess.run(
+            [binary, *args],
+            text=True,
+            capture_output=True,
+            env=env,
+            timeout=timeout,
+            check=False,
+        )
+    except subprocess.TimeoutExpired as exc:
+        raise CodexCommandError(f"codex {' '.join(args)} timed out") from exc
+    if check and proc.returncode != 0:
+        raise CodexCommandError(
+            proc.stderr.strip() or proc.stdout.strip() or "codex command failed"
+        )
+    return proc
+
+
+def codex_version(codex_bin: str = "codex") -> str:
+    proc = run_codex(["--version"], codex_bin=codex_bin, timeout=10)
+    if proc.returncode != 0:
+        raise CodexCommandError(proc.stderr.strip() or "Failed to read Codex CLI version")
+    return proc.stdout.strip() or proc.stderr.strip() or "unknown"
+
+
+def login_status(profile: Profile, codex_bin: str = "codex") -> bool:
+    if not profile.auth_file.is_file():
+        return False
+    proc = run_codex(["login", "status"], profile=profile, codex_bin=codex_bin, timeout=20)
+    output = f"{proc.stdout}\n{proc.stderr}".lower()
+    return proc.returncode == 0 and ("logged in" in output or "authenticated" in output)
+
+
+def login(
+    profile: Profile,
+    codex_bin: str = "codex",
+    *,
+    device_auth: bool = False,
+) -> int:
+    binary = find_codex(codex_bin)
+    ensure_private_dir(profile.codex_home, parents=True, exist_ok=True)
+    args = [binary, "login"]
+    if device_auth:
+        args.append("--device-auth")
+    return subprocess.call(args, env=codex_env(profile))
+
+
+def simple_exec(
+    profile: Profile,
+    prompt: str,
+    config: AppConfig,
+    *,
+    json_output: bool = True,
+    timeout: float = 180.0,
+) -> subprocess.CompletedProcess[str]:
+    args = [
+        "exec",
+        "--skip-git-repo-check",
+        "--ephemeral",
+        "-m",
+        config.default_model,
+        "-c",
+        f'model_reasoning_effort="{config.reasoning_effort}"',
+    ]
+    if json_output:
+        args.append("--json")
+    args.append(prompt)
+    return run_codex(args, profile=profile, codex_bin=config.codex_bin, timeout=timeout)

codex_quota/config.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import json
+from json import JSONDecodeError
+from pathlib import Path
+from typing import Any
+
+from .errors import ConfigError
+from .filesystem import ensure_private_dir
+from .models import AppConfig
+from .paths import default_paths
+
+CONFIG_KEYS = {"profiles_dir", "codex_bin", "default_model", "reasoning_effort", "refresh_seconds"}
+
+
+def create_default_config() -> AppConfig:
+    paths = default_paths()
+    return AppConfig(profiles_dir=paths.profiles_dir)
+
+
+def _as_path(value: object, fallback: Path) -> Path:
+    if not isinstance(value, str) or not value.strip():
+        return fallback
+    return Path(value).expanduser()
+
+
+def load_config(path: Path | None = None) -> AppConfig:
+    paths = default_paths()
+    config_file = path or paths.config_file
+    if not config_file.exists():
+        return create_default_config()
+
+    try:
+        data = json.loads(config_file.read_text(encoding="utf-8"))
+    except JSONDecodeError as exc:
+        raise ConfigError(f"Malformed config JSON: {config_file}: {exc}") from exc
+    except OSError as exc:
+        raise ConfigError(f"Could not read config: {config_file}: {exc}") from exc
+
+    if not isinstance(data, dict):
+        raise ConfigError(f"Config must be a JSON object: {config_file}")
+
+    refresh = data.get("refresh_seconds", 30.0)
+    try:
+        refresh_seconds = float(refresh)
+    except (TypeError, ValueError) as exc:
+        raise ConfigError("config refresh_seconds must be a number") from exc
+
+    extra = {key: value for key, value in data.items() if key not in CONFIG_KEYS}
+    return AppConfig(
+        profiles_dir=_as_path(data.get("profiles_dir"), paths.profiles_dir),
+        codex_bin=str(data.get("codex_bin", "codex")),
+        default_model=str(data.get("default_model", "gpt-5.4-mini")),
+        reasoning_effort=str(data.get("reasoning_effort", "low")),
+        refresh_seconds=refresh_seconds,
+        extra=extra,
+    )
+
+
+def config_to_json(config: AppConfig, *, preserve_extra: bool = True) -> dict[str, Any]:
+    data: dict[str, Any] = dict(config.extra) if preserve_extra else {}
+    data.update(
+        {
+            "profiles_dir": str(config.profiles_dir),
+            "codex_bin": config.codex_bin,
+            "default_model": config.default_model,
+            "reasoning_effort": config.reasoning_effort,
+            "refresh_seconds": config.refresh_seconds,
+        }
+    )
+    return data
+
+
+def save_config(
+    config: AppConfig,
+    path: Path | None = None,
+    *,
+    preserve_extra: bool = True,
+) -> Path:
+    paths = default_paths()
+    config_file = path or paths.config_file
+    ensure_private_dir(config_file.parent, parents=True, exist_ok=True)
+    data = config_to_json(config, preserve_extra=preserve_extra)
+    try:
+        config_file.write_text(
+            json.dumps(data, indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+    except OSError as exc:
+        raise ConfigError(f"Could not write config: {config_file}: {exc}") from exc
+    return config_file
+
+
+def ensure_config(path: Path | None = None, *, overwrite: bool = False) -> AppConfig:
+    paths = default_paths()
+    config_file = path or paths.config_file
+    ensure_private_dir(config_file.parent, parents=True, exist_ok=True)
+    if config_file.exists() and not overwrite:
+        return load_config(config_file)
+    config = create_default_config()
+    save_config(config, config_file, preserve_extra=False)
+    return config