agent-skills/skills/blitz/usage-budget.py

#!/usr/bin/env python3
"""Estimate provider pressure before blitz dispatches workers.

The provider bars are account-side. Claude's existing nightshift estimator
counts local Claude transcripts; Codex rollouts sometimes include the live
rate-limit envelope. Missing data is deliberately treated as pressure: an
unattended blitz should spend less when it cannot prove that a window is clear.
"""

from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import subprocess
from pathlib import Path


BAND_WRAP = 60.0
BAND_PARK = 80.0


def parse_time(value: str | None) -> dt.datetime | None:
    if not value:
        return None
    try:
        return dt.datetime.fromisoformat(value.replace("Z", "+00:00"))
    except ValueError:
        return None


def run_claude() -> dict:
    script = Path(__file__).parents[1] / "nightshift" / "usage-window.py"
    try:
        result = subprocess.run(
            ["python3", str(script), "--json"],
            check=False,
            capture_output=True,
            text=True,
            timeout=10,
        )
        report = json.loads(result.stdout)
        report["provider"] = "claude"
        report["source"] = "local Claude transcripts"
        return report
    except (OSError, subprocess.SubprocessError, json.JSONDecodeError) as exc:
        return {"provider": "claude", "verdict": "unknown", "error": str(exc)}


def latest_codex_limit() -> dict:
    root = Path(os.environ.get("CODEX_HOME", Path.home() / ".codex")) / "sessions"
    latest: tuple[dt.datetime, dict] | None = None
    if not root.is_dir():
        return {"provider": "codex", "verdict": "unknown", "error": f"missing {root}"}

    for path in root.glob("**/*.jsonl"):
        try:
            handle = path.open(encoding="utf-8", errors="replace")
        except OSError:
            continue
        with handle:
            for line in handle:
                try:
                    event = json.loads(line)
                except json.JSONDecodeError:
                    continue
                payload = event.get("payload") or {}
                limits = payload.get("rate_limits")
                stamp = parse_time(event.get("timestamp"))
                if not isinstance(limits, dict) or stamp is None:
                    continue
                if latest is None or stamp > latest[0]:
                    latest = (stamp, limits)

    if latest is None:
        return {"provider": "codex", "verdict": "unknown", "error": "no rate-limit event"}

    stamp, limits = latest
    result = {
        "provider": "codex",
        "observedAt": stamp.isoformat(),
        "source": "latest Codex rollout rate_limits event",
        "planType": limits.get("plan_type"),
        "limitId": limits.get("limit_id"),
        "rateLimitReached": limits.get("rate_limit_reached_type"),
    }
    if result["rateLimitReached"]:
        result["verdict"] = "park"
        return result

    # Codex commonly reports the weekly window as primary and the shorter
    # window as secondary. Only use a percentage as a dispatch verdict when
    # the window is plausibly short; weekly usage is reported separately.
    short = limits.get("secondary") or {}
    primary = limits.get("primary") or {}
    if short.get("used_percent") is not None:
        result["windowMinutes"] = short.get("window_minutes")
        result["percentUsed"] = float(short["used_percent"])
    elif primary.get("window_minutes", 0) <= 600 and primary.get("used_percent") is not None:
        result["windowMinutes"] = primary.get("window_minutes")
        result["percentUsed"] = float(primary["used_percent"])
    else:
        result["weeklyPercentUsed"] = primary.get("used_percent")
        result["weeklyWindowMinutes"] = primary.get("window_minutes")
        result["verdict"] = "unknown"
        result["reason"] = "Codex supplied no short-window percentage"
        return result

    percent = result["percentUsed"]
    result["verdict"] = "park" if percent >= BAND_PARK else "wrap-up" if percent >= BAND_WRAP else "clear"
    return result


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--max-sessions", type=int, default=2)
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    providers = [run_claude(), latest_codex_limit()]
    verdicts = {p.get("verdict", "unknown") for p in providers}
    if "park" in verdicts:
        sessions = 0
    elif verdicts == {"clear"}:
        sessions = args.max_sessions
    else:
        # Unknown is not permission to fan out. One medium worker preserves
        # progress while leaving room for a provider to be more used than the
        # local evidence can see.
        sessions = min(1, args.max_sessions)

    result = {
        "providers": providers,
        "recommendedMaxSessions": sessions,
        "preferMedium": sessions < args.max_sessions or "park" in verdicts,
        "reason": "all provider windows clear" if verdicts == {"clear"} else "provider pressure or unknown usage",
    }
    if args.json:
        print(json.dumps(result, indent=2))
    else:
        for provider in providers:
            print(f"{provider['provider']}: {provider.get('verdict', 'unknown')}")
        print(f"recommended max sessions: {sessions}")
        print(f"prefer medium models: {'yes' if result['preferMedium'] else 'no'}")


if __name__ == "__main__":
    main()