box/bin/invite.py

#!/usr/bin/env python3
"""invite.py — Muse.ai invite codes + usage via the agent browsers.

Find: GET /api/hatch/invite via in-page fetch (primary; proven live),
      Invite-button popover parse (DOM fallback).
Redeem: POST /api/hatch/invite-code/redeem via in-page fetch.
Usage: Settings > General DOM read (weekly reset, % used, additional
      tokens) through the hatch_menu tree (dialog/mouse/tabs);
      settings-nav primitives live there, this module keeps the
      invite API + popover flows and the flat CLI parse.

Server eligibility (observed live): one redemption per account
(already_redeemed); 48h window from joining (window_expired); codes
carry limited uses (used_up) and can be revoked. Inviter rewards
accrue regardless of the inviter's own redemption state.

Caller errors (unknown node, malformed code) raise InviteError before
any CDP traffic. Transport/CDP failures return {"ok": False, ...}.
"""

import re
import time

from approvals import VALID_NODES, cdp_evaluate, get_cdp_ws
from hatch_menu import dialog as menu_dialog
from hatch_menu.mouse import close as _close
from hatch_menu.mouse import escape as _escape
from hatch_menu.tabs import general as general_tab


class InviteError(ValueError):
    """Caller error: unknown node or malformed code. Raised before CDP."""


CODE_RE = re.compile(r"^[A-Z0-9]{6}$")
CODE_REVEALED_RE = re.compile(r"Invite code revealed:\s*([A-Z0-9]{6})")

JS_INVITE_GET = """(async () => {
  try {
    const r = await fetch('/api/hatch/invite',
      {method: 'GET', cache: 'no-store'});
    return {status: r.status, data: await r.json()};
  } catch (e) { return {error: String(e).slice(0, 200)}; }
})()"""

# %s is a validated [A-Z0-9]{6} code: injection-safe by construction.
JS_REDEEM_TMPL = """(async () => {
  try {
    const r = await fetch('/api/hatch/invite-code/redeem', {method: 'POST',
      headers: {'Content-Type': 'application/json'},
      body: JSON.stringify({code: '%s', supportsRedemptionStatus: true})});
    return {status: r.status, ok: r.ok, data: await r.json()};
  } catch (e) { return {error: String(e).slice(0, 200)}; }
})()"""

JS_INVITE_CLICK = ("(() => { const b = document.querySelector("
                   "'[data-testid=\"hatch-invite-friends-button\"]');"
                   " if (!b) return 'NO_BUTTON'; b.click();"
                   " return 'CLICKED'; })()")

JS_POPOVER_TEXT = ("(() => { const d = document.querySelector("
                   "'[data-slot=\"popover-content\"]');"
                   " return d ? d.innerText : null; })()")

def normalize_code(code):
    """Uppercase/strip a code; None unless 6-char A-Z0-9."""
    if not isinstance(code, str):
        return None
    c = code.strip().upper()
    return c if CODE_RE.fullmatch(c) else None


def _check_node(node):
    if node not in VALID_NODES:
        raise InviteError("unknown node: %r (valid: %s)"
                          % (node, ", ".join(VALID_NODES)))


def open_settings(ws):
    """Open the Settings dialog via the dock menu. True when open.

    Public shim over hatch_menu.dialog.open_settings (single copy).
    """
    return menu_dialog.open_settings(ws)


def click_settings_tab(ws, name):
    """Open a Settings dialog tab by visible name. True when open.

    Public shim over hatch_menu.dialog.goto_tab (single copy).
    """
    return menu_dialog.goto_tab(ws, name)


def _popover_code(ws):
    """Invite code via the main-chat popover (DOM fallback). None if absent."""
    try:
        if cdp_evaluate(ws, JS_INVITE_CLICK, timeout=5.0) != "CLICKED":
            return None
        time.sleep(1.0)
        text = cdp_evaluate(ws, JS_POPOVER_TEXT, timeout=5.0) or ""
    except Exception:
        return None
    finally:
        _escape(ws)
    m = CODE_REVEALED_RE.search(text)
    return m.group(1) if m else None


def get_invite(node, timeout=8.0):
    """Per-agent invite state. API first, popover fallback for the code."""
    _check_node(node)
    try:
        ws, _ = get_cdp_ws(node)
    except Exception as e:
        return {"ok": False, "node": node,
                "error": "CDP unreachable: %s" % e}
    try:
        try:
            res = cdp_evaluate(ws, JS_INVITE_GET, await_promise=True,
                               timeout=timeout)
        except Exception as e:
            res = {"error": "%s: %s" % (type(e).__name__, e)}
        if isinstance(res, dict) and res.get("status") == 200 \
                and isinstance(res.get("data"), dict) \
                and res["data"].get("code"):
            d = res["data"]
            return {"ok": True, "node": node, "source": "api",
                    "code": d.get("code"),
                    "has_redeemed": d.get("has_redeemed_invite_code"),
                    "uses_remaining": d.get("uses_remaining"),
                    "use_count": d.get("use_count"),
                    "reward": d.get("reward")}
        code = _popover_code(ws)
        if code:
            return {"ok": True, "node": node, "source": "dom",
                    "code": code, "has_redeemed": None,
                    "uses_remaining": None, "use_count": None,
                    "reward": None}
        if isinstance(res, dict):
            detail = res.get("error", "invite API failed")
        else:
            detail = "invite API failed"
        return {"ok": False, "node": node,
                "error": "invite lookup failed (%s); popover has no code"
                % detail}
    finally:
        _close(ws)


