opencode_usage_monitor.py

#!/usr/bin/env python3
"""Silent-first OpenCode Go quota monitor (stdlib only).

--keys-file explicitly selects a JSON array of {name, key}; otherwise use
OPENCODE_GO_API_KEY / OPENCODE_ZEN_API_KEY, then $HERMES_HOME/.env, ~/.hermes/.env.
--once never reads/writes monitor state. --json always emits every query result.
Percentages are quota percentages, NOT token counts or dollar amounts. The old
unverified dollar caps have been removed (including their environment override).
State/history base paths retain their environment overrides, but now always gain
full SHA256 key suffixes. Unattributed legacy state is deliberately not imported.
"""

import argparse
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timezone
import hashlib
import json
import math
import os
from pathlib import Path
import re
import sys
import tempfile
import urllib.error
import urllib.request

URL = 'https://opencode.ai/zen/go/v1/usage'
KEY_NAMES = ('OPENCODE_GO_API_KEY', 'OPENCODE_ZEN_API_KEY')
WINDOWS = ('rolling', 'weekly', 'monthly')
ALERT_PCT, WARN_PCT, DELTA_PP = 85, 70, 2
FAIL_COOLDOWN_SEC = 3600
UA = ('Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 '
      '(KHTML, like Gecko) Chrome/126 Safari/537.36')


def hermes_home():
    return os.path.expanduser(os.environ.get('HERMES_HOME') or '~/.hermes')


def resolve_paths():
    return tuple(os.path.expanduser(os.environ.get(env) or
                 os.path.join(hermes_home(), filename)) for env, filename in (
        ('OPENCODE_WATCHDOG_STATE', 'opencode-usage-state.json'),
        ('OPENCODE_WATCHDOG_HISTORY', 'opencode-usage-history.jsonl')))


def key_id(key):
    return hashlib.sha256(key.encode('utf-8')).hexdigest()


def paths_for_key(key):
    identity = key_id(key)
    return tuple(str(Path(p).with_name(f'{Path(p).stem}.{identity}{Path(p).suffix}'))
                 for p in resolve_paths())


def env_file_candidates():
    return list(dict.fromkeys([os.path.join(hermes_home(), '.env'),
                              os.path.expanduser('~/.hermes/.env')]))


def _key_from_file(path):
    try:
        lines = Path(path).read_text(encoding='utf-8').splitlines()
    except (OSError, UnicodeError):
        return None
    for name in KEY_NAMES:
        pattern = re.compile(r'\s*(?:export\s+)?' + name + r'\s*=\s*(.*)$')
        for line in lines:
            match = pattern.match(line)
            if not match:
                continue
            value = match[1].strip()
            if value.startswith(('"', "'")):
                value = value[1:].split(value[0], 1)[0]
            else:
                value = re.split(r'\s+#', value, maxsplit=1)[0].strip()
            if value:
                return value
    return None


def get_key():
    for name in KEY_NAMES:
        value = os.environ.get(name, '').strip()
        if value:
            return value
    for path in env_file_candidates():
        value = _key_from_file(path)
        if value:
            return value
    raise ValueError('未找到 API 密钥')


def load_keys(path=None):
    """Validate the entire explicit source before deduplication; first name wins."""
    try:
        entries = (json.loads(Path(path).expanduser().read_text(encoding='utf-8'))
                   if path is not None else [{'name': '默认', 'key': get_key()}])
        if not isinstance(entries, list) or not entries:
            raise ValueError
        for entry in entries:
            if not isinstance(entry, dict) or set(entry) != {'name', 'key'}:
                raise ValueError
            name, key = entry['name'], entry['key']
            if (not isinstance(name, str) or not re.fullmatch(r'[\w .-]{1,64}', name)
                    or not name.strip() or name != name.strip()
                    or not isinstance(key, str) or not key
                    or any(c.isspace() or not c.isprintable() for c in key)):
                raise ValueError
        secrets = [entry['key'] for entry in entries]
        if any(secret in entry['name'] for entry in entries for secret in secrets):
            raise ValueError
        unique = {}
        for entry in entries:
            unique.setdefault(entry['key'], entry)
        return list(unique.values())
    except (OSError, UnicodeError, ValueError, TypeError):
        raise ValueError('密钥配置无效或不可读；请检查来源、名称和密钥') from None


