agentpool-cli 0.1.0__py3-none-any.whl

agentpool/usage/ccusage.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+import os
+import shlex
+import shutil
+import subprocess
+from typing import Any
+
+from agentpool.models import CapacitySnapshot, Confidence, UsageStatus, UsageWindow, UsageWindowKind
+from agentpool.usage._common import (
+    ProbeError,
+    _clamp_percent,
+    _extract_json_payload,
+    _number,
+    _parse_datetime,
+    unavailable,
+    unknown,
+)
+
+
+def detect_ccusage(binary: str | None = None) -> dict[str, Any]:
+    command = _ccusage_command(binary)
+    if not command:
+        return {"installed": False, "path": None, "version": None, "safe_source": "local_claude_code_logs"}
+    version = None
+    try:
+        proc = subprocess.run([*command, "--version"], capture_output=True, text=True, timeout=5, check=False)
+        if proc.returncode == 0:
+            version = (proc.stdout or proc.stderr).strip().splitlines()[0][:200]
+    except (OSError, subprocess.TimeoutExpired):
+        version = None
+    return {
+        "installed": True,
+        "path": command[0],
+        "command": command,
+        "version": version,
+        "safe_source": "local_claude_code_logs",
+    }
+
+
+def ccusage_usage_snapshot(provider_id: str, binary: str | None = None) -> CapacitySnapshot:
+    if provider_id != "claude-code":
+        return unknown(provider_id, f"ccusage is only mapped for claude-code, not {provider_id}.", source="ccusage")
+    command = _ccusage_command(binary)
+    if not command:
+        return unavailable(
+            provider_id,
+            "ccusage CLI is not installed. Set AGENTPOOL_CCUSAGE_COMMAND to an explicit command if desired.",
+        )
+    try:
+        proc = subprocess.run(
+            [*command, "blocks", "--json", "--offline", "--active", "--no-color"],
+            capture_output=True,
+            text=True,
+            timeout=45,
+            check=False,
+        )
+    except (OSError, subprocess.TimeoutExpired) as exc:
+        return unknown(provider_id, f"ccusage probe failed: {exc}", source="ccusage")
+    text = "\n".join(part for part in [proc.stdout, proc.stderr] if part)
+    try:
+        payload = _extract_json_payload(text)
+        snapshot = parse_ccusage_blocks(provider_id, payload)
+        if proc.returncode != 0:
+            snapshot.warnings.append(f"ccusage exited with status {proc.returncode}.")
+        return snapshot
+    except ProbeError as exc:
+        return unknown(provider_id, f"ccusage probe failed: {exc}", source="ccusage")
+
+
+def parse_ccusage_blocks(provider_id: str, payload: Any) -> CapacitySnapshot:
+    if not isinstance(payload, dict):
+        raise ProbeError("ccusage output must be a JSON object.")
+    blocks = payload.get("blocks")
+    if not isinstance(blocks, list):
+        raise ProbeError("ccusage output did not include blocks.")
+    active = None
+    for block in blocks:
+        if isinstance(block, dict) and block.get("isActive") is True and not block.get("isGap"):
+            active = block
+    if active is None:
+        raise ProbeError("ccusage output did not include an active usage block.")
+    projection = active.get("projection") if isinstance(active.get("projection"), dict) else {}
+    burn_rate = active.get("burnRate") if isinstance(active.get("burnRate"), dict) else {}
+    token_counts = active.get("tokenCounts") if isinstance(active.get("tokenCounts"), dict) else {}
+    token_limit_status = active.get("tokenLimitStatus") if isinstance(active.get("tokenLimitStatus"), dict) else {}
+    token_limit_used = _number(token_limit_status.get("percentUsed"))
+    window = UsageWindow(
+        name="active_block",
+        kind=UsageWindowKind.FIVE_HOUR,
+        status="active_local_log_block",
+        used_percent=_clamp_percent(token_limit_used) if token_limit_used is not None else None,
+        used_units=_number(active.get("totalTokens")),
+        reset_at=_parse_datetime(active.get("endTime")),
+        confidence=Confidence.OBSERVED,
+        raw_text="ccusage:blocks:active",
+    )
+    raw = {
+        "source": "ccusage_local_logs",
+        "block_id": active.get("id"),
+        "start_time": active.get("startTime"),
+        "actual_end_time": active.get("actualEndTime"),
+        "entries": active.get("entries"),
+        "cost_usd": _number(active.get("costUSD")),
+        "models": active.get("models") if isinstance(active.get("models"), list) else [],
+        "token_counts": token_counts,
+        "burn_rate": burn_rate,
+        "projection": projection,
+        "token_limit_status": token_limit_status,
+    }
+    return CapacitySnapshot(
+        provider_id=provider_id,
+        status=UsageStatus.UNKNOWN,
+        confidence=Confidence.OBSERVED,
+        windows=[window],
+        warnings=["ccusage reads local Claude Code logs; it is telemetry, not authoritative quota."],
+        raw=raw,
+    )
+
+
+def _ccusage_command(binary: str | None = None) -> list[str] | None:
+    if binary:
+        return shlex.split(binary)
+    configured = os.environ.get("AGENTPOOL_CCUSAGE_COMMAND")
+    if configured:
+        return shlex.split(configured)
+    executable = shutil.which("ccusage")
+    if executable:
+        return [executable]
+    return None

