eink-push/codex_usage.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Codex 账户级用量 —— 通过官方 `codex app-server` JSON-RPC 读 ChatGPT 订阅的 Codex 额度（5 小时 / 7 天窗口）。

背景（2026-09 调研结论）：ChatGPT 订阅的**普通聊天**额度没有任何可编程的账户级读取途径（私有端点在
Sentinel/Cloudflare 门后、且只在被限流时才返回数值）；可监控的账户级实时余额只有 **Codex 额度池**——
Codex CLI / IDE 扩展 / Codex Web / ChatGPT 桌面 App 内的 Codex 共用同一池。本模块读的就是它。

为什么走 app-server RPC 而不是直调 chatgpt.com/backend-api/wham/usage：token 刷新由 Codex CLI 代管
（refresh token 轮换、外部进程自刷会踢掉 CLI 登录），官方 CI 文档也明示「不要自己调刷新接口」。
协议（已在本机 codex-cli 0.154 验证）：stdio、按行 JSON-RPC——
  → {"id":1,"method":"initialize","params":{"clientInfo":{...}}}   ← {"id":1,"result":{...}}
  → {"method":"initialized"}
  → {"id":2,"method":"account/rateLimits/read","params":{"excludeResetCreditDetails":true}}
  ← {"id":2,"result":{"rateLimits":{"primary":{usedPercent,windowDurationMins,resetsAt},"secondary":...,"planType":...},
                     "rateLimitsByLimitId":{...}}}
  API Key 登录时返回 error「chatgpt authentication required to read rate limits」。
登录：需 Codex 以 ChatGPT 账号 `codex login`（2026-09-14 用户在默认 ~/.codex 完成，故 CODEX_USAGE_HOME 默认留空跟随默认目录；
若日常 Codex 要保持 API Key/第三方中转登录，可设独立 CODEX_HOME 并在其下 `CODEX_HOME=… codex login`）。token 由 CLI 自管。
CODEX_BIN 由 config 解析为绝对路径（launchd 的 PATH 不含 /opt/homebrew/bin，曾因此报「找不到 codex 可执行文件」）。窗口归类按 windowDurationMins（Pro 计划常只回一个窗口、另一个为 null，
不能按 primary/secondary 位置判断）。输出快照与 Claude 同构，data.usage_from_raw 直接消费：
  {five_hour:{utilization, resets_at, window}, seven_day:{...}, plan, updated_at, source}
"""
import json
import os
import select
import subprocess
import time

import config
import usage_common

RPC_TIMEOUT = 25  # 秒：含 app-server 启动 + 一次远端查询


def _env():
    env = {k: v for k, v in os.environ.items()
           if k.upper() not in ("HTTPS_PROXY", "HTTP_PROXY", "ALL_PROXY", "NO_PROXY")}  # 直连，不走代理
    if config.CODEX_USAGE_HOME:
        home = os.path.expanduser(config.CODEX_USAGE_HOME)
        os.makedirs(home, exist_ok=True)  # codex 对不存在的 CODEX_HOME 直接退出；建空目录后它会给出清晰的「未登录」错误
        env["CODEX_HOME"] = home
    return env


def _rpc_read_rate_limits():
    """起一个 app-server 子进程，完成握手并读一次额度；返回 account/rateLimits/read 的 result。"""
    try:
        p = subprocess.Popen([config.CODEX_BIN, "app-server"], stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE, text=True, env=_env())
    except FileNotFoundError as e:
        raise RuntimeError(f"找不到 codex 可执行文件（CODEX_BIN={config.CODEX_BIN}）") from e
    deadline = time.time() + RPC_TIMEOUT

    def send(obj):
        p.stdin.write(json.dumps(obj) + "\n")
        p.stdin.flush()

    def wait_for(rid):
        while time.time() < deadline:
            r, _, _ = select.select([p.stdout], [], [], 0.2)
            if not r:
                if p.poll() is not None:
                    raise RuntimeError(f"app-server 提前退出：{p.stderr.read()[:300]}")
                continue
            line = p.stdout.readline()
            if not line:
                raise RuntimeError(f"app-server 关闭了输出：{p.stderr.read()[:300]}")
            try:
                msg = json.loads(line)
            except json.JSONDecodeError:
                continue
            if msg.get("id") != rid:
                continue  # 通知（remoteControl/status/changed 等）一律跳过
            if "error" in msg:
                err = msg["error"].get("message", str(msg["error"]))
                hint = ""
                if "authentication required" in err.lower():
                    hint = (f"；请先执行 {'CODEX_HOME=' + config.CODEX_USAGE_HOME + ' ' if config.CODEX_USAGE_HOME else ''}"
                            f"codex login 用 ChatGPT 账号登录")
                raise RuntimeError(f"RPC 错误：{err}{hint}")
            return msg.get("result") or {}
        raise RuntimeError(f"app-server 在 {RPC_TIMEOUT}s 内未响应 id={rid}")

    try:
        send({"id": 1, "method": "initialize",
              "params": {"clientInfo": {"name": "eink-push", "title": "eink-push usage", "version": "0.1"}}})
        wait_for(1)
        send({"method": "initialized"})
        send({"id": 2, "method": "account/rateLimits/read", "params": {"excludeResetCreditDetails": True}})
        return wait_for(2)
    finally:
        try:
            p.kill()
        except OSError:
            pass


def _classify(win, fallback_key):
    """按窗口长度归类：≤12 小时 → five_hour，≥5 天 → seven_day；缺 windowDurationMins 时按位置兜底。"""
    mins = win.get("windowDurationMins")
    if mins is None:
        return fallback_key
    if mins <= 12 * 60:
        return "five_hour"
    if mins >= 5 * 24 * 60:
        return "seven_day"
    return None


def normalize(result, now=None):
    now = now or time.time()
    rl = result.get("rateLimits") or {}
    out = {"five_hour": None, "seven_day": None, "plan": rl.get("planType"),
           "updated_at": int(now), "source": "codex-app-server"}
    for pos, fallback in (("primary", "five_hour"), ("secondary", "seven_day")):
        win = rl.get(pos)
        if not isinstance(win, dict) or win.get("usedPercent") is None:
            continue
        key = _classify(win, fallback)
        if key and out.get(key) is None:
            seg = {"utilization": float(win["usedPercent"]), "resets_at": win.get("resetsAt")}
            if win.get("windowDurationMins"):
                seg["window"] = int(win["windowDurationMins"]) * 60  # 秒，供 pace 用真实窗口长度
            out[key] = seg
    # 多桶视图原样附带（key 为 limit_id，如 codex），目前不渲染，留作观察真实数据后再决定是否加第三行
    by_id = result.get("rateLimitsByLimitId")
    if isinstance(by_id, dict) and by_id:
        out["buckets"] = {k: {"primary": (v or {}).get("primary"), "secondary": (v or {}).get("secondary"),
                              "limitName": (v or {}).get("limitName")} for k, v in by_id.items()}
    return out


def get_usage(force=False):
    """对外入口：带缓存/最小间隔/退避的 Codex 额度快照。未登录 ChatGPT 账号且无旧快照时抛出。"""
    snap, status = usage_common.cached_snapshot(
        config.CODEX_USAGE_CACHE_PATH, config.CODEX_USAGE_TTL, config.CODEX_USAGE_BACKOFF,
        lambda: normalize(_rpc_read_rate_limits()), force=force)
    snap = dict(snap)
    snap["status"] = status
    return snap


if __name__ == "__main__":
    import sys
    print(json.dumps(get_usage(force="--force" in sys.argv), ensure_ascii=False, indent=1))