def validate_usage(data):
    """Allowlist output: never copy arbitrary server strings or extra fields."""
    try:
        usage = {}
        for window in WINDOWS:
            item = data['usage'][window]
            percent, status, reset = item['percent'], item['status'], item['resetsAt']
            if (type(percent) not in (int, float) or not math.isfinite(percent)
                    or percent < 0 or not isinstance(status, str)
                    or not isinstance(reset, str) or 'T' not in reset):
                raise ValueError
            dt = datetime.fromisoformat(reset.replace('Z', '+00:00'))
            if dt.tzinfo is None or dt.utcoffset() is None:
                raise ValueError
            usage[window] = {'percent': percent, 'remainingPercent': max(0, 100 - percent),
                             'status': 'ok' if status == 'ok' else 'error',
                             'resetsAt': dt.astimezone(timezone.utc).isoformat()}
        return {'usage': usage}
    except (KeyError, TypeError, ValueError, OverflowError):
        raise ValueError('API 数据格式无效') from None


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        raise ValueError('不允许 API 重定向')


def fetch(key=None):
    req = urllib.request.Request(URL, headers={
        'Authorization': 'Bearer ' + (key if key is not None else get_key()),
        'User-Agent': UA})
    opener = urllib.request.build_opener(NoRedirect())
    with opener.open(req, timeout=30) as response:
        body = response.read(1024 * 1024 + 1)
        if len(body) > 1024 * 1024:
            raise ValueError('API 响应过大')
        return validate_usage(json.loads(body.decode('utf-8')))


def load_state(path):
    try:
        data = json.loads(Path(path).read_text(encoding='utf-8'))
        return data if isinstance(data, dict) else {}
    except (OSError, ValueError):
        return {}


def save_state(state, path):
    target = Path(path)
    target.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp = tempfile.mkstemp(dir=target.parent, prefix=target.name + '.')
    try:
        with os.fdopen(fd, 'w', encoding='utf-8') as stream:
            json.dump(state, stream, ensure_ascii=False, allow_nan=False)
        os.replace(tmp, path)
    finally:
        if os.path.exists(tmp):
            os.unlink(tmp)


def build_report(data, change=None, alert=None, warn=None,
                 digest=False, first=False, forced=False):
    usage = data['usage']
    alerts = alert or [w for w in WINDOWS if usage[w]['percent'] >= ALERT_PCT
                      or usage[w]['status'] != 'ok']
    warns = warn or [w for w in WINDOWS if WARN_PCT <= usage[w]['percent'] < ALERT_PCT]
    title = ('⚠️ 额度告警' if alerts else '⚡ 额度预警' if warns else
             '📊 每日汇总' if digest else '📡 首次报告' if first else
             '📡 当前额度' if forced else '🔄 额度变化')
    lines = [title]
    for window, label in zip(WINDOWS, ('滚动窗口', '每周', '每月')):
        item = usage[window]
        flag = '（状态异常）' if item['status'] != 'ok' else ''
        delta = f"（{change[window]:+g} 个百分点）" if change and window in change else ''
        remaining = max(0, 100 - item['percent'])
        lines.append(f"  {label}：已用 {item['percent']:g}%，剩余 {remaining:g}%{flag}{delta}；重置 {item['resetsAt']}")
    return '\n'.join(lines)


def tick(data, now=None, state_path=None, history_path=None):
    now = now or datetime.now(timezone.utc)
    defaults = resolve_paths()
    state_path, history_path = state_path or defaults[0], history_path or defaults[1]
    usage = validate_usage(data)['usage']
    prev = load_state(state_path)
    last = prev.get('last', {})
    if not isinstance(last, dict):
        last = {}
    first = not all(type(last.get(w)) in (int, float) for w in WINDOWS)
    change = {} if first else {w: usage[w]['percent'] - last[w] for w in WINDOWS
                              if abs(usage[w]['percent'] - last[w]) >= DELTA_PP}
    resets = {w: usage[w]['resetsAt'] for w in WINDOWS}
    digest = prev.get('day') != now.strftime('%Y-%m-%d')
    important = any(usage[w]['percent'] >= WARN_PCT or usage[w]['status'] != 'ok'
                    for w in WINDOWS)
    out = (build_report({'usage': usage}, change=change, digest=digest and not first,
                        first=first) if first or change or digest or important else None)
    Path(history_path).parent.mkdir(parents=True, exist_ok=True)
    with open(history_path, 'a', encoding='utf-8') as stream:
        stream.write(json.dumps({'ts': now.isoformat(), 'usage': usage}) + '\n')
    # Compare cumulative movement against the last *report*, not the last poll.
    save_state({'last': {w: usage[w]['percent'] for w in WINDOWS} if out else last,
                'resets': resets, 'day': now.strftime('%Y-%m-%d')}, state_path)
    return out