def redeem_invite(node, code, timeout=12.0):
    """Redeem CODE on node. Returns ok / server reason + detail."""
    c = normalize_code(code)
    if c is None:
        raise InviteError("malformed invite code: %r (want 6 chars A-Z0-9)"
                          % (code,))
    _check_node(node)
    js = JS_REDEEM_TMPL % c
    try:
        ws, _ = get_cdp_ws(node)
    except Exception as e:
        return {"ok": False, "node": node, "code": c,
                "error": "CDP unreachable: %s" % e}
    try:
        try:
            res = cdp_evaluate(ws, js, await_promise=True, timeout=timeout)
        except Exception as e:
            return {"ok": False, "node": node, "code": c,
                    "error": "%s: %s" % (type(e).__name__, e)}
        if not isinstance(res, dict) or "data" not in res:
            if isinstance(res, dict):
                detail = res.get("error", "empty redeem response")
            else:
                detail = "empty redeem response"
            return {"ok": False, "node": node, "code": c,
                    "reason": "unknown", "detail": detail}
        data = res.get("data") or {}
        if res.get("ok") and data.get("success"):
            return {"ok": True, "node": node, "code": c,
                    "redemption_status": data.get("redemptionStatus"),
                    "detail": data.get("detail")}
        return {"ok": False, "node": node, "code": c,
                "reason": data.get("reason", "unknown"),
                "detail": data.get("detail")}
    finally:
        _close(ws)


def parse_usage_text(text):
    """Flat usage parse (CLI contract) over the tree's General parser."""
    nested = general_tab.parse_usage_text(text)
    wk = (nested.get("weekly") if nested else None) or {}
    add = (nested.get("additional") if nested else None) or {}
    tokens = add.get("tokens_left")
    return {
        "weekly_reset": wk.get("resets_on"),
        "weekly_used_pct": wk.get("pct_used"),
        "additional_expires": "Never expires"
        if add.get("never_expires") else None,
        "additional_used_pct": add.get("pct_used"),
        "additional_left": ("%s tokens left" % tokens)
        if tokens else None,
        "has_redeemed": bool(add.get("never_expires") or tokens
                             or "additional tokens" in (text or "").lower()),
    }


def get_usage(node, timeout=8.0):
    """Usage limits for one agent via Settings General (DOM read)."""
    _check_node(node)
    try:
        ws, _ = get_cdp_ws(node)
    except Exception as e:
        return {"ok": False, "node": node,
                "error": "CDP unreachable: %s" % e}
    try:
        if not open_settings(ws):
            return {"ok": False, "node": node,
                    "error": "settings dialog did not open"}
        if not menu_dialog.goto_tab(ws, "General"):
            return {"ok": False, "node": node,
                    "error": "General tab did not open"}
        # Note: Usage stats and redeem field can load asynchronously in the React/Radix tree.
        # Poll up to timeout seconds for usage or redeem markers.
        text = None
        deadline = time.time() + timeout
        while time.time() < deadline:
            t = menu_dialog.dialog_text(ws)
            if t and ("Weekly limit" in t or "Additional tokens" in t or "Redeem invite code" in t or "tokens left" in t or "% used" in t):
                text = t
                break
            time.sleep(0.4)
        if not text:
            text = menu_dialog.dialog_text(ws)

        if not text:
            return {"ok": False, "node": node,
                    "error": "empty settings dialog"}
        out = {"ok": True, "node": node}
        out.update(parse_usage_text(text))
        return out
    finally:
        _escape(ws)
        _close(ws)


def fleet_invite_status(nodes=None):
    """Invite state per node. Never raises; per-node error dicts."""
    out = {}
    for n in (nodes or list(VALID_NODES)):
        try:
            out[n] = get_invite(n)
        except InviteError as e:
            out[n] = {"ok": False, "node": n, "error": str(e)}
    return out


def fleet_usage(nodes=None):
    """Usage limits per node. Never raises; per-node error dicts."""
    out = {}
    for n in (nodes or list(VALID_NODES)):
        try:
            out[n] = get_usage(n)
        except InviteError as e:
            out[n] = {"ok": False, "node": n, "error": str(e)}
    return out