From 0e03cd8eb749b4530d8aa2b4af317687ead35686 Mon Sep 17 00:00:00 2001 From: FelipeMayerDev Date: Wed, 19 Aug 2026 23:50:58 -0300 Subject: [PATCH 1/5] Add Z.ai GLM Coding Plan agent usage collector Adds omarchy-agent-usage-zai so the agents widget gains a Z.ai tab. Rate limits and plan tier come from Z.ai's coding-plan monitor endpoint (.../api/monitor/usage/quota/limit), mapping the rolling 5-hour and weekly TOKENS_LIMIT windows. Keys without a Coding Plan get "no coding plan" and the record falls back to local stats alone, the same way the claude collector degrades without a signed-in CLI. Local stats come from pi/omp sessions that ran through the Z.ai provider (glm-* models), mirroring how the codex collector reads pi/omp sessions. Credentials are discovered from ~/.config/omarchy/agents/zai.json, then ZAI_API_KEY/ZHIPU_API_KEY, then omp/pi .env files; both the global (api.z.ai) and China (open.bigmodel.cn) hosts are supported. --- bin/omarchy-agent-usage-zai | 400 ++++++++++++++++++++++++++++++++++++ 1 file changed, 400 insertions(+) create mode 100755 bin/omarchy-agent-usage-zai diff --git a/bin/omarchy-agent-usage-zai b/bin/omarchy-agent-usage-zai new file mode 100755 index 00000000000..0f82c489a9d --- /dev/null +++ b/bin/omarchy-agent-usage-zai @@ -0,0 +1,400 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Z.ai GLM Coding Plan usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Z.ai (GLM Coding Plan) usage into one display-ready JSON record. + +Rate limits and the plan tier come from Z.ai's coding-plan monitor endpoint +(…/api/monitor/usage/quota/limit), which reports the rolling 5-hour and weekly +token windows. That endpoint only answers for keys that carry a Coding Plan; a +generic pay-as-you-go key gets "no coding plan" and the record falls back to +local stats alone, exactly like the claude collector without a signed-in CLI. + +Local stats come from pi/omp sessions that ran through the Z.ai provider (the +`glm-*` models). The agents panel only ever reads the JSON this prints. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import urllib.error +import urllib.request +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Any + +AGENT_ID = "zai" +AGENT_NAME = "GLM Coding Plan" +AUTH_HELP_NO_KEY = ( + "Set ZAI_API_KEY (or add ~/.config/omarchy/agents/zai.json) to show Coding Plan limits." +) + +# Global (Z.ai) vs China (Zhipu / BigModel) split the same API under two hosts. +PLATFORMS = { + "zai": "https://api.z.ai", + "zhipu": "https://open.bigmodel.cn", +} +QUOTA_PATH = "/api/monitor/usage/quota/limit" + +# Coding-plan token windows are tagged by (unit, number); the monitor endpoint +# reuses TOKENS_LIMIT for both the 5-hour and the weekly window. +TOKEN_WINDOWS = { + (3, 5): "Session (5-hour)", + (6, 1): "Weekly (7-day)", +} + + +def number(value: Any) -> int: + try: + return max(0, int(value or 0)) + except (TypeError, ValueError): + return 0 + + +def local_day(value: Any) -> str: + if value is None: + return datetime.now().strftime("%Y-%m-%d") + if isinstance(value, (int, float)): + # pi/omp message timestamps are milliseconds. + if value > 10_000_000_000: + value = value / 1000 + return datetime.fromtimestamp(value).strftime("%Y-%m-%d") + text = str(value) + try: + if text.endswith("Z"): + dt = datetime.fromisoformat(text[:-1] + "+00:00") + else: + dt = datetime.fromisoformat(text) + if dt.tzinfo is not None: + dt = dt.astimezone() + return dt.strftime("%Y-%m-%d") + except ValueError: + return datetime.now().strftime("%Y-%m-%d") + + +def empty_bucket() -> dict[str, int]: + return { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + } + + +def empty_stats() -> dict[str, Any]: + return { + "todayPrompts": 0, + "todaySessions": 0, + "todayTotalTokens": 0, + "todayTokensByModel": {}, + "recentDays": [], + "totalPrompts": 0, + "totalSessions": 0, + "activeDays": 0, + "activeDates": [], + "modelUsage": {}, + } + + +def base_record(**overrides: Any) -> dict[str, Any]: + record: dict[str, Any] = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": False, + "hasLocalStats": False, + "hasPromptStats": True, + "tierLabel": "", + "usageStatusText": "", + "authHelpText": "", + "limits": [], + } + record.update(empty_stats()) + record.update(overrides) + return record + + +# -------------------------------------------------------------------------- +# Credentials +# -------------------------------------------------------------------------- + +def read_env_file(path: Path) -> dict[str, str]: + values: dict[str, str] = {} + try: + for line in path.read_text().splitlines(): + line = line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, _, val = line.partition("=") + values[key.strip()] = val.strip().strip("'\"") + except OSError: + return {} + return values + + +def read_config() -> dict[str, Any]: + config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) + try: + parsed = json.loads((config_home / "omarchy" / "agents" / "zai.json").read_text()) + return parsed if isinstance(parsed, dict) else {} + except (OSError, json.JSONDecodeError): + return {} + + +def credentials() -> tuple[str, str]: + """Return (api_key, platform). Platform is a PLATFORMS key.""" + config = read_config() + + # An explicit config file wins, then live env, then whatever omp/pi stored. + key = str(config.get("apiKey") or "").strip() + platform = str(config.get("platform") or "").strip().lower() + + if not key: + if os.environ.get("ZAI_API_KEY"): + key, platform = os.environ["ZAI_API_KEY"].strip(), platform or "zai" + elif os.environ.get("ZHIPU_API_KEY") or os.environ.get("ZHIPUAI_API_KEY"): + key = (os.environ.get("ZHIPU_API_KEY") or os.environ.get("ZHIPUAI_API_KEY")).strip() + platform = platform or "zhipu" + + if not key: + for env_path in ( + Path.home() / ".omp" / "agent" / ".env", + Path.home() / ".pi" / "agent" / ".env", + ): + values = read_env_file(env_path) + if values.get("ZAI_API_KEY"): + key, platform = values["ZAI_API_KEY"], platform or "zai" + break + if values.get("ZHIPU_API_KEY") or values.get("ZHIPUAI_API_KEY"): + key = values.get("ZHIPU_API_KEY") or values["ZHIPUAI_API_KEY"] + platform = platform or "zhipu" + break + + if platform not in PLATFORMS: + platform = "zai" + return key, platform + + +# -------------------------------------------------------------------------- +# Coding-plan limits +# -------------------------------------------------------------------------- + +def iso_from_ms(value: Any) -> str: + try: + ms = float(value) + except (TypeError, ValueError): + return "" + if ms <= 0: + return "" + return datetime.fromtimestamp(ms / 1000, timezone.utc).isoformat() + + +def fetch_limits(api_key: str, platform: str) -> dict[str, Any]: + """Return {"limits": [...], "tierLabel": str} or {"status": str} on miss. + + The monitor endpoint answers HTTP 200 even when the key has no plan, carrying + {"success": false, "msg": "…coding plan…"} instead of a data payload. + """ + url = PLATFORMS[platform] + QUOTA_PATH + request = urllib.request.Request( + url, + method="GET", + headers={ + # This endpoint takes the raw token, not a Bearer prefix. + "Authorization": api_key, + "Accept-Language": "en-US,en", + "Content-Type": "application/json", + }, + ) + try: + with urllib.request.urlopen(request, timeout=15) as response: + body = json.load(response) + except urllib.error.HTTPError as error: + if error.code in (401, 403): + return {"status": "Z.ai rejected the API key"} + return {"status": f"Z.ai monitor returned HTTP {error.code}"} + except (urllib.error.URLError, TimeoutError): + return {"status": "Could not reach the Z.ai monitor endpoint"} + except json.JSONDecodeError: + return {"status": "Z.ai returned an invalid monitor response"} + + data = body.get("data") if isinstance(body, dict) else None + if not isinstance(data, dict) or not isinstance(data.get("limits"), list): + msg = str((body or {}).get("msg") or "").strip() + status = "No GLM Coding Plan on this key" if "coding plan" in msg.lower() else ( + msg or "Coding Plan limits unavailable" + ) + return {"status": status} + + limits = [] + for item in data["limits"]: + if not isinstance(item, dict) or item.get("type") != "TOKENS_LIMIT": + continue + label = TOKEN_WINDOWS.get((item.get("unit"), item.get("number"))) + if not label: + continue + try: + percent = max(0.0, min(1.0, float(item.get("percentage") or 0) / 100.0)) + except (TypeError, ValueError): + percent = 0.0 + limits.append({ + "label": label, + "percent": percent, + "resetsAt": iso_from_ms(item.get("nextResetTime")), + }) + + tier = str(data.get("level") or "").strip() + return {"limits": limits, "tierLabel": tier[:1].upper() + tier[1:].lower() if tier else ""} + + +# -------------------------------------------------------------------------- +# Local session stats +# -------------------------------------------------------------------------- + +def scan_sessions() -> dict[str, Any]: + now = datetime.now() + today = now.strftime("%Y-%m-%d") + recent_dates = [(now - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)] + recent = {day: {"date": day, "messageCount": 0} for day in recent_dates} + + model_usage: dict[str, dict[str, int]] = {} + today_by_model: dict[str, int] = {} + active_dates: set[str] = set() + total_sessions: set[str] = set() + today_sessions: set[str] = set() + seen: set[str] = set() + totals = {"today_prompts": 0, "today_tokens": 0, "total_prompts": 0} + + def add_usage(day, session_key, model, inp, out, cread, cwrite): + total = inp + out + cread + cwrite + totals["total_prompts"] += 1 + total_sessions.add(session_key) + active_dates.add(day) + bucket = model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += inp + bucket["outputTokens"] += out + bucket["cacheReadInputTokens"] += cread + bucket["cacheCreationInputTokens"] += cwrite + if day in recent: + recent[day]["messageCount"] += total + if day == today: + totals["today_prompts"] += 1 + today_sessions.add(session_key) + totals["today_tokens"] += total + today_by_model[model] = today_by_model.get(model, 0) + total + + roots = [ + Path.home() / ".pi" / "agent" / "sessions", + Path.home() / ".omp" / "agent" / "sessions", + ] + for root in roots: + if not root.exists(): + continue + for path in root.rglob("*.jsonl"): + try: + handle = path.open(errors="replace") + except OSError: + continue + with handle: + for line in handle: + # Cheap substring gate before the JSON parse: an assistant turn on a + # Z.ai provider always carries both markers. + if '"usage"' not in line or ('"zai"' not in line and '"zhipu"' not in line): + continue + try: + entry = json.loads(line) + except json.JSONDecodeError: + continue + if not isinstance(entry, dict) or entry.get("type") != "message": + continue + message = entry.get("message") or {} + if message.get("role") != "assistant": + continue + provider = str(message.get("provider") or "").lower() + if "zai" not in provider and "zhipu" not in provider: + continue + usage = message.get("usage") or {} + inp = number(usage.get("input")) + out = number(usage.get("output")) + cread = number(usage.get("cacheRead")) + cwrite = number(usage.get("cacheWrite")) + if not (inp or out or cread or cwrite): + inp = number(usage.get("totalTokens")) + if not (inp or out or cread or cwrite): + continue + message_key = str(path) + ":" + str(entry.get("id") or "") + if message_key in seen: + continue + seen.add(message_key) + day = local_day(entry.get("timestamp") or message.get("ts")) + model = str(message.get("model") or "glm") + add_usage(day, str(path), model, inp, out, cread, cwrite) + + return { + "hasLocalStats": bool(model_usage), + "todayPrompts": totals["today_prompts"], + "todaySessions": len(today_sessions), + "todayTotalTokens": totals["today_tokens"], + "todayTokensByModel": today_by_model, + "recentDays": [recent[day] for day in recent_dates], + "totalPrompts": totals["total_prompts"], + "totalSessions": len(total_sessions), + "activeDays": len(active_dates), + "activeDates": sorted(active_dates), + "modelUsage": model_usage, + } + + +# -------------------------------------------------------------------------- +# Assembly +# -------------------------------------------------------------------------- + +def scan() -> dict[str, Any]: + record = base_record() + stats = scan_sessions() + record.update(stats) + + api_key, platform = credentials() + if not api_key: + record["authHelpText"] = AUTH_HELP_NO_KEY + record["ready"] = record["hasLocalStats"] + return record + + result = fetch_limits(api_key, platform) + if result.get("limits"): + record["limits"] = result["limits"] + if result.get("tierLabel"): + record["tierLabel"] = result["tierLabel"] + else: + # A key with no plan (or an endpoint error) still shows local usage. + record["usageStatusText"] = result.get("status", "Coding Plan limits unavailable") + record["authHelpText"] = AUTH_HELP_NO_KEY + + record["ready"] = bool(record["limits"] or record["hasLocalStats"]) + return record + + +def main() -> int: + parser = argparse.ArgumentParser(description="Print the Z.ai GLM usage record as JSON") + # Limits and stats come from one endpoint and a local scan, so there is no + # cache to force past or a faster limits-only path; the flags exist only so + # every collector accepts the same invocation. + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + parser.parse_args() + + try: + record = scan() + except Exception as error: # noqa: BLE001 - a collector must always emit a record + record = base_record(usageStatusText="Z.ai unavailable", authHelpText="Z.ai usage scan failed") + print(f"omarchy-agent-usage-zai: {type(error).__name__}", file=sys.stderr) + print(json.dumps(record, separators=(",", ":"))) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 465d7fdb1207dd31496fbbde0b577322a935083e Mon Sep 17 00:00:00 2001 From: FelipeMayerDev Date: Thu, 20 Aug 2026 00:12:46 -0300 Subject: [PATCH 2/5] Fix Z.ai limit parsing: monitor endpoint returns CREDIT_LIMIT The coding-plan monitor endpoint reports each 5-hour/weekly window as a CREDIT_LIMIT row, not TOKENS_LIMIT. Accept both types so real Coding Plan keys surface their Session and Weekly limits instead of an empty list. --- bin/omarchy-agent-usage-zai | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/bin/omarchy-agent-usage-zai b/bin/omarchy-agent-usage-zai index 0f82c489a9d..55239a9c2b6 100755 --- a/bin/omarchy-agent-usage-zai +++ b/bin/omarchy-agent-usage-zai @@ -39,9 +39,10 @@ PLATFORMS = { } QUOTA_PATH = "/api/monitor/usage/quota/limit" -# Coding-plan token windows are tagged by (unit, number); the monitor endpoint -# reuses TOKENS_LIMIT for both the 5-hour and the weekly window. -TOKEN_WINDOWS = { +# Coding-plan limit windows are tagged by (unit, number); the monitor endpoint +# reports one CREDIT_LIMIT row per window (older responses used TOKENS_LIMIT). +LIMIT_TYPES = {"CREDIT_LIMIT", "TOKENS_LIMIT"} +LIMIT_WINDOWS = { (3, 5): "Session (5-hour)", (6, 1): "Weekly (7-day)", } @@ -232,9 +233,9 @@ def fetch_limits(api_key: str, platform: str) -> dict[str, Any]: limits = [] for item in data["limits"]: - if not isinstance(item, dict) or item.get("type") != "TOKENS_LIMIT": + if not isinstance(item, dict) or item.get("type") not in LIMIT_TYPES: continue - label = TOKEN_WINDOWS.get((item.get("unit"), item.get("number"))) + label = LIMIT_WINDOWS.get((item.get("unit"), item.get("number"))) if not label: continue try: From bfc6f4b30265230bb05cfdc844ec092919c6957e Mon Sep 17 00:00:00 2001 From: FelipeMayerDev Date: Thu, 20 Aug 2026 00:12:46 -0300 Subject: [PATCH 3/5] Add Z.ai brand mark for the agents panel assets/zai.svg (white) for dark surfaces and assets/zai-light.svg (#111) for light surfaces, following the codex mark convention; the panel picks the twin by surface luminance. --- shell/plugins/agents/assets/zai-light.svg | 1 + shell/plugins/agents/assets/zai.svg | 1 + 2 files changed, 2 insertions(+) create mode 100644 shell/plugins/agents/assets/zai-light.svg create mode 100644 shell/plugins/agents/assets/zai.svg diff --git a/shell/plugins/agents/assets/zai-light.svg b/shell/plugins/agents/assets/zai-light.svg new file mode 100644 index 00000000000..040a6d88830 --- /dev/null +++ b/shell/plugins/agents/assets/zai-light.svg @@ -0,0 +1 @@ +Z.ai diff --git a/shell/plugins/agents/assets/zai.svg b/shell/plugins/agents/assets/zai.svg new file mode 100644 index 00000000000..48c8d1d7397 --- /dev/null +++ b/shell/plugins/agents/assets/zai.svg @@ -0,0 +1 @@ +Z.ai From acffe0c372686a3f7db91ea36aa1adba455c1d1f Mon Sep 17 00:00:00 2001 From: Felipe Mayer Date: Mon, 24 Aug 2026 11:05:42 -0300 Subject: [PATCH 4/5] Address review: limits/scan caching, token & timestamp fixes, tests - --limits-only/--force reuse a probe cache (15s min interval) and a local-stats scan cache; keep last-known limits while the window is open - fold reasoning tokens into output; date turns by message.timestamp - exact provider match + isinstance guards so one bad line can't drop stats - retryAdvised on transport failure; blame the key only on 401/403 - iso_from_ms accepts seconds; skip unparseable percentage; note unmapped windows - rename tab to Z.ai; README row + manifest providers default; .env export= parsing - add scanner and limits regression tests --- bin/omarchy-agent-usage-zai | 296 ++++++++++++++++--- shell/plugins/agents/README.md | 7 + shell/plugins/agents/manifest.json | 3 +- test/shell.d/agent-usage-zai-limits-test.sh | 141 +++++++++ test/shell.d/agent-usage-zai-scanner-test.sh | 90 ++++++ 5 files changed, 498 insertions(+), 39 deletions(-) create mode 100755 test/shell.d/agent-usage-zai-limits-test.sh create mode 100755 test/shell.d/agent-usage-zai-scanner-test.sh diff --git a/bin/omarchy-agent-usage-zai b/bin/omarchy-agent-usage-zai index 55239a9c2b6..7d7ce261739 100755 --- a/bin/omarchy-agent-usage-zai +++ b/bin/omarchy-agent-usage-zai @@ -11,15 +11,22 @@ generic pay-as-you-go key gets "no coding plan" and the record falls back to local stats alone, exactly like the claude collector without a signed-in CLI. Local stats come from pi/omp sessions that ran through the Z.ai provider (the -`glm-*` models). The agents panel only ever reads the JSON this prints. +`glm-*` models). Both the local scan and the network probe are cached: the +panel fires `--limits-only` on every open, so a repeated open must reuse a +recent scan and must not re-probe Z.ai more than once every few seconds. The +agents panel only ever reads the JSON this prints. """ from __future__ import annotations import argparse +import fcntl +import hashlib import json import os import sys +import tempfile +import time import urllib.error import urllib.request from datetime import datetime, timedelta, timezone @@ -27,10 +34,11 @@ from pathlib import Path from typing import Any AGENT_ID = "zai" -AGENT_NAME = "GLM Coding Plan" +AGENT_NAME = "Z.ai" AUTH_HELP_NO_KEY = ( "Set ZAI_API_KEY (or add ~/.config/omarchy/agents/zai.json) to show Coding Plan limits." ) +AUTH_HELP_BAD_KEY = "Z.ai rejected the saved API key; check ~/.config/omarchy/agents/zai.json." # Global (Z.ai) vs China (Zhipu / BigModel) split the same API under two hosts. PLATFORMS = { @@ -47,6 +55,20 @@ LIMIT_WINDOWS = { (6, 1): "Weekly (7-day)", } +# A scan this recent is only reused to dedup concurrent collector runs (the +# update command backgrounds one per agent while the panel refreshes on its +# own); every periodic widget refresh lands a real rescan. --limits-only +# promises only fresh limits, so it may reuse a scan for up to 15 minutes. +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 +# The panel opens and shuts fire a probe each; a recent probe result is reused +# for this window rather than turning every flick into a request. +PROBE_MIN_INTERVAL_SECONDS = 15 + +now = datetime.now() +today = now.strftime("%Y-%m-%d") +recent_dates = [(now - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)] + def number(value: Any) -> int: try: @@ -119,6 +141,46 @@ def base_record(**overrides: Any) -> dict[str, Any]: return record +# -------------------------------------------------------------------------- +# Cache primitives (shared by the local-stats scan and the limits probe) +# -------------------------------------------------------------------------- + +def cache_root() -> Path: + root = Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage" + root.mkdir(parents=True, exist_ok=True) + return root + + +def read_fresh_json(path: Path, max_age_seconds: float) -> Any: + if max_age_seconds <= 0 or not path.exists(): + return None + try: + # A negative age means the mtime is in the future: the clock moved + # backwards since the write, so the cache's freshness cannot be trusted. + age = time.time() - path.stat().st_mtime + if 0 <= age <= max_age_seconds: + return json.loads(path.read_text(encoding="utf-8")) + except Exception: + return None + return None + + +def write_json(path: Path, payload: Any) -> None: + # A temp name unique to this writer, not derived from the target: several + # collectors can run at once, and a shared temp path means the second + # replace finds the first one's file already moved away. + handle_fd, tmp_name = tempfile.mkstemp(dir=path.parent, prefix=path.name + ".", suffix=".tmp") + tmp = Path(tmp_name) + try: + with os.fdopen(handle_fd, "w", encoding="utf-8") as handle: + handle.write(json.dumps(payload, separators=(",", ":")) + "\n") + tmp.chmod(0o644) + tmp.replace(path) + except BaseException: + tmp.unlink(missing_ok=True) + raise + + # -------------------------------------------------------------------------- # Credentials # -------------------------------------------------------------------------- @@ -128,6 +190,8 @@ def read_env_file(path: Path) -> dict[str, str]: try: for line in path.read_text().splitlines(): line = line.strip() + if line.startswith("export "): + line = line[len("export "):].strip() if not line or line.startswith("#") or "=" not in line: continue key, _, val = line.partition("=") @@ -186,18 +250,28 @@ def credentials() -> tuple[str, str]: def iso_from_ms(value: Any) -> str: try: - ms = float(value) + ts = float(value) except (TypeError, ValueError): return "" - if ms <= 0: + if ts <= 0: + return "" + # The field is milliseconds, but a plain-seconds epoch (< year 33658) is + # accepted too rather than rendered as a date 1000x in the future. + if ts < 1e12: + ts *= 1000 + try: + return datetime.fromtimestamp(ts / 1000, timezone.utc).isoformat() + except (OverflowError, OSError, ValueError): return "" - return datetime.fromtimestamp(ms / 1000, timezone.utc).isoformat() def fetch_limits(api_key: str, platform: str) -> dict[str, Any]: - """Return {"limits": [...], "tierLabel": str} or {"status": str} on miss. + """Probe the monitor endpoint. - The monitor endpoint answers HTTP 200 even when the key has no plan, carrying + On success returns {"limits": [...], "tierLabel": str}. On a miss returns + {"status": str} plus, when applicable, "transport" (reached no server, so + worth an early retry) or "keyRejected" (the key itself is bad). The endpoint + answers HTTP 200 even when the key has no plan, carrying {"success": false, "msg": "…coding plan…"} instead of a data payload. """ url = PLATFORMS[platform] + QUOTA_PATH @@ -216,16 +290,16 @@ def fetch_limits(api_key: str, platform: str) -> dict[str, Any]: body = json.load(response) except urllib.error.HTTPError as error: if error.code in (401, 403): - return {"status": "Z.ai rejected the API key"} + return {"status": "Z.ai rejected the API key", "keyRejected": True} return {"status": f"Z.ai monitor returned HTTP {error.code}"} except (urllib.error.URLError, TimeoutError): - return {"status": "Could not reach the Z.ai monitor endpoint"} + return {"status": "Could not reach the Z.ai monitor endpoint", "transport": True} except json.JSONDecodeError: return {"status": "Z.ai returned an invalid monitor response"} data = body.get("data") if isinstance(body, dict) else None if not isinstance(data, dict) or not isinstance(data.get("limits"), list): - msg = str((body or {}).get("msg") or "").strip() + msg = str(body.get("msg") or "").strip() if isinstance(body, dict) else "" status = "No GLM Coding Plan on this key" if "coding plan" in msg.lower() else ( msg or "Coding Plan limits unavailable" ) @@ -237,11 +311,19 @@ def fetch_limits(api_key: str, platform: str) -> dict[str, Any]: continue label = LIMIT_WINDOWS.get((item.get("unit"), item.get("number"))) if not label: + # A window Z.ai adds later would vanish silently otherwise. + print( + f"omarchy-agent-usage-zai: unmapped limit window (unit={item.get('unit')}, " + f"number={item.get('number')})", + file=sys.stderr, + ) continue try: - percent = max(0.0, min(1.0, float(item.get("percentage") or 0) / 100.0)) + percent = max(0.0, min(1.0, float(item.get("percentage")) / 100.0)) except (TypeError, ValueError): - percent = 0.0 + # A missing/unparseable percentage is skipped: a false 0% meter reads + # as a fresh allowance, which is worse than showing one fewer window. + continue limits.append({ "label": label, "percent": percent, @@ -252,14 +334,84 @@ def fetch_limits(api_key: str, platform: str) -> dict[str, Any]: return {"limits": limits, "tierLabel": tier[:1].upper() + tier[1:].lower() if tier else ""} +# A cached percentage outlives the probe that measured it, but only until its +# window rolls over: once a window has reset, the figure describes a period +# that is over. A window with no reset time, or one that will not parse, is +# kept — an unreadable timestamp is no reason to throw away a real number. +def limit_window_open(entry: dict[str, Any], now_utc: datetime) -> bool: + raw = str(entry.get("resetsAt") or "") + if raw == "": + return True + try: + resets_at = datetime.fromisoformat(raw.replace("Z", "+00:00")) + except ValueError: + return True + if resets_at.tzinfo is None: + resets_at = resets_at.replace(tzinfo=timezone.utc) + return resets_at > now_utc + + +def usable_cached_limits(cached: dict[str, Any]) -> list[dict[str, Any]]: + entries = cached.get("limits") + if not isinstance(entries, list): + return [] + now_utc = datetime.now(timezone.utc) + return [entry for entry in entries if isinstance(entry, dict) and limit_window_open(entry, now_utc)] + + +def collect_limits(api_key: str, platform: str, force: bool) -> dict[str, Any]: + """Fresh limits when a probe is due, else the last known ones. + + Reuses a recent probe for PROBE_MIN_INTERVAL_SECONDS so repeated panel opens + do not spam the endpoint, and keeps the last known limits as the answer of + record when a later probe fails but its window has not yet reset. + """ + result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "retryAdvised": False} + + probe_cache = cache_root() / "zai-limits.json" + cached = read_fresh_json(probe_cache, float("inf")) or {} + fallback = usable_cached_limits(cached) if isinstance(cached, dict) else [] + cached_tier = str(cached.get("tierLabel") or "") if isinstance(cached, dict) else "" + + fetched_at = number(cached.get("fetchedAtMs")) / 1000 if isinstance(cached, dict) else 0 + if fallback and not force and time.time() - fetched_at < PROBE_MIN_INTERVAL_SECONDS: + result["limits"], result["tierLabel"] = fallback, cached_tier + return result + + probe = fetch_limits(api_key, platform) + if probe.get("limits"): + result["limits"] = probe["limits"] + result["tierLabel"] = probe.get("tierLabel", "") + try: + write_json(probe_cache, { + "fetchedAtMs": round(time.time() * 1000), + "limits": probe["limits"], + "tierLabel": result["tierLabel"], + }) + except Exception as exc: + print(f"omarchy-agent-usage-zai: could not write limits cache ({exc})", file=sys.stderr) + return result + + # No fresh limits. A transport failure reached no server — typically the + # seconds after login before the network is up — so ask the shell to retry + # sooner than its regular interval. + if probe.get("transport"): + result["retryAdvised"] = True + if fallback: + result["limits"], result["tierLabel"] = fallback, cached_tier + result["usageStatusText"] = probe.get("status", "") + # Only blame the key when the key is the problem; a working key with no plan, + # or an endpoint hiccup, must not tell the user to go set a key they have. + if probe.get("keyRejected"): + result["authHelpText"] = AUTH_HELP_BAD_KEY + return result + + # -------------------------------------------------------------------------- # Local session stats # -------------------------------------------------------------------------- def scan_sessions() -> dict[str, Any]: - now = datetime.now() - today = now.strftime("%Y-%m-%d") - recent_dates = [(now - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)] recent = {day: {"date": day, "messageCount": 0} for day in recent_dates} model_usage: dict[str, dict[str, int]] = {} @@ -312,15 +464,22 @@ def scan_sessions() -> dict[str, Any]: continue if not isinstance(entry, dict) or entry.get("type") != "message": continue - message = entry.get("message") or {} - if message.get("role") != "assistant": + # One malformed line (a non-dict where a dict belongs) must not take + # every real stat down with it, so each shape is checked explicitly. + message = entry.get("message") + if not isinstance(message, dict) or message.get("role") != "assistant": continue + # Exact provider match: a custom `zai-proxy` gateway is not this + # subscription, the same discipline the stock collectors keep. provider = str(message.get("provider") or "").lower() - if "zai" not in provider and "zhipu" not in provider: + if provider not in ("zai", "zhipu"): + continue + usage = message.get("usage") + if not isinstance(usage, dict): continue - usage = message.get("usage") or {} inp = number(usage.get("input")) - out = number(usage.get("output")) + # pi reports reasoning tokens apart from output; both are generated. + out = number(usage.get("output")) + number(usage.get("reasoning")) cread = number(usage.get("cacheRead")) cwrite = number(usage.get("cacheWrite")) if not (inp or out or cread or cwrite): @@ -331,7 +490,7 @@ def scan_sessions() -> dict[str, Any]: if message_key in seen: continue seen.add(message_key) - day = local_day(entry.get("timestamp") or message.get("ts")) + day = local_day(entry.get("timestamp") or message.get("timestamp")) model = str(message.get("model") or "glm") add_usage(day, str(path), model, inp, out, cread, cwrite) @@ -350,14 +509,75 @@ def scan_sessions() -> dict[str, Any]: } +def scan_cache_paths() -> tuple[Path, Path]: + # The digest covers every data path the scan reads (both session roots live + # under the home directory). + digest = hashlib.sha1(str(Path.home()).encode("utf-8")).hexdigest()[:16] + root = cache_root() + return root / f"zai-scan-{digest}.json", root / f"zai-scan-{digest}.lock" + + +def read_cached_stats(cache_file: Path, max_age_seconds: float) -> dict[str, Any] | None: + cached = read_fresh_json(cache_file, max_age_seconds) + if not isinstance(cached, dict) or cached.get("schemaVersion") != 1: + return None + # today* fields only mean "today" on the day they were scanned. A cache from + # another local date (midnight passed, or the clock moved) is a miss. + if cached.get("scanDate") != today: + return None + stats = cached.get("stats") + if not isinstance(stats, dict): + return None + if not all(key in stats for key in ("todayPrompts", "todayTotalTokens", "recentDays", "activeDates", "modelUsage")): + return None + return stats + + +def write_cached_stats(cache_file: Path, stats: dict[str, Any]) -> None: + try: + write_json(cache_file, {"schemaVersion": 1, "scanDate": today, "stats": stats}) + except Exception as exc: + print(f"omarchy-agent-usage-zai: could not write usage cache ({exc})", file=sys.stderr) + + +def cached_local_stats(max_age: float) -> dict[str, Any]: + """Local stats, with the cache as a pure optimization. + + Any cache-layer failure degrades to a direct scan and a warning: the JSON + record is the contract, the cache is not. + """ + try: + return _cached_local_stats(max_age) + except Exception as exc: + print(f"omarchy-agent-usage-zai: cache unavailable ({exc}); scanning directly", file=sys.stderr) + return scan_sessions() + + +def _cached_local_stats(max_age: float) -> dict[str, Any]: + cache_file, lock_file = scan_cache_paths() + + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + stats = scan_sessions() + write_cached_stats(cache_file, stats) + return stats + + # -------------------------------------------------------------------------- # Assembly # -------------------------------------------------------------------------- -def scan() -> dict[str, Any]: +def scan(force: bool, limits_only: bool) -> dict[str, Any]: record = base_record() - stats = scan_sessions() - record.update(stats) + max_age = 0 if force else (LIMITS_ONLY_REUSE_SECONDS if limits_only else SCAN_REUSE_SECONDS) + record.update(cached_local_stats(max_age)) api_key, platform = credentials() if not api_key: @@ -365,15 +585,15 @@ def scan() -> dict[str, Any]: record["ready"] = record["hasLocalStats"] return record - result = fetch_limits(api_key, platform) - if result.get("limits"): - record["limits"] = result["limits"] - if result.get("tierLabel"): - record["tierLabel"] = result["tierLabel"] - else: - # A key with no plan (or an endpoint error) still shows local usage. - record["usageStatusText"] = result.get("status", "Coding Plan limits unavailable") - record["authHelpText"] = AUTH_HELP_NO_KEY + limits = collect_limits(api_key, platform, force) + record["limits"] = limits["limits"] + if limits["tierLabel"]: + record["tierLabel"] = limits["tierLabel"] + record["usageStatusText"] = limits["usageStatusText"] + if limits["authHelpText"]: + record["authHelpText"] = limits["authHelpText"] + if limits["retryAdvised"]: + record["retryAdvised"] = True record["ready"] = bool(record["limits"] or record["hasLocalStats"]) return record @@ -381,15 +601,15 @@ def scan() -> dict[str, Any]: def main() -> int: parser = argparse.ArgumentParser(description="Print the Z.ai GLM usage record as JSON") - # Limits and stats come from one endpoint and a local scan, so there is no - # cache to force past or a faster limits-only path; the flags exist only so - # every collector accepts the same invocation. + # --force rescans and re-probes, ignoring both caches. --limits-only reuses a + # recent local scan (only the limits probe must be fresh); a bare run reuses a + # scan only long enough to dedup concurrent collector runs. parser.add_argument("--force", action="store_true") parser.add_argument("--limits-only", action="store_true") - parser.parse_args() + args = parser.parse_args() try: - record = scan() + record = scan(args.force, args.limits_only) except Exception as error: # noqa: BLE001 - a collector must always emit a record record = base_record(usageStatusText="Z.ai unavailable", authHelpText="Z.ai usage scan failed") print(f"omarchy-agent-usage-zai: {type(error).__name__}", file=sys.stderr) diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 51ad4e43c2d..3aae2f36580 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -55,6 +55,7 @@ light surfaces — and the bar glyph stands in when there is none. | `claude` | Anthropic's OAuth usage endpoint (5-hour session + 7-day weekly) | `~/.claude/projects` transcripts, opencode sessions on an Anthropic provider, plus `stats-cache.json` and `history.jsonl` as fallback | | `codex` | The Codex app-server RPC | native Codex CLI session files (plus pi and opencode sessions) | | `fireworks` | Estimated prepaid balance: configured funding minus rated account costs | Fireworks billing API, grouped by day and model for the last 30 days | +| `zai` | Z.ai's coding-plan monitor endpoint (5-hour session + 7-day weekly) | pi/omp sessions on the Z.ai provider (`glm-*` models) | Claude limits need a signed-in CLI; without credentials the panel says so and falls back to local stats only. A non-default Claude directory is honored via @@ -64,6 +65,12 @@ falls back to local stats only. A non-default Claude directory is honored via opencode stores in `~/.local/share/opencode/auth.json` when Fireworks is signed in there. +Z.ai limits need a key that carries a GLM Coding Plan; a pay-as-you-go key +reports "no coding plan" and the panel falls back to local stats only. The key +comes from `~/.config/omarchy/agents/zai.json`, then `ZAI_API_KEY` / +`ZHIPU_API_KEY` in the environment, then pi/omp's `.env`. Set `platform` to +`"zhipu"` in the config for the China (`open.bigmodel.cn`) host. + ### Fireworks balance The collector first asks the account's `:getBalance` endpoint for the real diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index 1743a798214..b9f80c3d751 100644 --- a/shell/plugins/agents/manifest.json +++ b/shell/plugins/agents/manifest.json @@ -21,7 +21,8 @@ "providers": { "claude": { "enabled": true }, "codex": { "enabled": true }, - "fireworks": { "enabled": true } + "fireworks": { "enabled": true }, + "zai": { "enabled": true } }, "refreshIntervalSec": 900, "syncMode": "Off", diff --git a/test/shell.d/agent-usage-zai-limits-test.sh b/test/shell.d/agent-usage-zai-limits-test.sh new file mode 100755 index 00000000000..7cda532f0e0 --- /dev/null +++ b/test/shell.d/agent-usage-zai-limits-test.sh @@ -0,0 +1,141 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +# fetch_limits/collect_limits reach Z.ai, so the readers that interpret the +# response are exercised on their own: the collector loads as a module and a +# recorded payload (or a raised error) stands in for the network. +CACHE_HOME=$(mktemp -d) +trap 'rm -rf "$CACHE_HOME"' EXIT + +read_limits() { + COLLECTOR="$ROOT/bin/omarchy-agent-usage-zai" PAYLOAD="$1" python3 - <<'PY' +import importlib.machinery, importlib.util, io, json, os + +loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) + +collector.urllib.request.urlopen = lambda request, timeout=None: io.BytesIO(os.environ["PAYLOAD"].encode()) +print(json.dumps(collector.fetch_limits("token", "zai"))) +PY +} + +# force, cached-json, mode (success|transport|http401), payload -> {result, probes, cached} +drive() { + COLLECTOR="$ROOT/bin/omarchy-agent-usage-zai" FORCE="$1" CACHED="$2" MODE="$3" PAYLOAD="$4" \ + XDG_CACHE_HOME="$CACHE_HOME" python3 - <<'PY' +import importlib.machinery, importlib.util, io, json, os, urllib.error + +loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) + +cache = collector.cache_root() / "zai-limits.json" +cached = os.environ["CACHED"] +if cached: + cache.write_text(cached, encoding="utf-8") +elif cache.exists(): + cache.unlink() + +probes = [] +mode = os.environ["MODE"] + +def urlopen(request, timeout=None): + probes.append(1) + if mode == "transport": + raise urllib.error.URLError("no route to host") + if mode == "http401": + raise urllib.error.HTTPError("u", 401, "no", {}, None) + return io.BytesIO(os.environ["PAYLOAD"].encode()) + +collector.urllib.request.urlopen = urlopen +result = collector.collect_limits("token", "zai", os.environ["FORCE"] == "true") +out = {"result": result, "probes": len(probes)} +if cache.exists(): + out["cached"] = json.loads(cache.read_text(encoding="utf-8")) +print(json.dumps(out)) +PY +} + +soon_ms=$(python3 -c 'import time; print(round((time.time() + 3*3600) * 1000))') +week_sec=$(python3 -c 'import time; print(round(time.time() + 5*86400))') +week_date=$(python3 -c "import datetime; print(datetime.datetime.fromtimestamp($week_sec, datetime.timezone.utc).date().isoformat())") +open_at=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) + dt.timedelta(hours=3)).isoformat())') +past_at=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) - dt.timedelta(hours=3)).isoformat())') + +# The mapped session + weekly windows, then the entries that must drop: a window +# tagged with an unmapped (unit, number) and one whose percentage will not parse. +# nextResetTime is milliseconds for the first and plain seconds for the second, +# to prove a seconds epoch is not rendered 1000x in the future. +payload=$(jq -nc --argjson sm "$soon_ms" --argjson ws "$week_sec" '{ + data: { level: "pro", limits: [ + { type: "CREDIT_LIMIT", unit: 3, number: 5, percentage: 42.0, nextResetTime: $sm }, + { type: "TOKENS_LIMIT", unit: 6, number: 1, percentage: 10, nextResetTime: $ws }, + { type: "CREDIT_LIMIT", unit: 9, number: 9, percentage: 5, nextResetTime: $sm }, + { type: "CREDIT_LIMIT", unit: 3, number: 5, percentage: null, nextResetTime: $sm } + ] } +}') +limits=$(read_limits "$payload") + +[[ $(jq -c '[.limits[] | {label, percent}]' <<<"$limits") == '[{"label":"Session (5-hour)","percent":0.42},{"label":"Weekly (7-day)","percent":0.1}]' ]] || + fail "Z.ai collector maps both windows and drops the unmapped and unparseable rows" "$limits" +[[ $(jq -r '.tierLabel' <<<"$limits") == "Pro" ]] || + fail "Z.ai collector titles the plan tier" "$limits" +[[ $(jq -r '.limits[1].resetsAt' <<<"$limits") == "$week_date"* ]] || + fail "Z.ai collector reads a plain-seconds reset time, not milliseconds" "$limits" +pass "Z.ai collector reads windows, tier, seconds/ms resets, and drops bad rows" + +# A key with no plan explains the empty section without blaming the key. +noplan=$(drive false "" success '{"success":false,"msg":"This API key has no coding plan"}') +[[ $(jq -r '.result.usageStatusText' <<<"$noplan") == "No GLM Coding Plan on this key" ]] || + fail "Z.ai collector reports a key without a coding plan" "$noplan" +[[ $(jq -r '.result.authHelpText' <<<"$noplan") == "" && $(jq -r '.result.retryAdvised' <<<"$noplan") == "false" ]] || + fail "Z.ai collector does not tell a working key to set a key" "$noplan" +pass "Z.ai collector reports a no-plan key without blaming it" + +# A rejected key is the one case that should point back at the key. +rejected=$(drive false "" http401 "") +[[ $(jq -r '.result.authHelpText' <<<"$rejected") == *"rejected the saved API key"* ]] || + fail "Z.ai collector blames the key only when the key is rejected" "$rejected" +pass "Z.ai collector points at the key only on a 401" + +# A transport failure with an open cached window keeps the last known figures, +# drops the window that already reset, and asks the shell to retry sooner. +cache=$(jq -nc --arg open "$open_at" --arg past "$past_at" '{ + fetchedAtMs: 1, tierLabel: "Pro", limits: [ + { label: "Session (5-hour)", percent: 0.31, resetsAt: $past }, + { label: "Weekly (7-day)", percent: 0.11, resetsAt: $open } + ] +}') +unreachable=$(drive false "$cache" transport "") +[[ $(jq -c '[.result.limits[].label]' <<<"$unreachable") == '["Weekly (7-day)"]' ]] || + fail "Z.ai collector keeps only cached windows that have not reset" "$unreachable" +[[ $(jq -r '.result.tierLabel' <<<"$unreachable") == "Pro" && $(jq -r '.result.retryAdvised' <<<"$unreachable") == "true" ]] || + fail "Z.ai collector keeps the cached tier and advises a retry after a transport failure" "$unreachable" +pass "Z.ai collector serves the last known limits and advises a retry when offline" + +# Repeated panel opens inside the probe interval share one answer; --force does not. +fresh=$(jq -nc --arg open "$open_at" --argjson now "$(python3 -c 'import time; print(round(time.time() * 1000))')" '{ + fetchedAtMs: $now, tierLabel: "Pro", limits: [{ label: "Weekly (7-day)", percent: 0.11, resetsAt: $open }] +}') +success_payload=$(jq -nc --argjson sm "$soon_ms" '{ data: { level: "max", limits: [ + { type: "CREDIT_LIMIT", unit: 3, number: 5, percentage: 44.0, nextResetTime: $sm } +] } }') + +reused=$(drive false "$fresh" success "$success_payload") +[[ $(jq -r '.probes' <<<"$reused") == "0" && $(jq -c '[.result.limits[].percent]' <<<"$reused") == "[0.11]" ]] || + fail "Z.ai collector reuses a cache younger than the probe interval" "$reused" +pass "Z.ai collector reuses a probe younger than the interval" + +forced=$(drive true "$fresh" success "$success_payload") +[[ $(jq -r '.probes' <<<"$forced") == "1" && $(jq -c '[.result.limits[].percent]' <<<"$forced") == "[0.44]" ]] || + fail "Z.ai collector re-probes on --force despite a fresh cache" "$forced" +[[ $(jq -c '[.cached.limits[].percent]' <<<"$forced") == "[0.44]" && $(jq -r '.cached.tierLabel' <<<"$forced") == "Max" ]] || + fail "Z.ai collector caches a successful probe for the next run" "$forced" +pass "Z.ai collector re-probes on --force and caches the result" diff --git a/test/shell.d/agent-usage-zai-scanner-test.sh b/test/shell.d/agent-usage-zai-scanner-test.sh new file mode 100755 index 00000000000..a83cab33fab --- /dev/null +++ b/test/shell.d/agent-usage-zai-scanner-test.sh @@ -0,0 +1,90 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +mkdir -p "$TEST_HOME/.pi/agent/sessions/project" + +today="$(date +%Y-%m-%d)" +timestamp="${today}T12:00:00Z" + +# One session file exercises the whole scan contract: +# a zai, today, with reasoning + cache tokens -> counted, today +# b zhipu, today -> counted, today +# c openai provider -> ignored +# d zai but a user turn -> ignored +# a' a repeated message id -> deduped +# f zai, no top-level timestamp, message.timestamp -> counted, NOT today +# e malformed JSON -> skipped, rest survive +session="$TEST_HOME/.pi/agent/sessions/project/pi.jsonl" +cat >"$session" </dev/null | head -n 1) +[[ -n $cache_file && $(jq -r '.schemaVersion' "$cache_file") == "1" && $(jq -r '.stats.todayTotalTokens' "$cache_file") == "50" ]] || + fail "Z.ai collector writes a versioned local-stats cache on first scan" "$result" +pass "Z.ai collector writes a local-stats cache on first scan" + +# A new turn arrives; --limits-only must reuse the cache instead of rescanning. +cat >>"$session" < Date: Sat, 12 Sep 2026 18:04:18 -0300 Subject: [PATCH 5/5] Harden Z.ai collector inputs and cache locks --- bin/omarchy-agent-usage-zai | 90 ++++++++++++++++----- test/shell.d/agent-usage-zai-limits-test.sh | 10 +++ 2 files changed, 78 insertions(+), 22 deletions(-) diff --git a/bin/omarchy-agent-usage-zai b/bin/omarchy-agent-usage-zai index 7d7ce261739..ed684aef791 100755 --- a/bin/omarchy-agent-usage-zai +++ b/bin/omarchy-agent-usage-zai @@ -20,10 +20,12 @@ agents panel only ever reads the JSON this prints. from __future__ import annotations import argparse +import contextlib import fcntl import hashlib import json import os +import stat import sys import tempfile import time @@ -31,7 +33,7 @@ import urllib.error import urllib.request from datetime import datetime, timedelta, timezone from pathlib import Path -from typing import Any +from typing import Any, Iterator AGENT_ID = "zai" AGENT_NAME = "Z.ai" @@ -181,6 +183,42 @@ def write_json(path: Path, payload: Any) -> None: raise +@contextlib.contextmanager +def exclusive_lock(path: Path) -> Iterator[None]: + """Hold flock on `path` for the block, or run the block unlocked.""" + fd: int | None = None + try: + flags = os.O_WRONLY | os.O_CREAT | os.O_NOFOLLOW | os.O_NONBLOCK | os.O_CLOEXEC + try: + fd = os.open(path, flags | os.O_EXCL, 0o600) + except FileExistsError: + fd = os.open(path, flags) + info = os.fstat(fd) + if ( + not stat.S_ISREG(info.st_mode) + or info.st_nlink == 0 + or info.st_uid != os.geteuid() + or info.st_mode & 0o022 + ): + os.close(fd) + fd = None + except OSError: + if fd is not None: + os.close(fd) + fd = None + if fd is not None: + try: + fcntl.flock(fd, fcntl.LOCK_EX) + except OSError: + os.close(fd) + fd = None + try: + yield + finally: + if fd is not None: + os.close(fd) + + # -------------------------------------------------------------------------- # Credentials # -------------------------------------------------------------------------- @@ -287,7 +325,7 @@ def fetch_limits(api_key: str, platform: str) -> dict[str, Any]: ) try: with urllib.request.urlopen(request, timeout=15) as response: - body = json.load(response) + body = json.loads(response.read(1_048_576)) except urllib.error.HTTPError as error: if error.code in (401, 403): return {"status": "Z.ai rejected the API key", "keyRejected": True} @@ -299,7 +337,7 @@ def fetch_limits(api_key: str, platform: str) -> dict[str, Any]: data = body.get("data") if isinstance(body, dict) else None if not isinstance(data, dict) or not isinstance(data.get("limits"), list): - msg = str(body.get("msg") or "").strip() if isinstance(body, dict) else "" + msg = str(body.get("msg") or "").strip()[:200] if isinstance(body, dict) else "" status = "No GLM Coding Plan on this key" if "coding plan" in msg.lower() else ( msg or "Coding Plan limits unavailable" ) @@ -330,7 +368,7 @@ def fetch_limits(api_key: str, platform: str) -> dict[str, Any]: "resetsAt": iso_from_ms(item.get("nextResetTime")), }) - tier = str(data.get("level") or "").strip() + tier = str(data.get("level") or "").strip()[:64] return {"limits": limits, "tierLabel": tier[:1].upper() + tier[1:].lower() if tier else ""} @@ -378,19 +416,28 @@ def collect_limits(api_key: str, platform: str, force: bool) -> dict[str, Any]: result["limits"], result["tierLabel"] = fallback, cached_tier return result - probe = fetch_limits(api_key, platform) - if probe.get("limits"): - result["limits"] = probe["limits"] - result["tierLabel"] = probe.get("tierLabel", "") - try: - write_json(probe_cache, { - "fetchedAtMs": round(time.time() * 1000), - "limits": probe["limits"], - "tierLabel": result["tierLabel"], - }) - except Exception as exc: - print(f"omarchy-agent-usage-zai: could not write limits cache ({exc})", file=sys.stderr) - return result + with exclusive_lock(scan_cache_paths()[2]): + if not force: + cached = read_fresh_json(probe_cache, float("inf")) or {} + fetched_at = number(cached.get("fetchedAtMs")) / 1000 if isinstance(cached, dict) else 0 + fresh = usable_cached_limits(cached) if isinstance(cached, dict) else [] + if fresh and time.time() - fetched_at < PROBE_MIN_INTERVAL_SECONDS: + result["limits"], result["tierLabel"] = fresh, str(cached.get("tierLabel") or "") + return result + + probe = fetch_limits(api_key, platform) + if probe.get("limits"): + result["limits"] = probe["limits"] + result["tierLabel"] = probe.get("tierLabel", "") + try: + write_json(probe_cache, { + "fetchedAtMs": round(time.time() * 1000), + "limits": probe["limits"], + "tierLabel": result["tierLabel"], + }) + except Exception as exc: + print(f"omarchy-agent-usage-zai: could not write limits cache ({exc})", file=sys.stderr) + return result # No fresh limits. A transport failure reached no server — typically the # seconds after login before the network is up — so ask the shell to retry @@ -509,12 +556,12 @@ def scan_sessions() -> dict[str, Any]: } -def scan_cache_paths() -> tuple[Path, Path]: +def scan_cache_paths() -> tuple[Path, Path, Path]: # The digest covers every data path the scan reads (both session roots live # under the home directory). digest = hashlib.sha1(str(Path.home()).encode("utf-8")).hexdigest()[:16] root = cache_root() - return root / f"zai-scan-{digest}.json", root / f"zai-scan-{digest}.lock" + return root / f"zai-scan-{digest}.json", root / f"zai-scan-{digest}.lock", root / "zai-limits.lock" def read_cached_stats(cache_file: Path, max_age_seconds: float) -> dict[str, Any] | None: @@ -554,14 +601,13 @@ def cached_local_stats(max_age: float) -> dict[str, Any]: def _cached_local_stats(max_age: float) -> dict[str, Any]: - cache_file, lock_file = scan_cache_paths() + cache_file, lock_file, _ = scan_cache_paths() cached = read_cached_stats(cache_file, max_age) if cached is not None: return cached - with lock_file.open("w") as lock: - fcntl.flock(lock, fcntl.LOCK_EX) + with exclusive_lock(lock_file): cached = read_cached_stats(cache_file, max_age) if cached is not None: return cached diff --git a/test/shell.d/agent-usage-zai-limits-test.sh b/test/shell.d/agent-usage-zai-limits-test.sh index 7cda532f0e0..277c03e329a 100755 --- a/test/shell.d/agent-usage-zai-limits-test.sh +++ b/test/shell.d/agent-usage-zai-limits-test.sh @@ -139,3 +139,13 @@ forced=$(drive true "$fresh" success "$success_payload") [[ $(jq -c '[.cached.limits[].percent]' <<<"$forced") == "[0.44]" && $(jq -r '.cached.tierLabel' <<<"$forced") == "Max" ]] || fail "Z.ai collector caches a successful probe for the next run" "$forced" pass "Z.ai collector re-probes on --force and caches the result" + +# A predictable lock path must not follow and truncate a planted symlink. +sentinel="$CACHE_HOME/sentinel" +printf intact >"$sentinel" +rm -f "$CACHE_HOME/omarchy/agent-usage/zai-limits.lock" +ln -s "$sentinel" "$CACHE_HOME/omarchy/agent-usage/zai-limits.lock" +drive true "$fresh" success "$success_payload" >/dev/null +[[ $(<"$sentinel") == "intact" ]] || + fail "Z.ai collector does not follow a planted lock symlink" +pass "Z.ai collector does not follow a planted lock symlink"