def tick_failure(err=None, now=None, state_path=None):
    now = now or datetime.now(timezone.utc)
    state_path = state_path or resolve_paths()[0]
    state = load_state(state_path)
    last_failed = state.get('last_failed', 0)
    if isinstance(last_failed, (int, float)) and last_failed >= now.timestamp() - FAIL_COOLDOWN_SEC:
        return None
    save_state({**state, 'last_failed': now.timestamp()}, state_path)
    return '⚠️ 查询失败，请检查网络、密钥或 API 状态'


class SafeParser(argparse.ArgumentParser):
    def error(self, message):
        self.exit(2, '参数无效；请使用 --help 查看用法\n')


def parse_args(argv=None):
    parser = SafeParser(description='OpenCode Go 额度监控；默认无变化时静默')
    parser.add_argument('--once', action='store_true', help='打印结果，不读取或写入状态/历史')
    parser.add_argument('--json', action='store_true', help='输出所有查询的 JSON（监控模式仍保存状态）')
    parser.add_argument('--keys-file', help='显式 JSON 密钥数组文件；不回退到环境变量')
    parser.add_argument('--workers', type=int, choices=range(1, 17), default=4,
                        metavar='1..16', help='并发数（默认 4）')
    parser.add_argument('--selftest', action='store_true', help='运行离线单元测试')
    return parser.parse_args(argv)


def query(entry):
    row = {'name': entry['name'], 'id': key_id(entry['key']),
           'timestamp': datetime.now(timezone.utc).isoformat(), 'success': False, 'error': None}
    try:
        row['usage'] = validate_usage(fetch(entry['key']))['usage']
        row['success'] = True
    except urllib.error.HTTPError as error:
        code = error.code if type(error.code) is int and 100 <= error.code <= 599 else 0
        row['error'] = f'HTTP {code} 查询失败'
    except Exception:
        row['error'] = '查询失败（网络、鉴权或 API 数据无效）'
    return row


def main(argv=None):
    args = parse_args(argv)
    if args.selftest:
        import unittest
        suite = unittest.defaultTestLoader.discover(str(Path(__file__).parent),
                                                   pattern='test_opencode_usage_monitor.py')
        if not suite.countTestCases():
            print('未找到离线测试文件', file=sys.stderr)
            return 1
        return 0 if unittest.TextTestRunner(verbosity=2).run(suite).wasSuccessful() else 1
    try:
        entries = load_keys(args.keys_file)
    except ValueError:
        message = '密钥配置无效或不可读；请检查来源、名称和密钥'
        if args.json:
            print(json.dumps({'count': 0, 'succeeded': 0, 'failed': 0, 'rows': [],
                              'timestamp': datetime.now(timezone.utc).isoformat(),
                              'error': message}, ensure_ascii=False))
        else:
            print(message, file=sys.stderr)
        return 1
    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        rows = list(pool.map(query, entries))
    for entry, row in zip(entries, rows):
        report = None
        try:
            if args.once:
                report = build_report(row, forced=True) if row['success'] else row['error']
            else:
                state, history = paths_for_key(entry['key'])
                report = (tick(row, state_path=state, history_path=history) if row['success']
                          else tick_failure(state_path=state))
        except Exception:
            row['success'], row['error'] = False, '本地状态或历史读写失败'
            report = row['error']
        if report and not args.json:
            print(f"[{row['name']}] {report}")
    failed = sum(not row['success'] for row in rows)
    if args.json:
        print(json.dumps({'timestamp': datetime.now(timezone.utc).isoformat(),
                          'count': len(rows), 'succeeded': len(rows) - failed,
                          'failed': failed, 'rows': rows}, ensure_ascii=False, allow_nan=False))
    return 2 if failed else 0


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