@0dai-dev/cli 4.3.6 → 4.3.8

package/lib/python/agent_quotas.py

@@ -0,0 +1,525 @@
+#!/usr/bin/env python3
+"""Agent subscription quota poller.
+
+Polls each known coding-agent CLI (codex, qoder, claude, opencode, gemini, capi)
+for its current subscription usage and caches the result in
+``ai/manifest/agent-quotas.json``. The cache TTL is 10 minutes; callers should
+treat the file as read-mostly runtime state (it is gitignored).
+
+Public surface
+--------------
+- ``get_quotas(target=None, force_refresh=False) -> dict``
+    Returns ``{agent: {plan, weekly_pct_remaining, hourly_pct_remaining,
+    resets_at, source, error}}`` for every known agent. Deprecated
+    ``*_pct_used`` aliases are still emitted for old dashboards/cache readers,
+    but subscription banners expose remaining quota.
+- ``poll_agent(agent) -> dict``
+    Polls a single agent and returns its quota record (no cache I/O).
+- ``CACHE_TTL_SECONDS`` - default 600s.
+
+The poller is defensive: a missing CLI, a non-zero exit, an unparseable
+banner, or a tmux failure all yield a record with ``error`` set, never
+an exception. The dispatcher (``swarm_lib.dispatch_candidates_for_task``)
+treats any agent without a clean quota record as "unknown" and keeps it
+in the candidate list.
+
+Codex usage banner is captured by attaching tmux to ``codex --status``
+output; the regex looks for the canonical
+``Context X% left * Context Y% used * 5h Z% * w W%`` shape.
+Qoder is queried via ``qodercli quota --json`` if the CLI is on PATH.
+Claude / OpenCode / Gemini / CAPI fall back to "unknown" gracefully when
+their respective CLIs do not expose a quota subcommand on this host.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import pathlib
+import re
+import shutil
+import subprocess
+import time
+from typing import Any, Callable
+
+log = logging.getLogger("0dai.agent_quotas")
+
+KNOWN_AGENTS: tuple[str, ...] = (
+    "codex",
+    "qoder",
+    "claude",
+    "opencode",
+    "gemini",
+    "capi",
+)
+
+CACHE_TTL_SECONDS = 600  # 10 minutes per the spec
+DEFAULT_TARGET = pathlib.Path.cwd()
+CACHE_RELATIVE = pathlib.Path("ai") / "manifest" / "agent-quotas.json"
+
+# Codex usage banner pattern, e.g.
+#   Context 41% left · Context 59% used · 5h 12% · w 38%
+# The trailing weekly group is optional because narrow tmux panes truncate
+# the banner to e.g. "5h 96% ·…" — we still want hourly + context%
+# even when weekly is cut off.
+CODEX_USAGE_RE = re.compile(
+    r"Context\s+(?P<ctx_left>\d+)%\s+left\s*[·•|-]\s*"
+    r"Context\s+(?P<ctx_used>\d+)%\s+used\s*[·•|-]\s*"
+    r"5h\s+(?P<hourly>\d+)%"
+    r"(?:\s*[·•|-]\s*w\s+(?P<weekly>\d+)%)?",
+    re.IGNORECASE,
+)
+
+# Generic "X% used / Y% remaining" fallback
+GENERIC_PCT_RE = re.compile(r"(?P<n>\d{1,3})\s*%")
+
+
+# ---------------------------------------------------------------------------
+# Cache I/O
+# ---------------------------------------------------------------------------
+
+
+def _resolve_target(target: str | os.PathLike[str] | None) -> pathlib.Path:
+    if target is None:
+        return DEFAULT_TARGET
+    return pathlib.Path(target).resolve()
+
+
+def _cache_path(target: pathlib.Path) -> pathlib.Path:
+    return target / CACHE_RELATIVE
+
+
+def _load_cache(path: pathlib.Path) -> dict[str, Any] | None:
+    try:
+        return json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+
+
+def _save_cache(path: pathlib.Path, payload: dict[str, Any]) -> None:
+    try:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
+    except OSError as exc:
+        log.warning("agent_quotas cache write failed at %s: %s", path, exc)
+
+
+def _is_fresh(cache: dict[str, Any] | None, now: float, ttl: int) -> bool:
+    if not cache:
+        return False
+    fetched = cache.get("fetched_at")
+    if not isinstance(fetched, (int, float)):
+        return False
+    return (now - float(fetched)) < ttl
+
+
+# ---------------------------------------------------------------------------
+# Per-agent pollers
+# ---------------------------------------------------------------------------
+
+
+def _empty_record(
+    agent: str, *, error: str | None = None, source: str = "unknown"
+) -> dict[str, Any]:
+    return {
+        "agent": agent,
+        "plan": "unknown",
+        "weekly_pct_remaining": None,
+        "hourly_pct_remaining": None,
+        # Deprecated compatibility aliases. The stored values are remaining,
+        # matching historical cache semantics despite the old field names.
+        "weekly_pct_used": None,
+        "hourly_pct_used": None,
+        "resets_at": None,
+        "source": source,
+        "error": error,
+    }
+
+
+def _coerce_pct(value: Any) -> int | None:
+    if value is None or isinstance(value, bool):
+        return None
+    try:
+        pct = int(value)
+    except (TypeError, ValueError):
+        return None
+    return max(0, min(100, pct))
+
+
+def _first_present(data: dict[str, Any], keys: tuple[str, ...]) -> Any:
+    for key in keys:
+        if key in data and data[key] is not None:
+            return data[key]
+    return None
+
+
+def _with_remaining_fields(
+    *,
+    weekly: Any = None,
+    hourly: Any = None,
+) -> dict[str, int | None]:
+    weekly_pct = _coerce_pct(weekly)
+    hourly_pct = _coerce_pct(hourly)
+    return {
+        "weekly_pct_remaining": weekly_pct,
+        "hourly_pct_remaining": hourly_pct,
+        "weekly_pct_used": weekly_pct,
+        "hourly_pct_used": hourly_pct,
+    }
+
+
+def quota_pct_remaining(record: dict[str, Any] | None, window: str) -> float | None:
+    """Return quota remaining for ``weekly``/``hourly`` records.
+
+    New records use ``*_pct_remaining``. Old cache records only have
+    ``*_pct_used``; those values were already remaining but mislabeled, so they
+    remain a fail-open compatibility fallback.
+    """
+    if not isinstance(record, dict):
+        return None
+    canonical = record.get(f"{window}_pct_remaining")
+    if isinstance(canonical, (int, float)) and not isinstance(canonical, bool):
+        return float(canonical)
+    legacy = record.get(f"{window}_pct_used")
+    if isinstance(legacy, (int, float)) and not isinstance(legacy, bool):
+        return float(legacy)
+    return None
+
+
+def _which(binary: str) -> str | None:
+    return shutil.which(binary)
+
+
+def _run(cmd: list[str], *, timeout: int = 8) -> tuple[int, str, str]:
+    """Run a subprocess and capture stdout/stderr without raising."""
+    try:
+        proc = subprocess.run(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            timeout=timeout,
+            check=False,
+        )
+        out = proc.stdout.decode("utf-8", errors="replace") if proc.stdout else ""
+        err = proc.stderr.decode("utf-8", errors="replace") if proc.stderr else ""
+        return proc.returncode, out, err
+    except (FileNotFoundError, subprocess.TimeoutExpired, OSError) as exc:
+        return 1, "", f"{type(exc).__name__}: {exc}"
+
+
+def _parse_codex_banner(text: str) -> dict[str, Any] | None:
+    m = CODEX_USAGE_RE.search(text)
+    if not m:
+        return None
+    weekly_raw = m.group("weekly")
+    return _with_remaining_fields(
+        weekly=weekly_raw if weekly_raw is not None else None,
+        hourly=m.group("hourly"),
+    )
+
+
+def poll_codex(
+    *, runner: Callable[[list[str]], tuple[int, str, str]] = _run
+) -> dict[str, Any]:
+    """Poll codex via tmux capture if available, otherwise via ``codex --status``."""
+    rec = _empty_record("codex")
+    bin_path = _which("codex")
+    if not bin_path:
+        rec["error"] = "codex CLI not found on PATH"
+        return rec
+    rec["plan"] = "subscription"
+
+    # Try direct status first - cheap and avoids tmux churn
+    code, out, err = runner([bin_path, "--status"])
+    text = (out or "") + "\n" + (err or "")
+    parsed = _parse_codex_banner(text)
+    if parsed:
+        rec.update(parsed)
+        rec["source"] = "codex --status"
+        return rec
+
+    # Fall back to tmux session capture. Spec: "tmux capture + regex".
+    # Use -x 220 so the codex TUI status bar fits on one line — narrower
+    # panes truncate "5h 96% · w 38%" to "5h 96% ·…" via Unicode ellipsis,
+    # which breaks the banner regex even with the weekly group made optional.
+    if _which("tmux"):
+        session = f"0dai-quota-codex-{int(time.time())}"
+        try:
+            runner(
+                [
+                    "tmux", "new-session", "-d", "-s", session,
+                    "-x", "220", "-y", "50",
+                    bin_path,
+                ]
+            )
+            time.sleep(5.0)  # codex TUI banner takes ~3-5s to render after spawn
+            cap_code, cap_out, _cap_err = runner(
+                ["tmux", "capture-pane", "-pt", session]
+            )
+            parsed = _parse_codex_banner(cap_out or "")
+            if parsed:
+                rec.update(parsed)
+                rec["source"] = "tmux capture"
+                return rec
+        finally:
+            runner(["tmux", "kill-session", "-t", session])
+
+    rec["error"] = "codex usage banner not found"
+    return rec
+
+
+def poll_qoder(
+    *, runner: Callable[[list[str]], tuple[int, str, str]] = _run
+) -> dict[str, Any]:
+    rec = _empty_record("qoder")
+    bin_path = _which("qodercli") or _which("qoder")
+    if not bin_path:
+        rec["error"] = "qoder CLI not found on PATH"
+        return rec
+    rec["plan"] = "subscription"
+
+    # Prefer JSON if the CLI supports it; fall back to plain text.
+    for argset in (["quota", "--json"], ["quota"], ["--quota"], ["status", "--json"]):
+        code, out, err = runner([bin_path, *argset])
+        if code != 0 and not out:
+            continue
+        try:
+            data = json.loads(out)
+        except (ValueError, TypeError):
+            data = None
+        if isinstance(data, dict):
+            weekly = _first_present(
+                data,
+                (
+                    "weekly_pct_remaining",
+                    "weekly_remaining_pct",
+                    "weekly_remaining",
+                    "weekly",
+                ),
+            )
+            hourly = _first_present(
+                data,
+                (
+                    "hourly_pct_remaining",
+                    "hourly_remaining_pct",
+                    "hourly_remaining",
+                    "hourly",
+                ),
+            )
+            weekly_used = _first_present(data, ("weekly_pct_used", "weekly_used_pct"))
+            hourly_used = _first_present(data, ("hourly_pct_used", "hourly_used_pct"))
+            if weekly is None and weekly_used is not None:
+                used = _coerce_pct(weekly_used)
+                weekly = None if used is None else 100 - used
+            if hourly is None and hourly_used is not None:
+                used = _coerce_pct(hourly_used)
+                hourly = None if used is None else 100 - used
+            plan = data.get("plan") or data.get("tier")
+            resets = data.get("resets_at") or data.get("reset_at")
+            if weekly is not None or hourly is not None or plan:
+                rec.update(_with_remaining_fields(weekly=weekly, hourly=hourly))
+                if plan:
+                    rec["plan"] = str(plan)
+                if resets:
+                    rec["resets_at"] = str(resets)
+                rec["source"] = f"{bin_path} {' '.join(argset)}"
+                return rec
+        # Plain-text fallback: scrape any "weekly: X%" / "hourly: Y%" lines
+        weekly_m = re.search(r"weekly[^0-9%]*(\d{1,3})\s*%", out, re.IGNORECASE)
+        hourly_m = re.search(r"(?:hourly|5h)[^0-9%]*(\d{1,3})\s*%", out, re.IGNORECASE)
+        if weekly_m or hourly_m:
+            rec.update(
+                _with_remaining_fields(
+                    weekly=weekly_m.group(1) if weekly_m else None,
+                    hourly=hourly_m.group(1) if hourly_m else None,
+                )
+            )
+            rec["source"] = f"{bin_path} {' '.join(argset)}"
+            return rec
+
+    rec["error"] = "qoder quota output not parseable"
+    return rec
+
+
+def _poll_unknown(agent: str) -> dict[str, Any]:
+    """Return a record indicating the agent has no programmatic quota source."""
+    bin_path = _which(agent)
+    rec = _empty_record(agent)
+    if not bin_path:
+        rec["error"] = f"{agent} CLI not found on PATH"
+        return rec
+    rec["plan"] = "subscription"
+    rec["source"] = "no-quota-api"
+    rec["error"] = f"{agent} CLI does not expose a quota endpoint"
+    return rec
+
+
+def poll_claude() -> dict[str, Any]:
+    return _poll_unknown("claude")
+
+
+def poll_opencode() -> dict[str, Any]:
+    return _poll_unknown("opencode")
+
+
+def poll_gemini() -> dict[str, Any]:
+    return _poll_unknown("gemini")
+
+
+def poll_capi() -> dict[str, Any]:
+    rec = _poll_unknown("capi")
+    # capi rides on the user's account-local token; treat as "unknown"
+    # but mark plan as 'managed' so dispatcher knows it's available.
+    if rec.get("error", "").endswith("not found on PATH"):
+        return rec
+    rec["plan"] = "managed"
+    return rec
+
+
+_POLLERS: dict[str, Callable[[], dict[str, Any]]] = {
+    "codex": poll_codex,
+    "qoder": poll_qoder,
+    "claude": poll_claude,
+    "opencode": poll_opencode,
+    "gemini": poll_gemini,
+    "capi": poll_capi,
+}
+
+
+def poll_agent(agent: str) -> dict[str, Any]:
+    poller = _POLLERS.get(agent)
+    if not poller:
+        return _empty_record(agent, error=f"unknown agent: {agent}")
+    try:
+        return poller()
+    except Exception as exc:  # noqa: BLE001 — intentional: defensive, never let a third-party poller raise out
+        log.warning("poll_agent(%s) failed: %s", agent, exc)
+        return _empty_record(agent, error=f"poller exception: {exc}")
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def refresh_quotas(target: str | os.PathLike[str] | None = None) -> dict[str, Any]:
+    """Force a poll of every known agent and persist to cache."""
+    target_path = _resolve_target(target)
+    now = time.time()
+    agents: dict[str, Any] = {}
+    for agent in KNOWN_AGENTS:
+        agents[agent] = poll_agent(agent)
+    payload = {
+        "fetched_at": now,
+        "fetched_at_iso": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(now)),
+        "ttl_seconds": CACHE_TTL_SECONDS,
+        "agents": agents,
+    }
+    _save_cache(_cache_path(target_path), payload)
+    return payload
+
+
+def get_quotas(
+    target: str | os.PathLike[str] | None = None,
+    *,
+    force_refresh: bool = False,
+    ttl_seconds: int = CACHE_TTL_SECONDS,
+) -> dict[str, Any]:
+    """Return cached quota data, refreshing if stale or forced.
+
+    Returns the ``agents`` mapping directly: ``{agent: record}``.
+    The full envelope (with ``fetched_at`` etc.) is available via
+    :func:`get_quotas_envelope`.
+    """
+    return get_quotas_envelope(
+        target, force_refresh=force_refresh, ttl_seconds=ttl_seconds
+    )["agents"]
+
+
+def get_quotas_envelope(
+    target: str | os.PathLike[str] | None = None,
+    *,
+    force_refresh: bool = False,
+    ttl_seconds: int = CACHE_TTL_SECONDS,
+) -> dict[str, Any]:
+    target_path = _resolve_target(target)
+    cache_file = _cache_path(target_path)
+    now = time.time()
+    cache = _load_cache(cache_file)
+    if not force_refresh and _is_fresh(cache, now, ttl_seconds):
+        return cache  # type: ignore[return-value]
+    return refresh_quotas(target_path)
+
+
+# ---------------------------------------------------------------------------
+# Helpers used by dispatcher
+# ---------------------------------------------------------------------------
+
+WEEKLY_EXHAUSTION_THRESHOLD = 5  # subscription banners expose remaining quota
+
+
+def is_agent_exhausted(
+    record: dict[str, Any] | None, threshold: int = WEEKLY_EXHAUSTION_THRESHOLD
+) -> bool:
+    """True iff the record reports weekly remaining quota at or below ``threshold``.
+
+    Records with no weekly data (None / unknown / error) are NOT considered
+    exhausted - the dispatcher should not penalize agents whose CLI lacks a
+    quota endpoint.
+    """
+    weekly = quota_pct_remaining(record, "weekly")
+    if weekly is None:
+        return False
+    return float(weekly) <= float(threshold)
+
+
+# ---------------------------------------------------------------------------
+# CLI entry point - useful for `python3 -m scripts.agent_quotas`
+# ---------------------------------------------------------------------------
+
+
+def _format_table(envelope: dict[str, Any]) -> str:
+    rows = [("AGENT", "PLAN", "WEEK%", "HOUR%", "RESETS", "SOURCE")]
+    for agent in KNOWN_AGENTS:
+        rec = envelope.get("agents", {}).get(agent) or _empty_record(agent)
+        weekly = quota_pct_remaining(rec, "weekly")
+        hourly = quota_pct_remaining(rec, "hourly")
+        rows.append(
+            (
+                agent,
+                str(rec.get("plan") or "unknown"),
+                "—" if weekly is None else f"{weekly}%",
+                "—" if hourly is None else f"{hourly}%",
+                str(rec.get("resets_at") or "—"),
+                str(rec.get("source") or rec.get("error") or "—"),
+            )
+        )
+    widths = [max(len(r[i]) for r in rows) for i in range(len(rows[0]))]
+    lines = []
+    for row in rows:
+        lines.append("  ".join(row[i].ljust(widths[i]) for i in range(len(row))))
+    return "\n".join(lines)
+
+
+def main(argv: list[str] | None = None) -> int:
+    import argparse
+
+    parser = argparse.ArgumentParser(description="0dai agent quota poller")
+    parser.add_argument("--target", default=None, help="Project root (default: cwd)")
+    parser.add_argument(
+        "--refresh", action="store_true", help="Force refresh, ignore cache TTL"
+    )
+    parser.add_argument("--json", action="store_true", help="Emit JSON envelope")
+    args = parser.parse_args(argv)
+
+    envelope = get_quotas_envelope(args.target, force_refresh=args.refresh)
+    if args.json:
+        print(json.dumps(envelope, indent=2, sort_keys=True))
+    else:
+        print(_format_table(envelope))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())