@agentikos/omega-os 0.1.0 → 0.2.0

package/omega/Agentik_Engine/omega_engine/account.py

@@ -0,0 +1,505 @@
+"""The Claude Code Max account pool + per-account billing.
+
+> Omega OS runs ONE engine process, not N tmux sessions. There is nothing to
+> "switch" globally. The Claude provider holds a POOL and assigns each agent
+> call to an account, so N Max accounts give the SUM of their rate limits as
+> usable throughput. See `docs/ACCOUNT-AND-BILLING.md` for the design.
+
+This module is the real implementation of that pool — plus the OAuth login
+flow, the encrypted-vault file format, and the billing aggregator that scans
+the event log for per-account token usage.
+
+Stdlib only (yaml is already a top-level dep of omega-engine).
+
+──── ON THE CLAUDE OAUTH DEVICE-CODE FLOW ────
+
+Anthropic does **not** publicly document an RFC 8628 device-authorization-grant
+endpoint for Claude Max accounts. The official Claude Code Max OAuth flow is
+browser-based PKCE through `https://claude.ai`, which a headless VPS cannot
+complete without launching a browser.
+
+`claude_device_code_flow` therefore implements RFC 8628 correctly against
+*configurable* endpoints (the spec is the spec — when Anthropic publishes the
+URLs, the code already works). If the endpoint is unreachable or returns a
+non-device-code response, the function raises with a clear message and the CLI
+falls back to a **manual paste flow**: the user logs in via browser on another
+machine, copies the OAuth token, pastes it into the prompt, and we write it
+into the vault with `chmod 600`. That manual fallback is the path that always
+works today — not a stub, a working credentials-into-the-vault flow.
+"""
+from __future__ import annotations
+
+import json
+import os
+import stat
+import time
+import urllib.error
+import urllib.parse
+import urllib.request
+from dataclasses import dataclass, field, asdict
+from pathlib import Path
+from typing import Any
+
+# yaml is a hard dependency of omega-engine — see pyproject.toml.
+import yaml
+
+
+__all__ = [
+    "ClaudeAccount",
+    "AccountPool",
+    "BillingAggregator",
+    "vault_path",
+    "read_token",
+    "write_token",
+    "claude_oauth_login_url",
+    "claude_device_code_flow",
+    "ACCOUNT_DEVICE_AUTH_URL",
+    "ACCOUNT_DEVICE_TOKEN_URL",
+]
+
+
+# ──── OAuth endpoint constants ────
+# Configurable via env so tests / future Anthropic endpoint changes do not
+# require a code edit. Defaults reflect the documented RFC 8628 shape; if
+# Anthropic publishes the actual URLs, set them with these env vars (or update
+# the constants here in one place).
+ACCOUNT_DEVICE_AUTH_URL = os.environ.get(
+    "OMEGA_CLAUDE_DEVICE_AUTH_URL",
+    "https://auth.anthropic.com/oauth/device/authorize",
+)
+ACCOUNT_DEVICE_TOKEN_URL = os.environ.get(
+    "OMEGA_CLAUDE_DEVICE_TOKEN_URL",
+    "https://auth.anthropic.com/oauth/device/token",
+)
+ACCOUNT_BROWSER_LOGIN_URL = "https://claude.ai/oauth"
+
+
+# ──── data model ────
+
+
+@dataclass
+class ClaudeAccount:
+    """One Claude Code Max account entry in the pool.
+
+    `secret_ref` is the only handle to the OAuth token — the token itself
+    lives in `Agentik_Extra/etc/secrets/<secret_ref>.env` with mode 600.
+    """
+
+    id: str
+    label: str = ""
+    secret_ref: str = ""
+    weight: int = 1
+    status: str = "active"           # active | resting | disabled
+    last_used_at: float = 0.0
+    tokens_used: int = 0
+
+    def to_yaml(self) -> dict[str, Any]:
+        """Serialize for accounts.yaml. Skips zero-valued runtime fields so the
+        on-disk file stays clean for accounts that have never been used."""
+        out: dict[str, Any] = {
+            "id": self.id,
+            "label": self.label,
+            "secret_ref": self.secret_ref,
+            "weight": int(self.weight),
+            "status": self.status,
+        }
+        if self.last_used_at:
+            out["last_used_at"] = float(self.last_used_at)
+        if self.tokens_used:
+            out["tokens_used"] = int(self.tokens_used)
+        return out
+
+    @classmethod
+    def from_yaml(cls, row: dict[str, Any]) -> "ClaudeAccount":
+        return cls(
+            id=str(row.get("id", "")).strip(),
+            label=str(row.get("label", "")),
+            secret_ref=str(row.get("secret_ref", "")),
+            weight=int(row.get("weight", 1)),
+            status=str(row.get("status", "active")),
+            last_used_at=float(row.get("last_used_at", 0.0) or 0.0),
+            tokens_used=int(row.get("tokens_used", 0) or 0),
+        )
+
+
+@dataclass
+class AccountPool:
+    """The pool of Claude Code Max accounts.
+
+    `selection` controls how `next()` picks the account for the next call:
+      - `round-robin` — strict rotation among active accounts
+      - `least-used` — pick the active account with the smallest `tokens_used`
+      - `by-quota`  — like least-used but weighted by `weight`
+    """
+
+    accounts: list[ClaudeAccount] = field(default_factory=list)
+    selection: str = "least-used"
+    version: int = 1
+    _rr_cursor: int = 0   # private; used by round-robin
+
+    # ──── loading & saving ────────────────────────────────────────────────
+
+    @classmethod
+    def load(cls, omega_home: str | Path) -> "AccountPool":
+        """Load the pool. Prefers `accounts.yaml`, falls back to
+        `accounts.example.yaml`. Returns an empty pool if neither exists."""
+        home = Path(omega_home)
+        cfg_dir = home / "Agentik_Providers" / "claude"
+        for name in ("accounts.yaml", "accounts.example.yaml"):
+            cfg = cfg_dir / name
+            if cfg.exists():
+                return cls._from_file(cfg)
+        return cls()
+
+    @classmethod
+    def _from_file(cls, path: Path) -> "AccountPool":
+        raw = yaml.safe_load(path.read_text()) or {}
+        accounts = [ClaudeAccount.from_yaml(r) for r in (raw.get("pool") or [])]
+        return cls(
+            accounts=accounts,
+            selection=str(raw.get("selection", "least-used")),
+            version=int(raw.get("version", 1) or 1),
+        )
+
+    def save(self, omega_home: str | Path) -> Path:
+        """Persist the pool to `accounts.yaml` (creates parent dirs).
+
+        Always writes to the canonical name (`accounts.yaml`), never to the
+        example file — that one is the template that ships with the repo.
+        """
+        home = Path(omega_home)
+        cfg_dir = home / "Agentik_Providers" / "claude"
+        cfg_dir.mkdir(parents=True, exist_ok=True)
+        cfg = cfg_dir / "accounts.yaml"
+        data = {
+            "version": self.version,
+            "selection": self.selection,
+            "pool": [a.to_yaml() for a in self.accounts],
+        }
+        # default_flow_style=False keeps it human-readable; sort_keys=False
+        # preserves the natural field order (id, label, secret_ref, ...).
+        cfg.write_text(yaml.safe_dump(data, sort_keys=False, default_flow_style=False))
+        return cfg
+
+    # ──── selection ───────────────────────────────────────────────────────
+
+    def _active(self) -> list[ClaudeAccount]:
+        return [a for a in self.accounts if a.status == "active"]
+
+    def next(self) -> ClaudeAccount:
+        """Pick the next account per the current `selection` strategy.
+
+        Raises RuntimeError if no account is active — that's a real error the
+        caller (the Claude provider) must surface, not silently mock out.
+        """
+        active = self._active()
+        if not active:
+            raise RuntimeError(
+                "no active Claude Max accounts in the pool — "
+                "run `omega account login` or `omega account use <id> active`"
+            )
+        if self.selection == "round-robin":
+            chosen = active[self._rr_cursor % len(active)]
+            self._rr_cursor = (self._rr_cursor + 1) % len(active)
+        elif self.selection == "by-quota":
+            # weighted least-used: score = tokens_used / max(weight, 1).
+            # smallest score wins. ties broken by last_used_at (older first).
+            chosen = min(
+                active,
+                key=lambda a: (a.tokens_used / max(a.weight, 1), a.last_used_at),
+            )
+        else:
+            # least-used (default): smallest tokens_used, then oldest last_used.
+            chosen = min(active, key=lambda a: (a.tokens_used, a.last_used_at))
+        chosen.last_used_at = time.time()
+        return chosen
+
+    # ──── mutation helpers ────────────────────────────────────────────────
+
+    def add(self, account: ClaudeAccount) -> None:
+        """Add or replace (by id) an account in the pool."""
+        if not account.id:
+            raise ValueError("account.id must be set")
+        existing = self.get(account.id)
+        if existing is None:
+            self.accounts.append(account)
+        else:
+            # replace in place so list ordering is preserved
+            idx = self.accounts.index(existing)
+            self.accounts[idx] = account
+
+    def get(self, account_id: str) -> ClaudeAccount | None:
+        for a in self.accounts:
+            if a.id == account_id:
+                return a
+        return None
+
+    def set_status(self, account_id: str, status: str) -> None:
+        if status not in ("active", "resting", "disabled"):
+            raise ValueError(
+                f"invalid status '{status}' — must be active|resting|disabled"
+            )
+        acc = self.get(account_id)
+        if acc is None:
+            raise KeyError(f"no account '{account_id}' in the pool")
+        acc.status = status
+
+    def usage_for(self, account_id: str, tokens: int) -> None:
+        """Record `tokens` usage on `account_id`. Idempotent at the pool level
+        (the event log is the source of truth — this just updates the cached
+        counter shown by `omega account list`)."""
+        acc = self.get(account_id)
+        if acc is None:
+            return
+        acc.tokens_used += max(int(tokens), 0)
+        acc.last_used_at = time.time()
+
+
+# ──── vault: where the OAuth tokens actually live ────────────────────────
+
+
+def vault_path(omega_home: str | Path, secret_ref: str) -> Path:
+    """Return the on-disk path for a secret reference.
+
+    Convention: `Agentik_Extra/etc/secrets/<secret_ref>.env`. The directory is
+    created on demand with mode 700; individual secret files are mode 600.
+
+    Refuses path-traversal — `secret_ref` cannot contain `/` or `..`.
+    """
+    if not secret_ref:
+        raise ValueError("secret_ref must not be empty")
+    if "/" in secret_ref or "\\" in secret_ref or ".." in secret_ref:
+        raise ValueError(f"secret_ref must be a flat name, got {secret_ref!r}")
+    home = Path(omega_home)
+    secrets_dir = home / "Agentik_Extra" / "etc" / "secrets"
+    secrets_dir.mkdir(parents=True, exist_ok=True)
+    try:
+        os.chmod(secrets_dir, 0o700)
+    except (OSError, PermissionError):
+        # best-effort — file mode is enforced below where it matters most
+        pass
+    return secrets_dir / f"{secret_ref}.env"
+
+
+def read_token(omega_home: str | Path, secret_ref: str) -> str:
+    """Read the OAuth token from the vault file.
+
+    Format: a tiny dotenv-style file with at least `CLAUDE_OAUTH_TOKEN=<value>`.
+    The reader is forgiving: it strips quotes, ignores blank lines and `#`
+    comments, and returns the FIRST value seen for `CLAUDE_OAUTH_TOKEN` (or the
+    canonical key formed by upper-casing the secret_ref).
+    """
+    path = vault_path(omega_home, secret_ref)
+    if not path.exists():
+        raise FileNotFoundError(f"vault file not found: {path}")
+    keys = ("CLAUDE_OAUTH_TOKEN", secret_ref.upper())
+    for line in path.read_text().splitlines():
+        s = line.strip()
+        if not s or s.startswith("#") or "=" not in s:
+            continue
+        k, _, v = s.partition("=")
+        k = k.strip()
+        v = v.strip().strip('"').strip("'")
+        if k in keys and v:
+            return v
+    raise ValueError(
+        f"no CLAUDE_OAUTH_TOKEN in vault file {path} — "
+        "expected a line like `CLAUDE_OAUTH_TOKEN=...`"
+    )
+
+
+def write_token(omega_home: str | Path, secret_ref: str, token: str) -> Path:
+    """Write the OAuth token to the vault file with mode 600.
+
+    Overwrites silently if a file already exists — rotating a token is the
+    expected use case. Returns the path written.
+    """
+    if not token or not token.strip():
+        raise ValueError("token must be a non-empty string")
+    path = vault_path(omega_home, secret_ref)
+    # Use os.open with O_CREAT|O_WRONLY|O_TRUNC and explicit mode so the file
+    # is born with 0o600 (no race window during which 0o644 would be visible).
+    fd = os.open(str(path), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+    try:
+        with os.fdopen(fd, "w") as fh:
+            fh.write(f"# Omega OS — Claude Max OAuth token for {secret_ref}\n")
+            fh.write(f"# Written: {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}\n")
+            fh.write(f"CLAUDE_OAUTH_TOKEN={token.strip()}\n")
+    finally:
+        # If fdopen succeeded, fd is closed by the context manager; if it
+        # raised, the fd is leaked — guard that by attempting a close.
+        try:
+            os.close(fd)
+        except OSError:
+            pass
+    # Belt-and-braces: re-assert mode in case some umask oddity sneaked in.
+    try:
+        os.chmod(path, stat.S_IRUSR | stat.S_IWUSR)  # 0o600
+    except OSError:
+        pass
+    return path
+
+
+# ──── OAuth login flow ───────────────────────────────────────────────────
+
+
+def claude_oauth_login_url() -> str:
+    """The URL a human visits to obtain a Claude OAuth token.
+
+    Returned as a plain string so the CLI can print it and instruct the user.
+    """
+    return ACCOUNT_BROWSER_LOGIN_URL
+
+
+def claude_device_code_flow(
+    client_id: str = "claude-code-cli",
+    scope: str = "claude:max",
+    device_auth_url: str | None = None,
+    token_url: str | None = None,
+    poll_timeout: float = 600.0,
+) -> dict[str, Any]:
+    """RFC 8628 device authorization flow against the Claude OAuth endpoints.
+
+    Returns a dict with at least `access_token`. If the endpoint does not
+    support this flow (404 / 4xx / non-JSON), raises RuntimeError with a clear
+    message so the CLI knows to fall back to manual paste.
+
+    This function performs ZERO disk I/O — the CLI is responsible for writing
+    the returned token into the vault. Keeps the function pure and testable.
+    """
+    auth_url = device_auth_url or ACCOUNT_DEVICE_AUTH_URL
+    tok_url = token_url or ACCOUNT_DEVICE_TOKEN_URL
+
+    # ---- step 1: request a device_code ----
+    body = urllib.parse.urlencode({"client_id": client_id, "scope": scope}).encode()
+    req = urllib.request.Request(
+        auth_url, data=body, method="POST",
+        headers={"Content-Type": "application/x-www-form-urlencoded",
+                 "Accept": "application/json"},
+    )
+    try:
+        with urllib.request.urlopen(req, timeout=30.0) as resp:
+            payload = json.loads(resp.read().decode("utf-8"))
+    except urllib.error.HTTPError as exc:
+        raise RuntimeError(
+            f"device authorize endpoint returned HTTP {exc.code}. "
+            "Anthropic may not support the device-code flow yet — "
+            "fall back to manual paste."
+        ) from exc
+    except (urllib.error.URLError, json.JSONDecodeError, TimeoutError) as exc:
+        raise RuntimeError(
+            f"could not reach the device authorize endpoint: {exc}. "
+            "Fall back to manual paste."
+        ) from exc
+
+    device_code = payload.get("device_code")
+    interval = float(payload.get("interval", 5))
+    expires_in = float(payload.get("expires_in", 600))
+    if not device_code:
+        raise RuntimeError(
+            f"device authorize response missing `device_code`: {payload!r}. "
+            "Fall back to manual paste."
+        )
+    verification = (
+        payload.get("verification_uri_complete")
+        or payload.get("verification_uri")
+        or "(see provider docs)"
+    )
+
+    # The caller (CLI) prints this — but in case it's invoked headlessly, surface
+    # it on the returned payload too.
+    payload["_human_instruction"] = (
+        f"open in a browser and approve:\n   {verification}"
+    )
+
+    # ---- step 2: poll for the token ----
+    deadline = min(time.monotonic() + poll_timeout, time.monotonic() + expires_in)
+    poll_body = urllib.parse.urlencode({
+        "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+        "device_code": device_code,
+        "client_id": client_id,
+    }).encode()
+    while time.monotonic() < deadline:
+        time.sleep(max(interval, 1.0))
+        poll_req = urllib.request.Request(
+            tok_url, data=poll_body, method="POST",
+            headers={"Content-Type": "application/x-www-form-urlencoded",
+                     "Accept": "application/json"},
+        )
+        try:
+            with urllib.request.urlopen(poll_req, timeout=30.0) as resp:
+                tok = json.loads(resp.read().decode("utf-8"))
+        except urllib.error.HTTPError as exc:
+            try:
+                err_payload = json.loads(exc.read().decode("utf-8"))
+            except Exception:  # noqa: BLE001 — best effort on body decode
+                err_payload = {}
+            err = err_payload.get("error", "")
+            # RFC 8628 §3.5: authorization_pending and slow_down mean keep polling.
+            if err == "authorization_pending":
+                continue
+            if err == "slow_down":
+                interval += 5
+                continue
+            raise RuntimeError(
+                f"device token endpoint refused: {err or exc.code} — "
+                f"{err_payload!r}"
+            ) from exc
+        except (urllib.error.URLError, json.JSONDecodeError, TimeoutError) as exc:
+            raise RuntimeError(f"could not poll token endpoint: {exc}") from exc
+
+        if tok.get("access_token"):
+            return {
+                "access_token": tok["access_token"],
+                "refresh_token": tok.get("refresh_token", ""),
+                "expires_in": int(tok.get("expires_in", 0) or 0),
+                "raw": tok,
+            }
+
+    raise RuntimeError(
+        "device-code flow timed out — user did not approve in time"
+    )
+
+
+# ──── billing aggregator ─────────────────────────────────────────────────
+
+
+class BillingAggregator:
+    """Aggregate per-account token usage from the event log.
+
+    Scans every `task.*` event whose payload includes a `usage.account_id` and
+    a positive `usage.input_tokens` / `usage.output_tokens`. The event log is
+    the source of truth — `AccountPool.tokens_used` is a cached view.
+    """
+
+    @staticmethod
+    def from_event_log(store: Any) -> dict[str, dict[str, int]]:
+        """Walk the store and aggregate. `store` is anything that yields
+        Event-like objects via `.all_events()` (see `omega_engine.store`).
+
+        Returns: { account_id: { input_tokens, output_tokens, total_calls } }
+        """
+        totals: dict[str, dict[str, int]] = {}
+        for ev in store.all_events():
+            # only billable events
+            etype = getattr(ev, "type", None)
+            type_val = etype.value if hasattr(etype, "value") else str(etype)
+            if not type_val.startswith("task."):
+                continue
+            payload = getattr(ev, "payload", None) or {}
+            usage = payload.get("usage") or {}
+            account_id = usage.get("account_id")
+            if not account_id:
+                continue
+            inp = int(usage.get("input_tokens", 0) or 0)
+            out = int(usage.get("output_tokens", 0) or 0)
+            if inp <= 0 and out <= 0:
+                continue
+            slot = totals.setdefault(
+                account_id,
+                {"input_tokens": 0, "output_tokens": 0, "total_calls": 0},
+            )
+            slot["input_tokens"] += inp
+            slot["output_tokens"] += out
+            slot["total_calls"] += 1
+        return totals