agentpool/usage/claude.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+import shutil
+
+from agentpool.models import CapacitySnapshot
+from agentpool.usage._common import _tmux_slash_usage_probe, unavailable
+from agentpool.usage.provider_parsers import parse_claude_usage
+
+
+def claude_code_usage_snapshot(provider_id: str, binary: str | None = None) -> CapacitySnapshot:
+    executable = binary or shutil.which("claude")
+    if not executable:
+        return unavailable(provider_id, "Claude Code is not installed.")
+    return _tmux_slash_usage_probe(
+        provider_id=provider_id,
+        command=[executable, "--allowed-tools", ""],
+        slash_command="/usage",
+        parser=parse_claude_usage,
+        source="claude_pty_usage",
+        startup_delay=1.2,
+        timeout=18.0,
+        extra_keys_after_match=[["PageDown"]],
+    )

agentpool/usage/codex.py

@@ -0,0 +1,210 @@
+from __future__ import annotations
+
+import json
+import select
+import shutil
+import subprocess
+import time
+from typing import Any
+
+from agentpool.models import CapacitySnapshot, Confidence, UsageWindow, UsageWindowKind
+from agentpool.usage._common import (
+    ProbeError,
+    _clamp_percent,
+    _epoch_seconds,
+    _int_number,
+    _number,
+    _safe_read_pipe,
+    _status_from_windows,
+    _terminate_process,
+    unavailable,
+    unknown,
+)
+
+
+def codex_cli_usage_snapshot(provider_id: str, binary: str | None = None) -> CapacitySnapshot:
+    executable = binary or shutil.which("codex")
+    if not executable:
+        return unavailable(provider_id, "Codex CLI is not installed.")
+    try:
+        response = _codex_rpc_rate_limits(executable)
+        snapshot = parse_codex_rate_limits(provider_id, response)
+        snapshot.raw["source"] = "codex_app_server"
+        return snapshot
+    except ProbeError as exc:
+        recovered = _recover_codex_error_snapshot(provider_id, str(exc))
+        if recovered:
+            recovered.warnings.append("Codex app-server returned usage inside an error response.")
+            return recovered
+        return unknown(provider_id, f"Codex app-server usage probe failed: {exc}", source="codex_app_server")
+
+
+def parse_codex_rate_limits(provider_id: str, payload: dict[str, Any]) -> CapacitySnapshot:
+    rate_limits = payload.get("rateLimits") if isinstance(payload.get("rateLimits"), dict) else payload
+    windows: list[UsageWindow] = []
+    for fallback_name, item in (("primary", rate_limits.get("primary")), ("secondary", rate_limits.get("secondary"))):
+        if not isinstance(item, dict):
+            continue
+        used = _number(item.get("usedPercent"))
+        if used is None:
+            continue
+        duration = _int_number(item.get("windowDurationMins"))
+        name = _codex_window_name(fallback_name, duration)
+        reset_at = _epoch_seconds(item.get("resetsAt"))
+        windows.append(
+            UsageWindow(
+                name=name,
+                kind=_codex_window_kind(name),
+                used_percent=_clamp_percent(used),
+                remaining_percent=_clamp_percent(100.0 - used),
+                reset_at=reset_at,
+                confidence=Confidence.OFFICIAL,
+                raw_text=f"{fallback_name}:{duration or 'unknown'}",
+            )
+        )
+    credits_remaining = None
+    credits = rate_limits.get("credits")
+    if isinstance(credits, dict) and not credits.get("unlimited"):
+        credits_remaining = _number(credits.get("balance"))
+    if not windows and credits_remaining is None:
+        raise ProbeError("No rate-limit windows or credits in Codex response.")
+    return CapacitySnapshot(
+        provider_id=provider_id,
+        status=_status_from_windows(windows),
+        confidence=Confidence.OFFICIAL,
+        windows=windows,
+        credits_remaining=credits_remaining,
+        raw={"credits": _safe_credit_summary(credits) if isinstance(credits, dict) else None},
+    )
+
+
+def _codex_rpc_rate_limits(executable: str) -> dict[str, Any]:
+    proc = subprocess.Popen(
+        [executable, "-s", "read-only", "-a", "untrusted", "app-server"],
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+    try:
+        _json_rpc_request(proc, 1, "initialize", {"clientInfo": {"name": "agentpool", "version": "0.1.0"}}, 8.0)
+        _json_rpc_notify(proc, "initialized")
+        return _json_rpc_request(proc, 2, "account/rateLimits/read", {}, 4.0)
+    finally:
+        _terminate_process(proc)
+
+
+def _json_rpc_request(
+    proc: subprocess.Popen[str],
+    request_id: int,
+    method: str,
+    params: dict[str, Any],
+    timeout: float,
+) -> dict[str, Any]:
+    _write_json_line(proc, {"id": request_id, "method": method, "params": params})
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        if proc.poll() is not None:
+            stderr = _safe_read_pipe(proc.stderr)
+            raise ProbeError(f"codex app-server exited during {method}: {stderr}".strip())
+        remaining = max(0.05, deadline - time.monotonic())
+        assert proc.stdout is not None
+        ready, _, _ = select.select([proc.stdout], [], [], remaining)
+        if not ready:
+            continue
+        line = proc.stdout.readline()
+        if not line:
+            continue
+        try:
+            message = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+        if message.get("id") != request_id:
+            continue
+        if "error" in message:
+            error = message["error"]
+            raise ProbeError(error.get("message", str(error)) if isinstance(error, dict) else str(error))
+        result = message.get("result")
+        if not isinstance(result, dict):
+            raise ProbeError(f"JSON-RPC {method} returned no object result.")
+        return result
+    _terminate_process(proc)
+    raise ProbeError(f"codex app-server timed out during {method}.")
+
+
+def _json_rpc_notify(proc: subprocess.Popen[str], method: str) -> None:
+    _write_json_line(proc, {"method": method, "params": {}})
+
+
+def _write_json_line(proc: subprocess.Popen[str], payload: dict[str, Any]) -> None:
+    if proc.stdin is None:
+        raise ProbeError("process stdin is unavailable")
+    proc.stdin.write(json.dumps(payload, separators=(",", ":")) + "\n")
+    proc.stdin.flush()
+
+
+def _recover_codex_error_snapshot(provider_id: str, message: str) -> CapacitySnapshot | None:
+    start = message.find("{")
+    end = message.rfind("}")
+    if start < 0 or end <= start:
+        return None
+    try:
+        payload = json.loads(message[start : end + 1])
+    except json.JSONDecodeError:
+        return None
+    rate_limit = payload.get("rate_limit")
+    if not isinstance(rate_limit, dict):
+        return None
+    windows = []
+    details = rate_limit.get("details") if isinstance(rate_limit.get("details"), list) else []
+    for item in details:
+        if not isinstance(item, dict):
+            continue
+        used = _number(item.get("used_percent"))
+        seconds = _int_number(item.get("limit_window_seconds"))
+        reset = _epoch_seconds(item.get("reset_at"))
+        if used is None:
+            continue
+        windows.append(
+            UsageWindow(
+                name=_codex_window_name("window", int(seconds / 60) if seconds else None),
+                kind=_codex_window_kind(_codex_window_name("window", int(seconds / 60) if seconds else None)),
+                used_percent=_clamp_percent(used),
+                remaining_percent=_clamp_percent(100.0 - used),
+                reset_at=reset,
+                confidence=Confidence.OFFICIAL,
+            )
+        )
+    if not windows:
+        return None
+    return CapacitySnapshot(
+        provider_id=provider_id,
+        status=_status_from_windows(windows),
+        confidence=Confidence.OFFICIAL,
+        windows=windows,
+        raw={"source": "codex_app_server_error_body"},
+    )
+
+
+def _codex_window_name(fallback: str, duration_mins: int | None) -> str:
+    if duration_mins == 300:
+        return "5h"
+    if duration_mins == 10080:
+        return "weekly"
+    return fallback
+
+
+def _codex_window_kind(name: str) -> UsageWindowKind:
+    if name == "5h":
+        return UsageWindowKind.FIVE_HOUR
+    if name == "weekly":
+        return UsageWindowKind.WEEKLY
+    return UsageWindowKind.UNKNOWN
+
+
+def _safe_credit_summary(credits: dict[str, Any]) -> dict[str, Any]:
+    return {
+        "hasCredits": bool(credits.get("hasCredits")),
+        "unlimited": bool(credits.get("unlimited")),
+        "hasBalance": credits.get("balance") is not None,
+    }

agentpool/usage/codexbar.py

@@ -0,0 +1,186 @@
+from __future__ import annotations
+
+import shutil
+import subprocess
+from typing import Any
+
+from agentpool.models import CapacitySnapshot, Confidence, UsageWindow, UsageWindowKind
+from agentpool.usage._common import (
+    ProbeError,
+    _clamp_percent,
+    _duration_window_kind,
+    _extract_json_payload,
+    _int_number,
+    _number,
+    _parse_datetime,
+    _status_from_windows,
+    _clean_optional_string,
+    unavailable,
+    unknown,
+)
+
+CODEXBAR_PROVIDER_MAP = {
+    "claude-code": "claude",
+    "codex-cli": "codex",
+    "copilot-cli": "copilot",
+    "cursor-cli": "cursor",
+    "opencode": "opencode",
+}
+
+CODEXBAR_SAFE_SOURCE_MAP = {
+    "claude-code": "cli",
+    "codex-cli": "cli",
+    "copilot-cli": "api",
+    "cursor-cli": "cli",
+    "opencode": "api",
+}
+
+
+def detect_codexbar(binary: str | None = None) -> dict[str, Any]:
+    executable = binary or shutil.which("codexbar")
+    if not executable:
+        return {
+            "installed": False,
+            "path": None,
+            "version": None,
+            "supported_agentpool_providers": sorted(CODEXBAR_PROVIDER_MAP),
+            "safe_sources": CODEXBAR_SAFE_SOURCE_MAP,
+        }
+    version = None
+    try:
+        proc = subprocess.run([executable, "--version"], capture_output=True, text=True, timeout=3, check=False)
+        if proc.returncode == 0:
+            version = (proc.stdout or proc.stderr).strip().splitlines()[0][:200]
+    except (OSError, subprocess.TimeoutExpired):
+        version = None
+    return {
+        "installed": True,
+        "path": executable,
+        "version": version,
+        "supported_agentpool_providers": sorted(CODEXBAR_PROVIDER_MAP),
+        "safe_sources": CODEXBAR_SAFE_SOURCE_MAP,
+    }
+
+
+def codexbar_usage_snapshot(
+    provider_id: str,
+    binary: str | None = None,
+    source: str | None = None,
+) -> CapacitySnapshot:
+    executable = binary or shutil.which("codexbar")
+    if not executable:
+        return unavailable(provider_id, "CodexBar CLI is not installed.")
+    codexbar_provider = CODEXBAR_PROVIDER_MAP.get(provider_id)
+    if not codexbar_provider:
+        return unknown(
+            provider_id,
+            f"CodexBar does not expose a safe mapped usage provider for {provider_id}.",
+            source="codexbar",
+        )
+    safe_source = source or CODEXBAR_SAFE_SOURCE_MAP.get(provider_id)
+    if not safe_source:
+        return unknown(provider_id, f"No safe CodexBar source is configured for {provider_id}.", source="codexbar")
+    command = [
+        executable,
+        "usage",
+        "--provider",
+        codexbar_provider,
+        "--source",
+        safe_source,
+        "--format",
+        "json",
+        "--json-only",
+        "--no-color",
+    ]
+    try:
+        proc = subprocess.run(command, capture_output=True, text=True, timeout=45, check=False)
+    except (OSError, subprocess.TimeoutExpired) as exc:
+        return unknown(provider_id, f"CodexBar usage probe failed: {exc}", source="codexbar")
+    text = "\n".join(part for part in [proc.stdout, proc.stderr] if part)
+    try:
+        payload = _extract_json_payload(text)
+        snapshot = parse_codexbar_usage(provider_id, payload, expected_provider=codexbar_provider)
+        snapshot.raw["source"] = "codexbar"
+        snapshot.raw["codexbar_provider"] = codexbar_provider
+        snapshot.raw["codexbar_source"] = safe_source
+        if proc.returncode != 0:
+            snapshot.warnings.append(f"CodexBar exited with status {proc.returncode}.")
+        return snapshot
+    except ProbeError as exc:
+        return unknown(provider_id, f"CodexBar usage probe failed: {exc}", source="codexbar")
+
+
+def parse_codexbar_usage(
+    provider_id: str,
+    payload: Any,
+    expected_provider: str | None = None,
+) -> CapacitySnapshot:
+    entries = payload if isinstance(payload, list) else [payload]
+    entry = None
+    for item in entries:
+        if not isinstance(item, dict):
+            continue
+        if expected_provider is None or item.get("provider") == expected_provider:
+            entry = item
+            break
+    if entry is None:
+        raise ProbeError("CodexBar output did not include the requested provider.")
+    error = entry.get("error")
+    if isinstance(error, dict):
+        raise ProbeError(_clean_optional_string(error.get("message")) or str(error))
+    usage = entry.get("usage")
+    if not isinstance(usage, dict):
+        raise ProbeError("CodexBar output did not include a usage object.")
+    windows = []
+    for name in ("primary", "secondary", "tertiary"):
+        item = usage.get(name)
+        if not isinstance(item, dict):
+            continue
+        used = _number(item.get("usedPercent"))
+        remaining = _number(item.get("remainingPercent"))
+        if used is None and remaining is None:
+            continue
+        if remaining is None and used is not None:
+            remaining = 100.0 - used
+        if used is None and remaining is not None:
+            used = 100.0 - remaining
+        duration = _int_number(item.get("windowMinutes"))
+        windows.append(
+            UsageWindow(
+                name=_codexbar_window_name(name, duration),
+                kind=_duration_window_kind(duration),
+                status=name,
+                used_percent=_clamp_percent(used) if used is not None else None,
+                remaining_percent=_clamp_percent(remaining) if remaining is not None else None,
+                reset_at=_parse_datetime(item.get("resetsAt")),
+                confidence=Confidence.LOCAL_CLI,
+                raw_text=f"codexbar:{name}:{duration or 'unknown'}",
+            )
+        )
+    credits_remaining = None
+    credits = entry.get("credits")
+    if isinstance(credits, dict):
+        credits_remaining = _number(credits.get("remaining"))
+    if not windows and credits_remaining is None:
+        raise ProbeError("CodexBar output did not include parseable windows or credits.")
+    raw: dict[str, Any] = {
+        "codexbar_source": entry.get("source"),
+        "version": entry.get("version"),
+        "login_method": _clean_optional_string(usage.get("loginMethod")),
+        "account_email": _clean_optional_string(usage.get("accountEmail")),
+    }
+    return CapacitySnapshot(
+        provider_id=provider_id,
+        status=_status_from_windows(windows),
+        confidence=Confidence.LOCAL_CLI,
+        windows=windows,
+        credits_remaining=credits_remaining,
+        raw={key: value for key, value in raw.items() if value is not None},
+    )
+
+
+def _codexbar_window_name(fallback: str, duration_mins: int | None) -> str:
+    kind = _duration_window_kind(duration_mins)
+    if kind != UsageWindowKind.UNKNOWN:
+        return kind.value
+    return fallback

agentpool/usage/combine.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from typing import Any
+
+from agentpool.models import CapacitySnapshot, UsageStatus, UsageWindow
+from agentpool.usage._common import _status_from_windows
+
+
+def combine_usage_snapshots(
+    native: CapacitySnapshot,
+    codexbar: CapacitySnapshot,
+    ccusage: CapacitySnapshot | None = None,
+) -> CapacitySnapshot:
+    if codexbar.status in {UsageStatus.UNKNOWN, UsageStatus.UNAVAILABLE, UsageStatus.UNAUTHENTICATED}:
+        snapshot = native.model_copy(deep=True)
+        snapshot.raw["alternate_sources"] = [_usage_source_summary(codexbar)]
+        if codexbar.warnings:
+            snapshot.warnings.extend(f"CodexBar: {warning}" for warning in codexbar.warnings)
+        return _add_ccusage_enrichment(snapshot, ccusage)
+    if native.status in {UsageStatus.UNKNOWN, UsageStatus.UNAVAILABLE, UsageStatus.UNAUTHENTICATED}:
+        snapshot = codexbar.model_copy(deep=True)
+        snapshot.raw["source"] = "combined"
+        snapshot.raw["primary_source"] = "codexbar"
+        snapshot.raw["alternate_sources"] = [_usage_source_summary(native)]
+        if native.warnings:
+            snapshot.warnings.extend(f"Native probe: {warning}" for warning in native.warnings)
+        return _add_ccusage_enrichment(snapshot, ccusage)
+    snapshot = native.model_copy(deep=True)
+    existing = {_usage_window_key(window) for window in snapshot.windows}
+    for window in codexbar.windows:
+        key = _usage_window_key(window)
+        if key not in existing:
+            snapshot.windows.append(window)
+            existing.add(key)
+    if snapshot.credits_remaining is None and codexbar.credits_remaining is not None:
+        snapshot.credits_remaining = codexbar.credits_remaining
+    snapshot.status = _status_from_windows(snapshot.windows)
+    snapshot.raw["source"] = "combined"
+    snapshot.raw["primary_source"] = native.raw.get("source", "native")
+    snapshot.raw["alternate_sources"] = [_usage_source_summary(codexbar)]
+    return _add_ccusage_enrichment(snapshot, ccusage)
+
+
+def _add_ccusage_enrichment(
+    snapshot: CapacitySnapshot,
+    ccusage: CapacitySnapshot | None = None,
+) -> CapacitySnapshot:
+    if ccusage is None:
+        return snapshot
+    enriched = snapshot.model_copy(deep=True)
+    enriched.raw.setdefault("alternate_sources", [])
+    enriched.raw["alternate_sources"].append(_usage_source_summary(ccusage))
+    if ccusage.status in {UsageStatus.UNAVAILABLE, UsageStatus.UNAUTHENTICATED}:
+        return enriched
+    if ccusage.warnings:
+        enriched.warnings.extend(f"ccusage: {warning}" for warning in ccusage.warnings)
+    enriched.raw["ccusage"] = ccusage.raw
+    return enriched
+
+
+def _usage_window_key(window: UsageWindow) -> tuple[str, str, str | None]:
+    return (window.kind.value, window.name, window.reset_at.isoformat() if window.reset_at else None)
+
+
+def _usage_source_summary(snapshot: CapacitySnapshot) -> dict[str, Any]:
+    return {
+        "source": snapshot.raw.get("source"),
+        "status": snapshot.status.value if hasattr(snapshot.status, "value") else str(snapshot.status),
+        "confidence": snapshot.confidence.value if hasattr(snapshot.confidence, "value") else str(snapshot.confidence),
+        "windows": len(snapshot.windows),
+    }