742 lines
28 KiB
Python
742 lines
28 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
usage_collector.py — OpenCode Go 订阅用量采集脚本 (Hybrid v3.0)
|
|
================================================================
|
|
混合采集模式:
|
|
主模式 — HTTP + stored cookies (全部 4 账号同时采集)
|
|
辅助模式 — CDP (仅当前 Chrome 登录的账号,用于补充/验证)
|
|
|
|
采集流程:
|
|
1. 遍历 accounts.json 中所有配置了 workspace_id 的账号
|
|
2. 对每个账号:
|
|
a. 主模式: 读取 cookies/{key_id}.json → 构建 Cookie header →
|
|
HTTP GET https://opencode.ai/workspace/{ws_id}/go → 解析 SSR HTML
|
|
b. 辅助模式(仅对 Chrome 当前登录账号): CDP fetch(location.href) → 解析 SSR HTML
|
|
3. 合并结果到 gateway/temp/usage_stats.json
|
|
|
|
Cookie 提取流程 (一次性,每个账号登录后执行):
|
|
python extract_cookies.py key1 --workspace-id wrk_XXX
|
|
→ cookies/key1.json (含 httpOnly auth cookie)
|
|
|
|
用法:
|
|
python usage_collector.py # 一次性采集 (全部账号)
|
|
python usage_collector.py --daemon # 后台守护,每 5 分钟采集一次
|
|
python usage_collector.py --print # 采集 + 打印结果
|
|
python usage_collector.py --cdp-only # 仅 CDP 采集当前登录账号
|
|
|
|
相关文件:
|
|
- usage_monitor/accounts.json — 账号配置
|
|
- usage_monitor/cookies/{key_id}.json — 存储的 cookies (gitignored)
|
|
- usage_monitor/extract_cookies.py — Cookie 提取工具
|
|
- gateway/temp/usage_stats.json — 采集结果缓存
|
|
- xmpp_agent_core.py /usage endpoint 触发本脚本
|
|
- CDP proxy (cdp-proxy.mjs) localhost:3456
|
|
"""
|
|
import os
|
|
import sys
|
|
import re
|
|
import json
|
|
import time
|
|
import logging
|
|
import argparse
|
|
import urllib.request
|
|
import urllib.error
|
|
import urllib.parse
|
|
import http.cookiejar
|
|
from pathlib import Path
|
|
from datetime import datetime, timezone
|
|
|
|
# ── Path bootstrap (platform-agnostic) ──
|
|
_SCRIPT_DIR = Path(__file__).resolve().parent # gateway/scripts/
|
|
_GATEWAY_DIR = _SCRIPT_DIR.parent # gateway/
|
|
_PROJECT_DIR = _GATEWAY_DIR.parent # AgentsMeeting/
|
|
|
|
USAGE_MONITOR_DIR = _SCRIPT_DIR / "usage_monitor"
|
|
ACCOUNTS_FILE = USAGE_MONITOR_DIR / "accounts.json"
|
|
COOKIES_DIR = USAGE_MONITOR_DIR / "cookies"
|
|
TEMP_DIR = _GATEWAY_DIR / "temp"
|
|
OUTPUT_FILE = TEMP_DIR / "usage_stats.json"
|
|
LOG_FILE = _GATEWAY_DIR / "logs" / "usage_collector.log"
|
|
|
|
# ── Ensure log dir exists before FileHandler ──
|
|
LOG_FILE.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
# ── Logging setup ──
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s [usage] %(levelname)s: %(message)s',
|
|
handlers=[
|
|
logging.FileHandler(str(LOG_FILE), encoding='utf-8'),
|
|
logging.StreamHandler(sys.stdout)
|
|
]
|
|
)
|
|
log = logging.getLogger('usage_collector')
|
|
|
|
# ── Constants ──
|
|
DEFAULT_POLL_INTERVAL_SEC = 300 # 5 minutes
|
|
CDP_PROXY_URL = os.environ.get('CDP_PROXY_URL', 'http://localhost:3456')
|
|
CDP_EVAL_TIMEOUT = 20
|
|
CDP_TARGETS_TIMEOUT = 5
|
|
HTTP_TIMEOUT = 15
|
|
OPENCODE_BASE_URL = 'https://opencode.ai'
|
|
|
|
# ── SSR HTML parsing ──
|
|
#
|
|
# opencode.ai SSR hydration format:
|
|
# lite.subscription.get["wrk_XXX"]).p.v = $R[N] = {
|
|
# rollingUsage: $R[M] = {status:"ok", resetInSec:16602, usagePercent:37},
|
|
# weeklyUsage: $R[M] = {status:"ok", resetInSec:402661, usagePercent:15},
|
|
# monthlyUsage: $R[M] = {status:"ok", resetInSec:2674309, usagePercent:7}
|
|
# }
|
|
|
|
_METRIC_RE = re.compile(
|
|
r'(?P<metric>rollingUsage|weeklyUsage|monthlyUsage)'
|
|
r'\s*:\s*'
|
|
r'(?:\$R\[\d+\]\s*=\s*)?'
|
|
r'\{(?P<inner>[^}]*)\}',
|
|
re.DOTALL
|
|
)
|
|
|
|
_WORKSPACE_LITERAL_RE = re.compile(r'(?P<workspace>wrk_[A-Z0-9]{20,})')
|
|
|
|
|
|
def ensure_dirs():
|
|
"""Make sure all expected directories exist."""
|
|
USAGE_MONITOR_DIR.mkdir(parents=True, exist_ok=True)
|
|
COOKIES_DIR.mkdir(parents=True, exist_ok=True)
|
|
TEMP_DIR.mkdir(parents=True, exist_ok=True)
|
|
LOG_FILE.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
def load_accounts_config():
|
|
"""Load accounts.json."""
|
|
if not ACCOUNTS_FILE.exists():
|
|
return None
|
|
try:
|
|
with open(str(ACCOUNTS_FILE), 'r', encoding='utf-8-sig') as f:
|
|
return json.load(f)
|
|
except Exception as e:
|
|
log.error(f'Failed to read accounts.json: {e}')
|
|
return None
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# PRIMARY MODE: HTTP + Stored Cookies
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
def load_stored_cookies(key_id):
|
|
"""
|
|
Load cookies from cookies/{key_id}.json.
|
|
Returns list of cookie dicts or None if file doesn't exist.
|
|
"""
|
|
cookie_file = COOKIES_DIR / f'{key_id}.json'
|
|
if not cookie_file.exists():
|
|
return None
|
|
try:
|
|
with open(str(cookie_file), 'r', encoding='utf-8') as f:
|
|
data = json.load(f)
|
|
return data.get('cookies', [])
|
|
except Exception as e:
|
|
log.error(f'Failed to load cookies for {key_id}: {e}')
|
|
return None
|
|
|
|
|
|
def _cookie_dict_to_jar(cookies_list, workspace_id=None):
|
|
"""Convert stored cookie dicts → http.cookiejar.CookieJar."""
|
|
jar = http.cookiejar.CookieJar()
|
|
for c in cookies_list:
|
|
domain = c.get('domain', '')
|
|
try:
|
|
cookie = http.cookiejar.Cookie(
|
|
version=0, name=c.get('name', ''), value=c.get('value', ''),
|
|
port=None, port_specified=False,
|
|
domain=domain, domain_specified=bool(domain),
|
|
domain_initial_dot=domain.startswith('.'),
|
|
path=c.get('path', '/'), path_specified=True,
|
|
secure=c.get('secure', False), expires=c.get('expires'),
|
|
discard=False, comment=None, comment_url=None,
|
|
rest={'HttpOnly': None} if c.get('httpOnly') else {},
|
|
rfc2109=False,
|
|
)
|
|
jar.set_cookie(cookie)
|
|
except Exception as e:
|
|
log.warning(f'Failed to convert cookie {c.get("name","?")}: {e}')
|
|
return jar
|
|
|
|
|
|
def _jar_to_cookie_dicts(jar):
|
|
"""Convert CookieJar → list of cookie dicts for JSON storage."""
|
|
out = []
|
|
for c in jar:
|
|
out.append({
|
|
'name': c.name, 'value': c.value,
|
|
'domain': c.domain, 'path': c.path,
|
|
'secure': c.secure, 'expires': c.expires,
|
|
'httpOnly': 'HttpOnly' in (c._rest or {}),
|
|
})
|
|
return out
|
|
|
|
|
|
def _save_cookies_back(key_id, jar, workspace_id):
|
|
"""Save updated cookies from jar back to cookies/{key_id}.json."""
|
|
cookies_list = _jar_to_cookie_dicts(jar)
|
|
if not cookies_list:
|
|
return
|
|
cookie_file = COOKIES_DIR / f'{key_id}.json'
|
|
data = {
|
|
'key_id': key_id, 'workspace_id': workspace_id,
|
|
'extracted_at': datetime.now(timezone.utc).isoformat(timespec='seconds'),
|
|
'cookie_count': len(cookies_list), 'cookies': cookies_list,
|
|
}
|
|
try:
|
|
with open(str(cookie_file), 'w', encoding='utf-8') as f:
|
|
json.dump(data, f, ensure_ascii=False, indent=2)
|
|
log.info(f'[keepalive] Updated cookies for {key_id} ({len(cookies_list)} cookies)')
|
|
except Exception as e:
|
|
log.error(f'[keepalive] Failed to save cookies for {key_id}: {e}')
|
|
|
|
|
|
def fetch_ssr_via_http(workspace_id, cookies_list, key_id):
|
|
"""
|
|
HTTP GET with CookieJar (auto-captures Set-Cookie for keepalive).
|
|
Returns (html_text, error_message, updated_jar).
|
|
"""
|
|
url = f'{OPENCODE_BASE_URL}/workspace/{workspace_id}/go'
|
|
jar = _cookie_dict_to_jar(cookies_list, workspace_id)
|
|
opener = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(jar))
|
|
req = urllib.request.Request(url, method='GET')
|
|
req.add_header('User-Agent', 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36')
|
|
req.add_header('Accept', 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8')
|
|
|
|
try:
|
|
with opener.open(req, timeout=HTTP_TIMEOUT) as resp:
|
|
status = resp.getcode()
|
|
final_url = resp.geturl()
|
|
|
|
if 'auth.opencode.ai' in final_url or 'auth/authorize' in final_url:
|
|
return None, 'session expired (redirected to auth.opencode.ai)', jar
|
|
|
|
if status != 200:
|
|
return None, f'HTTP {status}', jar
|
|
|
|
html = resp.read().decode('utf-8', errors='replace')
|
|
return html, None, jar
|
|
except urllib.error.HTTPError as e:
|
|
return None, f'HTTP {e.code}: {e.reason}', jar
|
|
except urllib.error.URLError as e:
|
|
return None, f'URL error: {e.reason}', jar
|
|
except Exception as e:
|
|
return None, f'Unexpected: {e}', jar
|
|
|
|
|
|
def _cookies_changed(original, jar):
|
|
"""Check if jar cookies differ from original list (Set-Cookie was received)."""
|
|
jar_cookies = _jar_to_cookie_dicts(jar)
|
|
if len(jar_cookies) != len(original):
|
|
return True
|
|
for a, b in zip(jar_cookies, original):
|
|
if a.get('value') != b.get('value') or a.get('expires') != b.get('expires'):
|
|
return True
|
|
return False
|
|
|
|
|
|
def collect_via_http(key_id, account_cfg):
|
|
"""
|
|
Collect usage data for one account via HTTP + stored cookies.
|
|
Captures Set-Cookie for keepalive (saves refreshed cookies back to JSON).
|
|
"""
|
|
workspace_id = account_cfg.get('workspace_id', '').strip()
|
|
if not workspace_id:
|
|
return {
|
|
'key_id': key_id, 'label': account_cfg.get('label', key_id),
|
|
'workspace_id': None, 'error': 'no workspace_id configured',
|
|
'session_expired': False,
|
|
'last_update_iso': datetime.now(timezone.utc).isoformat(timespec='seconds'),
|
|
}
|
|
|
|
cookies = load_stored_cookies(key_id)
|
|
if not cookies:
|
|
return {
|
|
'key_id': key_id, 'label': account_cfg.get('label', key_id),
|
|
'workspace_id': workspace_id,
|
|
'error': 'no stored cookies — run extract_cookies.py after logging in',
|
|
'session_expired': False,
|
|
'last_update_iso': datetime.now(timezone.utc).isoformat(timespec='seconds'),
|
|
}
|
|
|
|
log.info(f'[HTTP] Fetching {workspace_id}/go for {key_id}...')
|
|
|
|
html, err, jar = fetch_ssr_via_http(workspace_id, cookies, key_id)
|
|
|
|
# Cookie keepalive: if server sent Set-Cookie, save refreshed cookies
|
|
if _cookies_changed(cookies, jar):
|
|
_save_cookies_back(key_id, jar, workspace_id)
|
|
|
|
if err:
|
|
session_expired = 'auth' in err.lower() or 'session' in err.lower()
|
|
log.warning(f'[HTTP] {key_id} failed: {err}')
|
|
return {
|
|
'key_id': key_id, 'label': account_cfg.get('label', key_id),
|
|
'workspace_id': workspace_id, 'error': err,
|
|
'session_expired': session_expired,
|
|
'last_update_iso': datetime.now(timezone.utc).isoformat(timespec='seconds'),
|
|
}
|
|
|
|
parsed = parse_usage_from_html(html, workspace_id)
|
|
if not parsed:
|
|
return {
|
|
'key_id': key_id, 'label': account_cfg.get('label', key_id),
|
|
'workspace_id': workspace_id,
|
|
'error': 'no usage metrics in SSR HTML',
|
|
'session_expired': False,
|
|
'last_update_iso': datetime.now(timezone.utc).isoformat(timespec='seconds'),
|
|
}
|
|
|
|
result = {
|
|
'key_id': key_id, 'label': account_cfg.get('label', key_id),
|
|
'workspace_id': parsed['workspace_id'],
|
|
'rolling': parsed['rolling'], 'weekly': parsed['weekly'],
|
|
'monthly': parsed['monthly'], 'session_expired': False,
|
|
'last_update_iso': datetime.now(timezone.utc).isoformat(timespec='seconds'),
|
|
'error': None,
|
|
}
|
|
|
|
log.info(f'[HTTP] SUCCESS {key_id}: '
|
|
f'rolling={result["rolling"]["usage_percent"]}% '
|
|
f'weekly={result["weekly"]["usage_percent"]}% '
|
|
f'monthly={result["monthly"]["usage_percent"]}%')
|
|
|
|
return result
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# BONUS MODE: CDP (currently logged-in account only)
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
def cdp_get_targets():
|
|
"""Query CDP proxy for Chrome targets."""
|
|
url = f'{CDP_PROXY_URL}/targets'
|
|
try:
|
|
req = urllib.request.Request(url, method='GET')
|
|
with urllib.request.urlopen(req, timeout=CDP_TARGETS_TIMEOUT) as resp:
|
|
data = json.loads(resp.read().decode('utf-8'))
|
|
if isinstance(data, list):
|
|
return data
|
|
if isinstance(data, dict):
|
|
return data.get('targets', data.get('result', []))
|
|
return []
|
|
except Exception as e:
|
|
log.error(f'CDP /targets error: {e}')
|
|
return []
|
|
|
|
|
|
def cdp_eval(target_id, js_code):
|
|
"""Execute JS in a Chrome tab via CDP proxy."""
|
|
url = f'{CDP_PROXY_URL}/eval?target={target_id}'
|
|
data = js_code.encode('utf-8')
|
|
try:
|
|
req = urllib.request.Request(url, data=data, method='POST')
|
|
with urllib.request.urlopen(req, timeout=CDP_EVAL_TIMEOUT) as resp:
|
|
return json.loads(resp.read().decode('utf-8'))
|
|
except Exception as e:
|
|
log.error(f'CDP /eval error: {e}')
|
|
return None
|
|
|
|
|
|
def cdp_navigate(target_id, url):
|
|
"""Navigate a Chrome tab to a URL via CDP proxy."""
|
|
full_url = f'{CDP_PROXY_URL}/navigate?target={target_id}&url={urllib.parse.quote(url, safe="")}'
|
|
try:
|
|
req = urllib.request.Request(full_url, method='GET')
|
|
with urllib.request.urlopen(req, timeout=10) as resp:
|
|
return json.loads(resp.read().decode('utf-8'))
|
|
except Exception as e:
|
|
log.error(f'CDP /navigate error: {e}')
|
|
return None
|
|
|
|
|
|
def find_opencode_tab():
|
|
"""Find an opencode.ai workspace tab in Chrome."""
|
|
targets = cdp_get_targets()
|
|
if not targets:
|
|
log.warning('No Chrome targets — CDP proxy may be down or Chrome not running')
|
|
return None
|
|
|
|
for t in targets:
|
|
url = t.get('url', '')
|
|
if t.get('type') == 'page' and 'opencode.ai/workspace/' in url:
|
|
return t
|
|
for t in targets:
|
|
url = t.get('url', '')
|
|
if t.get('type') == 'page' and 'opencode.ai' in url:
|
|
return t
|
|
for t in targets:
|
|
if t.get('type') == 'page':
|
|
return t
|
|
return None
|
|
|
|
|
|
def fetch_ssr_html_via_cdp(target_id, workspace_url=None):
|
|
"""Use CDP /eval to run fetch(location.href) inside Chrome tab."""
|
|
if workspace_url:
|
|
log.info(f'[CDP] Navigating to {workspace_url}')
|
|
cdp_navigate(target_id, workspace_url)
|
|
time.sleep(3)
|
|
|
|
js = """
|
|
(async () => {
|
|
try {
|
|
var r = await fetch(location.href, {credentials:'include'});
|
|
var html = await r.text();
|
|
var wsMatch = html.match(/wrk_[A-Z0-9]{20,}/);
|
|
var wsId = wsMatch ? wsMatch[0] : null;
|
|
return JSON.stringify({
|
|
ok: true,
|
|
status: r.status,
|
|
url: location.href,
|
|
html: html,
|
|
workspace_id: wsId
|
|
});
|
|
} catch(e) {
|
|
return JSON.stringify({ok: false, error: e.message});
|
|
}
|
|
})()
|
|
""".strip()
|
|
|
|
result = cdp_eval(target_id, js)
|
|
if not result:
|
|
return None, 'CDP /eval returned no result'
|
|
|
|
value = None
|
|
if isinstance(result, dict):
|
|
value = result.get('value') or result.get('result')
|
|
elif isinstance(result, str):
|
|
value = result
|
|
|
|
if not value:
|
|
return None, f'CDP /eval returned empty: {result}'
|
|
|
|
try:
|
|
parsed = json.loads(value)
|
|
except (json.JSONDecodeError, TypeError) as e:
|
|
return None, f'Failed to parse /eval return: {e}'
|
|
|
|
if not parsed.get('ok'):
|
|
return None, f'JS fetch failed: {parsed.get("error", "unknown")}'
|
|
|
|
html = parsed.get('html', '')
|
|
ws_id = parsed.get('workspace_id')
|
|
status = parsed.get('status', 0)
|
|
url = parsed.get('url', '')
|
|
|
|
log.info(f'[CDP] fetch: status={status} url={url[:120]} html_len={len(html)} ws_id={ws_id}')
|
|
|
|
if 'auth.opencode.ai' in url or 'auth/authorize' in url:
|
|
return None, 'session expired (CDP tab on auth page)'
|
|
if status != 200 or not html:
|
|
return None, f'CDP fetch status={status} html_len={len(html)}'
|
|
return ws_id, html
|
|
|
|
|
|
def collect_via_cdp(accounts_cfg):
|
|
"""
|
|
Collect usage for the currently logged-in Chrome account via CDP.
|
|
Returns (key_id, result_dict) or (None, error_dict).
|
|
"""
|
|
tab = find_opencode_tab()
|
|
if not tab:
|
|
return None, {
|
|
'error': 'no Chrome tab (CDP proxy down or Chrome not running)',
|
|
'session_expired': False,
|
|
}
|
|
|
|
target_id = tab.get('targetId') or tab.get('id') or ''
|
|
if not target_id:
|
|
return None, {'error': 'tab has no targetId', 'session_expired': False}
|
|
|
|
tab_url = tab.get('url', '')
|
|
|
|
# If on auth page, try navigating to a workspace URL
|
|
workspace_url = None
|
|
if 'auth.opencode.ai' in tab_url or 'auth/authorize' in tab_url:
|
|
if accounts_cfg:
|
|
for acc in accounts_cfg.get('accounts', []):
|
|
ws_id = acc.get('workspace_id', '').strip()
|
|
if ws_id:
|
|
workspace_url = f'{OPENCODE_BASE_URL}/workspace/{ws_id}/go'
|
|
break
|
|
if not workspace_url:
|
|
return None, {'error': 'session expired + no workspace URLs', 'session_expired': True}
|
|
elif 'opencode.ai/workspace/' in tab_url:
|
|
workspace_url = None # fetch(location.href) will work
|
|
|
|
ws_id, err_or_html = fetch_ssr_html_via_cdp(target_id, workspace_url)
|
|
if not ws_id:
|
|
error_msg = err_or_html or 'unknown fetch error'
|
|
return None, {
|
|
'error': error_msg,
|
|
'session_expired': 'session' in error_msg.lower() or 'auth' in error_msg.lower(),
|
|
}
|
|
|
|
html_text = err_or_html
|
|
parsed = parse_usage_from_html(html_text, ws_id)
|
|
if not parsed:
|
|
return ws_id, {
|
|
'workspace_id': ws_id,
|
|
'error': 'no usage metrics (no Go subscription?)',
|
|
'session_expired': False,
|
|
}
|
|
|
|
# Match to account
|
|
label = ws_id
|
|
key_id = ws_id
|
|
if accounts_cfg:
|
|
for acc in accounts_cfg.get('accounts', []):
|
|
if acc.get('workspace_id') == ws_id:
|
|
label = acc.get('label', ws_id)
|
|
key_id = acc.get('key_id', ws_id)
|
|
break
|
|
|
|
result = {
|
|
'key_id': key_id,
|
|
'label': label,
|
|
'workspace_id': parsed['workspace_id'],
|
|
'rolling': parsed['rolling'],
|
|
'weekly': parsed['weekly'],
|
|
'monthly': parsed['monthly'],
|
|
'session_expired': False,
|
|
'last_update_iso': datetime.now(timezone.utc).isoformat(timespec='seconds'),
|
|
'error': None,
|
|
}
|
|
|
|
log.info(f'[CDP] SUCCESS {key_id}: '
|
|
f'rolling={result["rolling"]["usage_percent"]}% '
|
|
f'weekly={result["weekly"]["usage_percent"]}% '
|
|
f'monthly={result["monthly"]["usage_percent"]}%')
|
|
|
|
return key_id, result
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# SSR Parsing (shared by both modes)
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
def parse_usage_from_html(html_text, workspace_id_hint):
|
|
"""
|
|
Parse SSR HTML for usage data using per-metric regex.
|
|
Returns dict with workspace_id, rolling, weekly, monthly or None.
|
|
"""
|
|
metrics = {}
|
|
for m in _METRIC_RE.finditer(html_text):
|
|
metric_name = m.group('metric')
|
|
if metric_name in metrics:
|
|
continue
|
|
inner = m.group('inner')
|
|
status_match = re.search(r'status\s*:\s*"([^"]*)"', inner)
|
|
pct_match = re.search(r'usagePercent\s*:\s*(\d+)', inner)
|
|
reset_match = re.search(r'resetInSec\s*:\s*(\d+)', inner)
|
|
metrics[metric_name] = {
|
|
'usage_percent': int(pct_match.group(1)) if pct_match else None,
|
|
'reset_in_sec': int(reset_match.group(1)) if reset_match else None,
|
|
'status': status_match.group(1) if status_match else None,
|
|
}
|
|
|
|
ws_match = _WORKSPACE_LITERAL_RE.search(html_text)
|
|
workspace_id = ws_match.group('workspace') if ws_match else workspace_id_hint
|
|
|
|
if not metrics:
|
|
return None
|
|
|
|
log.info(f'Parsed {workspace_id}: '
|
|
f'rolling={metrics.get("rollingUsage", {}).get("usage_percent")}% '
|
|
f'weekly={metrics.get("weeklyUsage", {}).get("usage_percent")}% '
|
|
f'monthly={metrics.get("monthlyUsage", {}).get("usage_percent")}%')
|
|
|
|
return {
|
|
'workspace_id': workspace_id,
|
|
'rolling': metrics.get('rollingUsage', {'usage_percent': None, 'reset_in_sec': None, 'status': None}),
|
|
'weekly': metrics.get('weeklyUsage', {'usage_percent': None, 'reset_in_sec': None, 'status': None}),
|
|
'monthly': metrics.get('monthlyUsage', {'usage_percent': None, 'reset_in_sec': None, 'status': None}),
|
|
}
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# Cache merge logic
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
def load_existing_cache():
|
|
"""Load existing usage_stats.json or return empty structure."""
|
|
if not OUTPUT_FILE.exists():
|
|
return {'accounts': {}}
|
|
try:
|
|
with open(str(OUTPUT_FILE), 'r', encoding='utf-8') as f:
|
|
data = json.load(f)
|
|
if 'accounts' in data and isinstance(data['accounts'], list):
|
|
accounts_dict = {}
|
|
for acc in data['accounts']:
|
|
kid = acc.get('key_id', acc.get('workspace_id', '?'))
|
|
accounts_dict[kid] = acc
|
|
data['accounts'] = accounts_dict
|
|
return data
|
|
except Exception:
|
|
return {'accounts': {}}
|
|
|
|
|
|
def merge_result_into_cache(cache, key_id, new_result):
|
|
"""Update cache with new result for one account."""
|
|
if 'accounts' not in cache or not isinstance(cache['accounts'], dict):
|
|
cache['accounts'] = {}
|
|
cache['accounts'][key_id] = new_result
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# Main collection (Hybrid)
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
def collect_all(cdp_only=False):
|
|
"""
|
|
Hybrid collection: HTTP for all accounts + CDP bonus for current login.
|
|
"""
|
|
ensure_dirs()
|
|
accounts_cfg = load_accounts_config()
|
|
if not accounts_cfg:
|
|
log.error('No accounts.json found')
|
|
return {'ok': False, 'error': 'no accounts.json', 'accounts': []}
|
|
|
|
cache = load_existing_cache()
|
|
collected_count = 0
|
|
cdp_key_id = None
|
|
|
|
# ── PRIMARY: HTTP collection for all accounts ──
|
|
if not cdp_only:
|
|
for acc_cfg in accounts_cfg.get('accounts', []):
|
|
key_id = acc_cfg.get('key_id', '')
|
|
if not key_id:
|
|
continue
|
|
|
|
result = collect_via_http(key_id, acc_cfg)
|
|
merge_result_into_cache(cache, key_id, result)
|
|
|
|
if result.get('error') is None:
|
|
collected_count += 1
|
|
elif result.get('session_expired'):
|
|
log.warning(f'{key_id} session expired — needs re-login + re-extract cookies')
|
|
|
|
# ── BONUS: CDP collection for currently logged-in account ──
|
|
try:
|
|
cdp_kid, cdp_result = collect_via_cdp(accounts_cfg)
|
|
if cdp_kid:
|
|
cdp_key_id = cdp_kid
|
|
# CDP result overrides HTTP for this account (more reliable — live cookies)
|
|
merge_result_into_cache(cache, cdp_kid, cdp_result)
|
|
if cdp_result.get('error') is None:
|
|
collected_count += 1
|
|
log.info(f'CDP override for {cdp_kid}')
|
|
except Exception as e:
|
|
log.warning(f'CDP bonus collection skipped: {e}')
|
|
|
|
# ── Build output ──
|
|
accounts_list = []
|
|
for acc_cfg in accounts_cfg.get('accounts', []):
|
|
kid = acc_cfg.get('key_id', acc_cfg.get('workspace_id', '?'))
|
|
if kid in cache.get('accounts', {}):
|
|
# Merge in account config fields
|
|
cached = cache['accounts'][kid]
|
|
cached['key_id'] = kid
|
|
cached['label'] = acc_cfg.get('label', kid)
|
|
cached['workspace_id'] = acc_cfg.get('workspace_id') or cached.get('workspace_id')
|
|
accounts_list.append(cached)
|
|
else:
|
|
accounts_list.append({
|
|
'key_id': kid,
|
|
'label': acc_cfg.get('label', kid),
|
|
'workspace_id': acc_cfg.get('workspace_id'),
|
|
'rolling': None,
|
|
'weekly': None,
|
|
'monthly': None,
|
|
'session_expired': False,
|
|
'last_update_iso': None,
|
|
'error': 'not yet collected',
|
|
})
|
|
|
|
output = {
|
|
'ok': collected_count > 0,
|
|
'last_refresh_iso': datetime.now(timezone.utc).isoformat(timespec='seconds'),
|
|
'account_count': len(accounts_list),
|
|
'collected_count': collected_count,
|
|
'cdp_override_key_id': cdp_key_id,
|
|
'accounts': accounts_list,
|
|
}
|
|
|
|
save_result(output)
|
|
|
|
ok = sum(1 for a in accounts_list if a.get('rolling') and a['rolling'].get('usage_percent') is not None)
|
|
err = sum(1 for a in accounts_list if a.get('error'))
|
|
log.info(f'Done: {len(accounts_list)} accounts, {ok} with data, {err} errored.')
|
|
|
|
return output
|
|
|
|
|
|
def save_result(data):
|
|
"""Persist to gateway/temp/usage_stats.json."""
|
|
ensure_dirs()
|
|
try:
|
|
with open(str(OUTPUT_FILE), 'w', encoding='utf-8') as f:
|
|
json.dump(data, f, ensure_ascii=False, indent=2)
|
|
log.info(f'Saved to {OUTPUT_FILE}')
|
|
except Exception as e:
|
|
log.error(f'Failed to save: {e}')
|
|
|
|
|
|
def load_result():
|
|
"""Return cached result or empty dict."""
|
|
if not OUTPUT_FILE.exists():
|
|
return {
|
|
'ok': False,
|
|
'error': 'no usage_stats.json yet',
|
|
'last_refresh_iso': None,
|
|
'accounts': []
|
|
}
|
|
try:
|
|
with open(str(OUTPUT_FILE), 'r', encoding='utf-8') as f:
|
|
return json.load(f)
|
|
except Exception as e:
|
|
return {'ok': False, 'error': str(e), 'accounts': []}
|
|
|
|
|
|
def run_once(cdp_only=False):
|
|
"""One-shot collection."""
|
|
return collect_all(cdp_only=cdp_only)
|
|
|
|
|
|
def run_daemon(interval=DEFAULT_POLL_INTERVAL_SEC):
|
|
"""Daemon mode."""
|
|
ensure_dirs()
|
|
log.info(f'Daemon started, interval={interval}s')
|
|
while True:
|
|
try:
|
|
collect_all()
|
|
except Exception as e:
|
|
log.exception(f'Daemon error: {e}')
|
|
time.sleep(interval)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='OpenCode Go usage collector (Hybrid v3.0)')
|
|
parser.add_argument('--daemon', action='store_true', help='Run as daemon')
|
|
parser.add_argument('--interval', type=int, default=DEFAULT_POLL_INTERVAL_SEC,
|
|
help=f'Polling interval in seconds (default: {DEFAULT_POLL_INTERVAL_SEC})')
|
|
parser.add_argument('--print', action='store_true', help='Print result JSON')
|
|
parser.add_argument('--cdp-only', action='store_true', help='Only collect via CDP (skip HTTP)')
|
|
args = parser.parse_args()
|
|
|
|
if args.daemon:
|
|
run_daemon(interval=args.interval)
|
|
else:
|
|
result = run_once(cdp_only=args.cdp_only)
|
|
if args.print:
|
|
print(json.dumps(result, ensure_ascii=False, indent=2))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main() |