@agentikos/omega-os 0.1.0 → 0.19.5

package/omega/Agentik_Engine/omega_engine/account.py

@@ -0,0 +1,502 @@
+"""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.
+
+    Resolution order, first match wins:
+      1. ``vault_read(<secret_ref>_TOKEN)`` — the canonical key the new
+         ``write_token`` uses. If age is the active backend the read is
+         transparently decrypted.
+      2. Legacy dotenv file at ``Agentik_Extra/etc/secrets/<secret_ref>.env``
+         with a ``CLAUDE_OAUTH_TOKEN=...`` line. Kept for backward compat
+         with installs that pre-date the encrypted vault.
+    """
+    # Local import to avoid an import cycle at module load time.
+    from omega_engine.vault import vault_read
+
+    canon = f"{secret_ref}_TOKEN"
+    try:
+        token = vault_read(omega_home, canon)
+    except ValueError:
+        token = None
+    if token:
+        return token.strip()
+
+    path = vault_path(omega_home, secret_ref)
+    if not path.exists():
+        raise FileNotFoundError(f"vault entry not found for ref {secret_ref!r}")
+    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 through the vault layer.
+
+    Uses ``omega_engine.vault.vault_write``: age-encrypted if the system has
+    ``age`` + ``age-keygen``, else mode-600 plaintext. Either way, the file
+    is owner-only on disk. Overwrites silently — rotating a token is the
+    expected use case. Returns the actual on-disk path written.
+    """
+    if not token or not token.strip():
+        raise ValueError("token must be a non-empty string")
+    from omega_engine.vault import vault_write
+    return vault_write(omega_home, f"{secret_ref}_TOKEN", token.strip())
+    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

package/omega/Agentik_Engine/omega_engine/agent_messages.py

@@ -0,0 +1,167 @@
+"""Worker ↔ Oracle messaging — the missing back-channel.
+
+Today a worker writes a single ``.done.json`` and that's it. If it's
+stuck mid-task and needs the oracle's input, it has no way to ask. Then
+either it guesses wrong (and the audit catches it later — wasted work)
+or it gives up early (status=pending — wasted dispatch).
+
+This primitive is a small JSON-lines file per task at
+``Agentik_Runtime/sessions/<task_id>/messages.jsonl``. The worker
+appends a question, the oracle reads, appends an answer, the worker
+polls and continues.
+
+The protocol is intentionally minimal — one JSON object per line:
+
+  ``{from, to, ts, type, text}``
+
+``type`` is one of:
+
+  * ``question``  — worker asks the oracle
+  * ``answer``    — oracle's reply
+  * ``proposal``  — worker proposes a course change
+  * ``decision``  — oracle says yes/no/pivot
+
+Polling is the operator's choice — the worker can ``read_messages``
+between Bash/Edit iterations. For long-running missions the oracle's
+session can sit in ``omega message wait <task_id>`` to surface
+questions as they arrive.
+"""
+from __future__ import annotations
+
+import json
+import os
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+VALID_TYPES = {"question", "answer", "proposal", "decision"}
+
+
+@dataclass
+class Message:
+    from_: str
+    to: str
+    type: str
+    text: str
+    ts: int = field(default_factory=lambda: int(time.time()))
+
+
+def _msg_path(omega_home: str | Path, task_id: str) -> Path:
+    home = Path(omega_home)
+    p = home / "Agentik_Runtime" / "sessions" / task_id / "messages.jsonl"
+    p.parent.mkdir(parents=True, exist_ok=True)
+    return p
+
+
+def send(
+    omega_home: str | Path,
+    task_id: str,
+    *,
+    from_: str,
+    to: str,
+    type: str,
+    text: str,
+) -> Message:
+    """Append one message to the per-task JSONL. Returns the persisted Message."""
+    if type not in VALID_TYPES:
+        raise ValueError(f"type must be one of {VALID_TYPES}, got {type!r}")
+    msg = Message(from_=from_, to=to, type=type, text=text)
+    p = _msg_path(omega_home, task_id)
+    with p.open("a") as fh:
+        fh.write(json.dumps({
+            "from": msg.from_, "to": msg.to,
+            "type": msg.type, "text": msg.text,
+            "ts": msg.ts,
+        }) + "\n")
+    return msg
+
+
+def read(
+    omega_home: str | Path, task_id: str,
+    *,
+    since_ts: int = 0,
+    for_recipient: str | None = None,
+) -> list[Message]:
+    """Read messages for a task. Returns oldest-first."""
+    p = _msg_path(omega_home, task_id)
+    if not p.exists():
+        return []
+    out: list[Message] = []
+    for line in p.read_text().splitlines():
+        if not line.strip():
+            continue
+        try:
+            data = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+        if int(data.get("ts", 0)) < since_ts:
+            continue
+        if for_recipient and data.get("to") != for_recipient:
+            continue
+        out.append(Message(
+            from_=str(data.get("from", "?")),
+            to=str(data.get("to", "?")),
+            type=str(data.get("type", "question")),
+            text=str(data.get("text", "")),
+            ts=int(data.get("ts", 0)),
+        ))
+    return out
+
+
+def wait_for_reply(
+    omega_home: str | Path,
+    task_id: str,
+    *,
+    for_recipient: str,
+    since_ts: int | None = None,
+    timeout_s: int = 600,
+    poll_interval_s: int = 5,
+) -> Message | None:
+    """Block until a message addressed to ``for_recipient`` arrives or timeout.
+
+    Use case: the worker asks a question, then waits for the oracle's
+    answer before continuing. Returns the new message, or None on
+    timeout.
+    """
+    since_ts = since_ts if since_ts is not None else int(time.time())
+    end = time.time() + timeout_s
+    while time.time() < end:
+        msgs = read(omega_home, task_id, since_ts=since_ts,
+                    for_recipient=for_recipient)
+        if msgs:
+            return msgs[0]
+        time.sleep(poll_interval_s)
+    return None
+
+
+def list_open_questions(
+    omega_home: str | Path,
+) -> list[tuple[str, Message]]:
+    """Scan every per-task messages.jsonl for unanswered questions.
+
+    A "question" is unanswered when no later message of type=answer
+    addresses the same conversation. Returns (task_id, question_msg)
+    pairs sorted oldest first.
+    """
+    home = Path(omega_home)
+    sess = home / "Agentik_Runtime" / "sessions"
+    if not sess.is_dir():
+        return []
+    open_qs: list[tuple[str, Message]] = []
+    for child in sess.iterdir():
+        if not child.is_dir():
+            continue
+        p = child / "messages.jsonl"
+        if not p.exists():
+            continue
+        msgs = read(home, child.name)
+        # Questions whose ts is later than the latest answer in the same file
+        latest_answer = max(
+            (m.ts for m in msgs if m.type == "answer"), default=0,
+        )
+        for m in msgs:
+            if m.type == "question" and m.ts > latest_answer:
+                open_qs.append((child.name, m))
+    open_qs.sort(key=lambda pair: pair[1].ts)
+    return open_qs