souveraine/tools/provider-balance.py

#!/usr/bin/env python3
"""Probe every configured inference provider and report which agent roles are
about to fail.

Why this exists
---------------
Between 2026-08-12 and 2026-08-13 three providers ran out of credit and in
every case the first thing that told us was a failed subconscious pass:

  * glm-5.2 (zai)      -- 429 code 1113 "insufficient balance"
  * deepseek           -- 402 "Insufficient Balance", total_balance -0.04
  * opencode zen       -- 401 CreditsError

Balance exhaustion is invisible until something dies.  Absence of complaint
reads as health.  This is the instrument that speaks first.

Method
------
It does NOT ask whether an endpoint is reachable.  Reachability is not an
answer -- GET /models on opencode returns 200 with a valid key and a bankrupt
workspace.  So the probe sends a real (tiny) completion and classifies the
reply.  Where a provider exposes a native balance endpoint we report the
actual number as well.

Exit status
-----------
  0  every provider backing a configured role can answer
  1  at least one role-backing provider cannot answer
  2  could not run (config unreadable)

Usage
-----
  provider-balance.py            # probe role-backing providers (default)
  provider-balance.py --all      # probe every configured provider
  provider-balance.py --json     # machine-readable
"""
from __future__ import annotations

import argparse
import json
import subprocess
import sys
import tomllib
from pathlib import Path

CONFIG = Path.home() / ".souveraine" / "config.toml"
TIMEOUT = 45

OK = "OK"
NO_BALANCE = "NO_BALANCE"
RATE_LIMITED = "RATE_LIMITED"
AUTH = "AUTH"
UNREACHABLE = "UNREACHABLE"
SKIPPED = "SKIPPED"

# Substrings that mean "your account is out of money", collected from the real
# error bodies of three different providers.  Matching on wording rather than
# status code because the codes disagree: zai says 429, deepseek 402,
# opencode 401 -- all for the same condition.
BROKE_MARKERS = (
    "insufficient balance",
    "creditserror",
    "out of credit",
    "quota exceeded",
    "billing",
    "payment required",
)

RATE_MARKERS = (
    "rate limit",
    "freeusagelimiterror",
    "too many requests",
)


def curl_json(url: str, key: str, payload: dict | None = None,
              virtual_key: str = "") -> tuple[int, str]:
    """POST/GET via curl.

    curl rather than urllib on purpose: opencode sits behind Cloudflare, which
    answers Python-urllib's user-agent with 403 code 1010.  That failure looks
    exactly like a rejected key and is not one.
    """
    cmd = ["curl", "-sS", "--max-time", str(TIMEOUT), "-w", "\n%{http_code}",
           "-H", f"Authorization: Bearer {key}"]
    if virtual_key:
        cmd += ["-H", f"x-bf-vk: {virtual_key}"]
    if payload is not None:
        cmd += ["-H", "Content-Type: application/json", "-d", json.dumps(payload)]
    cmd.append(url)
    try:
        out = subprocess.run(cmd, capture_output=True, text=True, timeout=TIMEOUT + 10)
    except subprocess.TimeoutExpired:
        return 0, "timed out"
    body = out.stdout
    if "\n" in body:
        body, _, code = body.rpartition("\n")
    else:
        code = "0"
    try:
        return int(code or 0), body
    except ValueError:
        return 0, body


def classify(code: int, body: str) -> str:
    low = body.lower()
    if code == 0:
        return UNREACHABLE
    if any(m in low for m in BROKE_MARKERS):
        return NO_BALANCE
    if any(m in low for m in RATE_MARKERS):
        return RATE_LIMITED
    if code == 200:
        return OK
    if code in (401, 403):
        return AUTH
    return UNREACHABLE


def native_balance(name: str, base_url: str, key: str) -> str | None:
    """Ask for an actual number where the provider exposes one."""
    if "api.deepseek.com" in base_url:
        code, body = curl_json("https://api.deepseek.com/user/balance", key)
        if code == 200:
            try:
                d = json.loads(body)
                info = (d.get("balance_infos") or [{}])[0]
                return f"{info.get('total_balance')} {info.get('currency', '')}".strip()
            except Exception:
                return None
    return None


def probe_provider(name: str, cfg: dict) -> dict:
    ptype = cfg.get("type", "")
    key = cfg.get("api_key", "")
    base = cfg.get("base_url", "").rstrip("/")
    model = cfg.get("primary_model", "")

    if ptype == "claude-subscription":
        files = cfg.get("credential_files") or [cfg.get("credential_file")]
        present = [f for f in files if f and Path(f).expanduser().exists()]
        status = OK if present else AUTH
        return {"provider": name, "status": status, "model": model,
                "detail": f"{len(present)}/{len(files)} credential files present"}

    if not key:
        return {"provider": name, "status": SKIPPED, "model": model,
                "detail": "no api_key configured"}

    payload = {"model": model,
               "messages": [{"role": "user", "content": "ping"}],
               "max_tokens": 1}
    code, body = curl_json(f"{base}/chat/completions", key, payload,
                           virtual_key=cfg.get("virtual_key", ""))
    status = classify(code, body)

    detail = f"HTTP {code}"
    bal = native_balance(name, base, key)
    if bal:
        detail += f" | balance {bal}"
    if status != OK:
        snippet = " ".join(body.split())[:160]
        if snippet:
            detail += f" | {snippet}"
    return {"provider": name, "status": status, "model": model, "detail": detail}


def roles(cfg: dict) -> dict[str, str]:
    """Which model each configured role depends on."""
    out = {}
    if m := cfg.get("subconscious", {}).get("model"):
        out["subconscious"] = m
    if m := cfg.get("reflection", {}).get("model"):
        out["reflection"] = m
    if m := cfg.get("archivist", {}).get("compression_model"):
        out["archivist"] = m
    return out


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--all", action="store_true",
                    help="probe every provider, not just role-backing ones")
    ap.add_argument("--json", action="store_true")
    args = ap.parse_args()

    try:
        with CONFIG.open("rb") as fh:
            cfg = tomllib.load(fh)
    except Exception as exc:
        print(f"cannot read {CONFIG}: {exc}", file=sys.stderr)
        return 2

    providers = cfg.get("providers", {})
    models = cfg.get("models", {})
    role_models = roles(cfg)

    # model alias -> provider name
    needed = {}
    for role, model in role_models.items():
        entry = models.get(model)
        if not entry:
            needed.setdefault("<unregistered:" + model + ">", []).append(role)
        else:
            needed.setdefault(entry.get("provider", "?"), []).append(role)

    targets = list(providers) if args.all else [p for p in needed if p in providers]
    results = [probe_provider(n, providers[n]) for n in targets]

    for r in results:
        r["roles"] = needed.get(r["provider"], [])

    unregistered = [k for k in needed if k.startswith("<unregistered:")]
    failing = [r for r in results if r["roles"] and r["status"] != OK]

    if args.json:
        print(json.dumps({"results": results, "unregistered": unregistered,
                          "failing": [r["provider"] for r in failing]}, indent=2))
    else:
        width = max((len(r["provider"]) for r in results), default=8)
        for r in results:
            tag = ",".join(r["roles"]) or "-"
            print(f"{r['provider']:<{width}}  {r['status']:<13} {tag:<28} {r['detail']}")
        for u in unregistered:
            print(f"{u}  NOT IN [models.*] -- role(s): {','.join(needed[u])}")
        if failing:
            print()
            for r in failing:
                print(f"!! {','.join(r['roles'])} will fail: "
                      f"{r['provider']} is {r['status']}")

    return 1 if (failing or unregistered) else 0


if __name__ == "__main__":
    sys.exit(main())