From 2a42aa72acbf65f3cf0627b6b34dbcec0f7a4a46 Mon Sep 17 00:00:00 2001 From: steph4n <284007938+steph4n-gh@users.noreply.github.com> Date: Tue, 15 Sep 2026 13:58:46 -0400 Subject: [PATCH 1/7] Add Google Antigravity usage collector and panel integration Add the omarchy-agent-usage-antigravity collector to monitor Google Antigravity usage alongside Claude, Codex, and Fireworks in the native agents bar panel. Authoritative subscription tier and quota limits (rolling 5-hour session, weekly limits, and scoped 3P model allowances) are probed from the Google Cloud Code API via Secret Service keyring credentials with rolling local history estimation as fallback. Also teach the agents panel to watch ~/.config/omarchy/defaults/agent so the panel selects the user's default agent on launch. --- bin/omarchy-agent-usage-antigravity | 737 ++++++++++++++++++ shell/plugins/agents/Panel.qml | 30 +- shell/plugins/agents/README.md | 9 +- .../agents/assets/antigravity-light.svg | 1 + shell/plugins/agents/assets/antigravity.svg | 1 + shell/plugins/agents/manifest.json | 3 +- .../agent-usage-antigravity-limits-test.sh | 136 ++++ .../agent-usage-antigravity-scanner-test.sh | 115 +++ test/shell.d/agents-panel-test.sh | 2 + 9 files changed, 1024 insertions(+), 10 deletions(-) create mode 100755 bin/omarchy-agent-usage-antigravity create mode 100644 shell/plugins/agents/assets/antigravity-light.svg create mode 100644 shell/plugins/agents/assets/antigravity.svg create mode 100755 test/shell.d/agent-usage-antigravity-limits-test.sh create mode 100755 test/shell.d/agent-usage-antigravity-scanner-test.sh diff --git a/bin/omarchy-agent-usage-antigravity b/bin/omarchy-agent-usage-antigravity new file mode 100755 index 00000000000..4f205e1dd76 --- /dev/null +++ b/bin/omarchy-agent-usage-antigravity @@ -0,0 +1,737 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Antigravity usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Google Antigravity usage into one display-ready JSON record. + +Local stats come from Antigravity conversation transcripts and history logs +under ~/.gemini/antigravity-cli/. Limits, tier, and session window calculations +come from the Google Cloud Code API (v1internal:loadCodeAssist and +v1internal:retrieveUserQuotaSummary) with local rolling calculations as fallback. +The agents panel only ever reads the JSON this prints. +""" + +from __future__ import annotations + +import argparse +import base64 +import fcntl +import json +import os +import shutil +import sqlite3 +import subprocess +import sys +import tempfile +import time +import urllib.error +import urllib.request +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Any + +AGENT_ID = "antigravity" +AGENT_NAME = "Antigravity" +AUTH_HELP = "Run `agy` to authenticate." + +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 +PROBE_MIN_INTERVAL_SECONDS = 15 + + +def app_root() -> Path: + return Path(os.environ.get("AGY_DIR", Path.home() / ".gemini/antigravity-cli")).resolve() + + +def cache_root() -> Path: + root = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) / "omarchy" / "agent-usage" + root.mkdir(parents=True, exist_ok=True) + return root + + +def scan_cache_paths() -> tuple[Path, Path]: + root = cache_root() + return root / f"{AGENT_ID}-stats-cache.json", root / f"{AGENT_ID}-stats-cache.lock" + + +def read_fresh_json(path: Path, max_age_seconds: float) -> Any: + try: + stat = path.stat() + except OSError: + return None + age = time.time() - stat.st_mtime + if age > max_age_seconds: + return None + try: + with path.open("r", encoding="utf-8") as f: + return json.load(f) + except Exception: + return None + + +def write_json(path: Path, payload: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp_path = tempfile.mkstemp(prefix=f".{path.name}.", dir=str(path.parent)) + tmp = Path(tmp_path) + try: + with os.fdopen(fd, "w", encoding="utf-8") as f: + json.dump(payload, f, separators=(",", ":")) + f.write("\n") + tmp.chmod(0o644) + tmp.replace(path) + except BaseException: + tmp.unlink(missing_ok=True) + raise + + +def check_auth() -> tuple[bool, str, str, str, int]: + """Check authentication state via Secret Service keyring.""" + try: + raw = subprocess.check_output( + ["secret-tool", "lookup", "service", "gemini", "username", "antigravity"], + text=True, + stderr=subprocess.DEVNULL, + timeout=3, + ) + if raw.strip(): + data = json.loads(raw) + auth_method = data.get("auth_method", "consumer") + token_dict = data.get("token") or {} + access_token = str(token_dict.get("access_token") or data.get("access_token") or "") + expiry_raw = str(token_dict.get("expiry") or data.get("expiry") or "") + expires_at_ms = 0 + if expiry_raw: + try: + if expiry_raw.isdigit(): + expires_at_ms = int(expiry_raw) + if expires_at_ms < 1e12: + expires_at_ms *= 1000 + else: + dt = datetime.fromisoformat(expiry_raw.replace("Z", "+00:00")) + expires_at_ms = int(dt.timestamp() * 1000) + except Exception: + expires_at_ms = 0 + + id_token = data.get("id_token") + email = "" + if id_token: + parts = id_token.split(".") + if len(parts) >= 2: + padded = parts[1] + "=" * ((4 - len(parts[1]) % 4) % 4) + claims = json.loads(base64.urlsafe_b64decode(padded)) + email = claims.get("email", "") + return True, email, auth_method, access_token, expires_at_ms + except Exception: + pass + return False, "", "consumer", "", 0 + + +def normalize_tier_label(tier_name: str, tier_id: str = "") -> str: + tier_id_lower = tier_id.lower() + tier_name_lower = tier_name.lower() + if "ultra" in tier_id_lower or "ultra" in tier_name_lower: + return "Ultra" + if "pro" in tier_id_lower or "pro" in tier_name_lower: + return "Pro" + if "free" in tier_id_lower or "free" in tier_name_lower: + return "Free" + if tier_name.lower().startswith("google ai "): + return tier_name[10:].strip() + if tier_name: + return tier_name + if tier_id: + return tier_id.replace("-tier", "").capitalize() + return "Free" + + +def extract_tier(data: dict[str, Any], auth_method: str = "consumer") -> str: + paid_tier = data.get("paidTier") + if isinstance(paid_tier, dict): + t_name = str(paid_tier.get("name") or "").strip() + t_id = str(paid_tier.get("id") or "").strip() + if t_name or t_id: + return normalize_tier_label(t_name, t_id) + + current_tier = data.get("currentTier") + if isinstance(current_tier, dict): + c_id = str(current_tier.get("id") or "").strip() + c_name = str(current_tier.get("name") or "").strip() + if c_id or c_name: + return normalize_tier_label(c_name, c_id) + + if auth_method == "enterprise": + return "Enterprise" + if auth_method in ("cloud", "vertex"): + return "Cloud" + return "Free" + + +def normalize_reset_at(value: Any) -> str: + if value is None: + return "" + raw = str(value).strip() + if raw == "": + return "" + if raw.isdigit(): + ts = int(raw) + if ts < 1e12: + ts *= 1000 + try: + return datetime.fromtimestamp(ts / 1000, timezone.utc).isoformat() + except Exception: + return raw + try: + parsed = datetime.fromisoformat(raw.replace("Z", "+00:00")) + return parsed.isoformat() + except Exception: + return raw + + +def extract_limits(data: dict[str, Any]) -> list[dict[str, Any]]: + groups = data.get("groups") + if not isinstance(groups, list): + return [] + + gemini_group = None + other_groups = [] + for g in groups: + if not isinstance(g, dict): + continue + name = str(g.get("displayName") or "").lower() + if "gemini" in name or gemini_group is None: + gemini_group = g + else: + other_groups.append(g) + + limits = [] + # Primary group: gemini models + if gemini_group: + buckets = gemini_group.get("buckets") or [] + bucket_5h = None + bucket_weekly = None + for b in buckets: + if not isinstance(b, dict): + continue + w = str(b.get("window") or "").lower() + bid = str(b.get("bucketId") or "").lower() + dname = str(b.get("displayName") or "").lower() + if w == "5h" or "5h" in bid or "five" in dname: + bucket_5h = b + elif w == "weekly" or "weekly" in bid or "week" in dname: + bucket_weekly = b + + # Enforce Session (5-hour) first, then Weekly (7-day) + if bucket_5h: + rem = float(bucket_5h.get("remainingFraction", 1.0)) + used = round(max(0.0, min(1.0, 1.0 - rem)), 4) + limits.append({ + "label": "Session (5-hour)", + "title": "Session", + "percent": used, + "resetsAt": normalize_reset_at(bucket_5h.get("resetTime")), + }) + if bucket_weekly: + rem = float(bucket_weekly.get("remainingFraction", 1.0)) + used = round(max(0.0, min(1.0, 1.0 - rem)), 4) + limits.append({ + "label": "Weekly (7-day)", + "title": "Weekly", + "percent": used, + "resetsAt": normalize_reset_at(bucket_weekly.get("resetTime")), + }) + + # Scoped limits for other groups (e.g. Claude and GPT models) + for g in other_groups: + g_name = str(g.get("displayName") or "").strip() + short_name = "Claude/GPT" if "claude" in g_name.lower() else (g_name or "3P") + for b in (g.get("buckets") or []): + if not isinstance(b, dict): + continue + rem = float(b.get("remainingFraction", 1.0)) + used = round(max(0.0, min(1.0, 1.0 - rem)), 4) + # Only show secondary model limits if there has been usage + if used > 0: + w = str(b.get("window") or "").lower() + w_title = "Weekly" if "week" in w else "Session" + title = f"{short_name} {w_title}" + limits.append({ + "label": title, + "title": title, + "percent": used, + "resetsAt": normalize_reset_at(b.get("resetTime")), + }) + + return limits + + +def probe_antigravity(access_token: str, endpoint_base: str = "") -> dict[str, Any]: + if not endpoint_base: + endpoint_base = os.environ.get("CLOUD_CODE_URL", "https://daily-cloudcode-pa.googleapis.com").rstrip("/") + + headers = { + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json", + "User-Agent": "antigravity", + } + + result: dict[str, Any] = { + "ok": False, + "tierLabel": "", + "limits": [], + "helpText": "", + "rateLimited": False, + "transport": False, + } + + tier = "" + try: + req = urllib.request.Request(f"{endpoint_base}/v1internal:loadCodeAssist", data=b"{}", headers=headers) + with urllib.request.urlopen(req, timeout=8) as resp: + payload = json.loads(resp.read().decode("utf-8", errors="replace")) + tier = extract_tier(payload) + except urllib.error.HTTPError as err: + if err.code == 429: + result["rateLimited"] = True + elif err.code == 401: + result["helpText"] = "Antigravity authentication expired. Run `agy` to authenticate." + return result + except Exception: + result["transport"] = True + result["helpText"] = "Could not reach Antigravity service." + + limits = [] + try: + req = urllib.request.Request(f"{endpoint_base}/v1internal:retrieveUserQuotaSummary", data=b"{}", headers=headers) + with urllib.request.urlopen(req, timeout=8) as resp: + payload = json.loads(resp.read().decode("utf-8", errors="replace")) + limits = extract_limits(payload) + except urllib.error.HTTPError as err: + if err.code == 429: + result["rateLimited"] = True + elif err.code == 401: + result["helpText"] = "Antigravity authentication expired. Run `agy` to authenticate." + return result + except Exception: + result["transport"] = True + + result["ok"] = bool(tier or limits) + result["tierLabel"] = tier + result["limits"] = limits + return result + + +def limit_window_open(entry: dict[str, Any], now: datetime) -> bool: + raw = str(entry.get("resetsAt") or "") + if raw == "": + return True + try: + resets_at = datetime.fromisoformat(raw.replace("Z", "+00:00")) + except Exception: + return True + if resets_at.tzinfo is None: + resets_at = resets_at.replace(tzinfo=timezone.utc) + return resets_at > now + + +def usable_cached_limits(cached: dict[str, Any]) -> list[dict[str, Any]]: + entries = cached.get("limits") + if not isinstance(entries, list): + return [] + now = datetime.now(timezone.utc) + return [entry for entry in entries if isinstance(entry, dict) and limit_window_open(entry, now)] + + +def detect_tier_label( + root_dir: Path, + auth_ok: bool, + email: str, + auth_method: str = "consumer", + probed_tier: str = "", + cached_tier: str = "", +) -> str: + if not auth_ok: + return "" + + if os.environ.get("AGY_TIER"): + return os.environ["AGY_TIER"] + + cfg_file = Path.home() / ".config" / "omarchy" / "agents" / "antigravity.json" + if cfg_file.exists(): + try: + with cfg_file.open("r", encoding="utf-8") as f: + data = json.load(f) + if data.get("tier"): + return str(data["tier"]) + except Exception: + pass + + if probed_tier: + return probed_tier + + if cached_tier: + return cached_tier + + if auth_method == "enterprise": + return "Enterprise" + if auth_method in ("cloud", "vertex"): + return "Cloud" + + return "Pro" + + +def collect_local_stats(root_dir: Path, now: datetime) -> dict[str, Any]: + today_str = now.strftime("%Y-%m-%d") + recent_dates = [(now - timedelta(days=i)).strftime("%Y-%m-%d") for i in range(6, -1, -1)] + recent = {d: {"date": d, "messageCount": 0} for d in recent_dates} + + history_file = root_dir / "history.jsonl" + brain_dir = root_dir / "brain" + + today_prompts = 0 + total_prompts = 0 + active_dates: set[str] = set() + today_sessions: set[str] = set() + total_sessions: set[str] = set() + + if history_file.exists(): + try: + with history_file.open("r", encoding="utf-8", errors="ignore") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + entry = json.loads(line) + ts = entry.get("timestamp") + if not ts: + continue + try: + if isinstance(ts, (int, float)): + dt = datetime.fromtimestamp(ts / 1000.0, timezone.utc) + else: + dt = datetime.fromisoformat(str(ts).replace("Z", "+00:00")) + except Exception: + continue + d_str = dt.strftime("%Y-%m-%d") + active_dates.add(d_str) + total_prompts += 1 + conv_id = entry.get("conversationId") or entry.get("conversation_id") + if conv_id: + total_sessions.add(conv_id) + if d_str == today_str: + today_prompts += 1 + if conv_id: + today_sessions.add(conv_id) + except Exception: + continue + except Exception: + pass + + # Fallback to conversation summaries if history.jsonl is sparse + summaries_db = root_dir / "conversation_summaries.db" + if summaries_db.exists(): + try: + conn = sqlite3.connect(f"file:{summaries_db}?mode=ro", uri=True) + cur = conn.cursor() + for row in cur.execute("SELECT conversation_id, last_modified_time FROM conversation_summaries"): + cid, lmt = row[0], str(row[1] or "") + if cid: + total_sessions.add(cid) + if lmt.startswith(today_str): + today_sessions.add(cid) + conn.close() + except Exception: + pass + + # Compute token counts across conversation transcripts + model_usage: dict[str, dict[str, int]] = {} + today_tokens = 0 + today_tokens_by_model: dict[str, int] = {} + + if brain_dir.exists(): + try: + for conv_folder in brain_dir.iterdir(): + if not conv_folder.is_dir(): + continue + tr_file = conv_folder / ".system_generated/logs/transcript.jsonl" + if not tr_file.exists(): + continue + try: + with tr_file.open("r", encoding="utf-8", errors="ignore") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + step = json.loads(line) + m_name = step.get("model") or "gemini-3.8-flash-high" + if m_name not in model_usage: + model_usage[m_name] = { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + } + + step_usage = step.get("usage") + is_model = step.get("source") == "MODEL" or step.get("type") == "PLANNER_RESPONSE" + if isinstance(step_usage, dict): + inp = int(step_usage.get("prompt_token_count") or step_usage.get("input_tokens") or 0) + out = int(step_usage.get("candidates_token_count") or step_usage.get("output_tokens") or 0) + tok = inp + out + model_usage[m_name]["inputTokens"] += inp + model_usage[m_name]["outputTokens"] += out + else: + c_text = step.get("content") or "" + t_text = step.get("thinking") or "" + tok = max(1, (len(c_text) + len(t_text)) // 4) + if is_model: + model_usage[m_name]["outputTokens"] += tok + else: + model_usage[m_name]["inputTokens"] += tok + + created = step.get("created_at") + is_today = False + if created: + step_d = created[:10] + if step_d in recent: + recent[step_d]["messageCount"] += tok + if step_d == today_str: + is_today = True + else: + is_today = True + + if is_today: + today_tokens += tok + today_tokens_by_model[m_name] = today_tokens_by_model.get(m_name, 0) + tok + except Exception: + continue + except Exception: + continue + except Exception: + pass + + return { + "todayPrompts": today_prompts, + "todaySessions": len(today_sessions), + "todayTotalTokens": today_tokens, + "todayTokensByModel": today_tokens_by_model, + "recentDays": [recent[d] for d in recent_dates], + "totalPrompts": total_prompts, + "totalSessions": len(total_sessions), + "activeDays": len(active_dates), + "activeDates": sorted(list(active_dates)), + "modelUsage": model_usage, + } + + +def calculate_limits(root_dir: Path, now: datetime) -> list[dict[str, Any]]: + """Fallback calculation of rolling limits from local history files.""" + history_file = root_dir / "history.jsonl" + five_hours_ago_ts = int((now - timedelta(hours=5)).timestamp() * 1000) + + prompts_5h = 0 + oldest_5h_ts = None + prompts_7d = 0 + + seven_days_ago_ts = int((now - timedelta(days=7)).timestamp() * 1000) + + if history_file.exists(): + try: + with history_file.open("r", encoding="utf-8", errors="ignore") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + entry = json.loads(line) + ts = entry.get("timestamp") + if not ts: + continue + if ts >= seven_days_ago_ts: + prompts_7d += 1 + if ts >= five_hours_ago_ts: + prompts_5h += 1 + if oldest_5h_ts is None or ts < oldest_5h_ts: + oldest_5h_ts = ts + except Exception: + continue + except Exception: + pass + + # Rolling 5h reset + if oldest_5h_ts: + reset_5h = datetime.fromtimestamp(oldest_5h_ts / 1000.0, timezone.utc) + timedelta(hours=5) + else: + reset_5h = now + timedelta(hours=5) + + # Check if a 429 / resource exhausted error happened recently + cli_log = root_dir / "cli.log" + rate_limited = False + if cli_log.exists(): + try: + stat = cli_log.stat() + if (time.time() - stat.st_mtime) < 18000: + with cli_log.open("r", encoding="utf-8", errors="ignore") as f: + content = f.read()[-8192:] + if "RESOURCE_EXHAUSTED" in content or "rate limit exceeded" in content.lower(): + rate_limited = True + except Exception: + pass + + session_pct = 1.0 if rate_limited else round(min(1.0, prompts_5h / 50.0), 2) + weekly_pct = round(min(1.0, prompts_7d / 250.0), 2) + + return [ + { + "label": "Session (5-hour)", + "title": "Session", + "percent": session_pct, + "resetsAt": reset_5h.strftime("%Y-%m-%dT%H:%M:%S+00:00"), + }, + { + "label": "Weekly (7-day)", + "title": "Weekly", + "percent": weekly_pct, + "resetsAt": (now + timedelta(days=7)).strftime("%Y-%m-%dT%H:%M:%S+00:00"), + }, + ] + + +def collect_limits( + root_dir: Path, + now: datetime, + auth_ok: bool, + access_token: str, + expires_at_ms: int, + auth_method: str, + force: bool, +) -> tuple[list[dict[str, Any]], str, str, str]: + """Collect limits via API probe, falling back to cached or rolling local stats.""" + limits_cache = cache_root() / f"{AGENT_ID}-limits.json" + cached = read_fresh_json(limits_cache, float("inf")) or {} + fallback = usable_cached_limits(cached) + cached_tier = str(cached.get("tierLabel") or "") + + if not auth_ok: + return [], "", "Sign in to see limits", AUTH_HELP + + # If access token is empty, we cannot probe the remote endpoint + if not access_token: + limits = fallback if fallback else calculate_limits(root_dir, now) + return limits, cached_tier, "", AUTH_HELP + + # Check reuse window for limits probe (PROBE_MIN_INTERVAL_SECONDS) + fetched_at = float(cached.get("fetchedAtMs") or 0) / 1000.0 + if fallback and not force and (time.time() - fetched_at) < PROBE_MIN_INTERVAL_SECONDS: + return fallback, cached_tier, "", AUTH_HELP + + # Perform live probe + probe = probe_antigravity(access_token) + if probe["ok"]: + probed_limits = probe["limits"] + probed_tier = probe["tierLabel"] + write_json( + limits_cache, + { + "fetchedAtMs": round(time.time() * 1000), + "limits": probed_limits, + "tierLabel": probed_tier, + }, + ) + return probed_limits, probed_tier, "", AUTH_HELP + + # Probe failed + help_text = probe["helpText"] if probe.get("helpText") else AUTH_HELP + usage_status = "" + if probe.get("rateLimited"): + usage_status = "Antigravity limits rate limited" + elif probe.get("transport"): + usage_status = "" + + if fallback: + return fallback, cached_tier, usage_status, help_text + + # Fallback to local rolling calculation + local_limits = calculate_limits(root_dir, now) + return local_limits, cached_tier, usage_status, help_text + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--force", action="store_true", help="Bypass local caches and force a fresh scan") + parser.add_argument("--limits-only", action="store_true", help="Recompute only rate limits if cache is fresh") + args = parser.parse_args() + + root_dir = app_root() + now = datetime.now(timezone.utc) + today_str = now.strftime("%Y-%m-%d") + + auth_ok, email, auth_method, access_token, expires_at_ms = check_auth() + + cache_file, lock_file = scan_cache_paths() + stats = None + + if not args.force: + max_age = LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS + cached = read_fresh_json(cache_file, max_age) + if isinstance(cached, dict) and cached.get("scanDate") == today_str: + stats = cached.get("stats") + + if stats is None: + try: + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + stats = collect_local_stats(root_dir, now) + write_json(cache_file, {"schemaVersion": 1, "scanDate": today_str, "stats": stats}) + except Exception as exc: + print(f"omarchy-agent-usage-antigravity: cache lock failed ({exc}), scanning directly", file=sys.stderr) + stats = collect_local_stats(root_dir, now) + + limits, probed_tier, usage_status_text, auth_help_text = collect_limits( + root_dir, + now, + auth_ok, + access_token, + expires_at_ms, + auth_method, + args.force, + ) + + cached_limits_file = cache_root() / f"{AGENT_ID}-limits.json" + cached_payload = read_fresh_json(cached_limits_file, float("inf")) or {} + cached_tier = str(cached_payload.get("tierLabel") or "") + tier_label = detect_tier_label(root_dir, auth_ok, email, auth_method, probed_tier, cached_tier) + + total_prompts = stats.get("totalPrompts", 0) + ready = auth_ok or total_prompts > 0 or len(limits) > 0 + has_local_stats = total_prompts > 0 + + record = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": now.strftime("%Y-%m-%dT%H:%M:%S.%f+00:00"), + "ready": ready, + "hasLocalStats": has_local_stats, + "todayPrompts": stats.get("todayPrompts", 0), + "todaySessions": stats.get("todaySessions", 0), + "todayTotalTokens": stats.get("todayTotalTokens", 0), + "todayTokensByModel": stats.get("todayTokensByModel", {}), + "recentDays": stats.get("recentDays", []), + "totalPrompts": total_prompts, + "totalSessions": stats.get("totalSessions", 0), + "activeDays": stats.get("activeDays", 0), + "activeDates": stats.get("activeDates", []), + "modelUsage": stats.get("modelUsage", {}), + "limits": limits, + "tierLabel": tier_label, + "usageStatusText": usage_status_text, + "authHelpText": auth_help_text, + } + + print(json.dumps(record, separators=(",", ":"))) + + +if __name__ == "__main__": + main() diff --git a/shell/plugins/agents/Panel.qml b/shell/plugins/agents/Panel.qml index f4ecdd9a09b..620b1c5cbe2 100644 --- a/shell/plugins/agents/Panel.qml +++ b/shell/plugins/agents/Panel.qml @@ -23,9 +23,35 @@ Panel { // provider whose first scan lands while the panel is open would otherwise // shift the list underneath you and swap out what you were reading. property string selectedProviderId: "" + property string defaultAgentRaw: "" + + FileView { + path: (Quickshell.env("HOME") || "") + "/.config/omarchy/defaults/agent" + watchChanges: true + printErrors: false + onLoaded: root.defaultAgentRaw = String(text() || "").trim() + onFileChanged: reload() + } + + function normalizeAgentId(raw) { + var id = String(raw || "").trim().toLowerCase() + if (id === "agy" || id === "antigravity" || id === "antigravity-cli" || id === "gemini" || id === "gemini-cli") + return "antigravity" + if (id === "claude-code") + return "claude" + return id + } + readonly property int providerIndex: { - for (var i = 0; i < providers.length; i++) - if (providers[i].providerId === selectedProviderId) return i + if (selectedProviderId !== "") { + for (var i = 0; i < providers.length; i++) + if (providers[i].providerId === selectedProviderId) return i + } + var defId = normalizeAgentId(defaultAgentRaw) + if (defId !== "") { + for (var j = 0; j < providers.length; j++) + if (providers[j].providerId === defId) return j + } return 0 } readonly property var provider: providers.length > 0 ? providers[providerIndex] : null diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 51ad4e43c2d..4210e453572 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -52,17 +52,12 @@ light surfaces — and the bar glyph stands in when there is none. | Collector | Limits | Local stats | |---|---|---| +| `antigravity` | Google Cloud Code API (`loadCodeAssist` and `retrieveUserQuotaSummary`) with rolling session/weekly calculation fallback | Transcripts and history logs under `~/.gemini/antigravity-cli/` (`AGY_DIR`) | | `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 | -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 -`CLAUDE_CONFIG_DIR`, Codex via `CODEX_HOME`. Fireworks reads -`FIREWORKS_API_KEY` and `FIREWORKS_ACCOUNT_ID` first, then -`~/.fireworks/auth.ini` (which `firectl set-api-key` creates), then the key -opencode stores in `~/.local/share/opencode/auth.json` when Fireworks is -signed in there. +Antigravity checks authentication via the Secret Service keyring (`gemini`/`antigravity`), queries Google's Cloud Code endpoints for account tier and quota limits (with rolling local stats as fallback), and honors custom app roots via `AGY_DIR` or tier overrides via `AGY_TIER` / `~/.config/omarchy/agents/antigravity.json`. 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 `CLAUDE_CONFIG_DIR`, Codex via `CODEX_HOME`. Fireworks reads `FIREWORKS_API_KEY` and `FIREWORKS_ACCOUNT_ID` first, then `~/.fireworks/auth.ini` (which `firectl set-api-key` creates), then the key opencode stores in `~/.local/share/opencode/auth.json` when Fireworks is signed in there. ### Fireworks balance diff --git a/shell/plugins/agents/assets/antigravity-light.svg b/shell/plugins/agents/assets/antigravity-light.svg new file mode 100644 index 00000000000..67ced694468 --- /dev/null +++ b/shell/plugins/agents/assets/antigravity-light.svg @@ -0,0 +1 @@ +Antigravity diff --git a/shell/plugins/agents/assets/antigravity.svg b/shell/plugins/agents/assets/antigravity.svg new file mode 100644 index 00000000000..3743be8edf4 --- /dev/null +++ b/shell/plugins/agents/assets/antigravity.svg @@ -0,0 +1 @@ +Antigravity diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index 1743a798214..305b2ca8c4c 100644 --- a/shell/plugins/agents/manifest.json +++ b/shell/plugins/agents/manifest.json @@ -5,7 +5,7 @@ "version": "1.0.0", "author": "Omarchy", "license": "MIT", - "description": "Claude Code, Codex, and Fireworks usage, limits, and pace in a native Omarchy bar panel.", + "description": "Antigravity, Claude Code, Codex, and Fireworks usage, limits, and pace in a native Omarchy bar panel.", "kinds": ["bar-widget"], "activation": "on-demand", "entryPoints": { @@ -19,6 +19,7 @@ "allowMultiple": false, "defaults": { "providers": { + "antigravity": { "enabled": true }, "claude": { "enabled": true }, "codex": { "enabled": true }, "fireworks": { "enabled": true } diff --git a/test/shell.d/agent-usage-antigravity-limits-test.sh b/test/shell.d/agent-usage-antigravity-limits-test.sh new file mode 100755 index 00000000000..89c0dbdc945 --- /dev/null +++ b/test/shell.d/agent-usage-antigravity-limits-test.sh @@ -0,0 +1,136 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +COLLECTOR="$ROOT/bin/omarchy-agent-usage-antigravity" + +# probe_antigravity reaches Google's Cloud Code endpoints, so the readers that +# interpret the answers are exercised with recorded payloads standing in. + +test_tier() { + COLLECTOR="$COLLECTOR" PAYLOAD="$1" python3 - <<'PY' +import importlib.machinery, importlib.util, 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) + +payload = json.loads(os.environ["PAYLOAD"]) +print(collector.extract_tier(payload)) +PY +} + +ultra_tier=$(test_tier '{ + "currentTier": { "id": "free-tier", "name": "Antigravity" }, + "paidTier": { "id": "g1-ultra-tier", "name": "Google AI Ultra" } +}') +[[ $ultra_tier == "Ultra" ]] || + fail "Antigravity collector extracts Ultra tier from paidTier" "$ultra_tier" +pass "Antigravity collector extracts Ultra tier from paidTier" + +pro_tier=$(test_tier '{ + "currentTier": { "id": "free-tier", "name": "Antigravity" }, + "paidTier": { "id": "g1-pro-tier", "name": "Google AI Pro" } +}') +[[ $pro_tier == "Pro" ]] || + fail "Antigravity collector extracts Pro tier from paidTier" "$pro_tier" +pass "Antigravity collector extracts Pro tier from paidTier" + +free_tier=$(test_tier '{ + "currentTier": { "id": "free-tier", "name": "Antigravity" } +}') +[[ $free_tier == "Free" ]] || + fail "Antigravity collector extracts Free tier when no paidTier" "$free_tier" +pass "Antigravity collector extracts Free tier when no paidTier" + +# Test quota extraction +test_limits() { + COLLECTOR="$COLLECTOR" PAYLOAD="$1" python3 - <<'PY' +import importlib.machinery, importlib.util, 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) + +payload = json.loads(os.environ["PAYLOAD"]) +print(json.dumps(collector.extract_limits(payload))) +PY +} + +quota_payload='{ + "groups": [ + { + "displayName": "Gemini Models", + "buckets": [ + { + "bucketId": "gemini-weekly", + "displayName": "Weekly Limit Remaining", + "window": "weekly", + "resetTime": "2026-09-17T18:19:42Z", + "remainingFraction": 0.95 + }, + { + "bucketId": "gemini-5h", + "displayName": "Five Hour Limit Remaining", + "window": "5h", + "resetTime": "2026-09-15T19:31:46Z", + "remainingFraction": 0.90 + } + ] + }, + { + "displayName": "Claude and GPT models", + "buckets": [ + { + "bucketId": "3p-weekly", + "displayName": "Weekly Limit Remaining", + "window": "weekly", + "resetTime": "2026-09-18T12:56:14Z", + "remainingFraction": 0.98 + }, + { + "bucketId": "3p-5h", + "displayName": "Five Hour Limit Remaining", + "window": "5h", + "resetTime": "2026-09-15T22:16:22Z", + "remainingFraction": 1.0 + } + ] + } + ] +}' + +limits=$(test_limits "$quota_payload") + +[[ $(jq -r '.[0].label' <<<"$limits") == "Session (5-hour)" ]] || + fail "Antigravity collector places Session window first" "$limits" +pass "Antigravity collector places Session window first" + +[[ $(jq -r '.[0].percent' <<<"$limits") == "0.1" ]] || + fail "Antigravity collector calculates used percentage from remainingFraction" "$limits" +pass "Antigravity collector calculates used percentage from remainingFraction" + +[[ $(jq -r '.[1].label' <<<"$limits") == "Weekly (7-day)" ]] || + fail "Antigravity collector places Weekly window second" "$limits" +pass "Antigravity collector places Weekly window second" + +[[ $(jq -r '.[1].percent' <<<"$limits") == "0.05" ]] || + fail "Antigravity collector calculates weekly used percentage" "$limits" +pass "Antigravity collector calculates weekly used percentage" + +[[ $(jq -r '.[2].title' <<<"$limits") == "Claude/GPT Weekly" ]] || + fail "Antigravity collector includes used scoped 3P limits" "$limits" +pass "Antigravity collector includes used scoped 3P limits" + +[[ $(jq -r '.[2].percent' <<<"$limits") == "0.02" ]] || + fail "Antigravity collector calculates 3P used percentage" "$limits" +pass "Antigravity collector calculates 3P used percentage" + +[[ $(jq -r 'length' <<<"$limits") == "3" ]] || + fail "Antigravity collector omits unused 3P session limit" "$limits" +pass "Antigravity collector omits unused 3P session limit" diff --git a/test/shell.d/agent-usage-antigravity-scanner-test.sh b/test/shell.d/agent-usage-antigravity-scanner-test.sh new file mode 100755 index 00000000000..48f6fed5119 --- /dev/null +++ b/test/shell.d/agent-usage-antigravity-scanner-test.sh @@ -0,0 +1,115 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +TEST_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME"' EXIT + +AGY_DIR="$TEST_HOME/.gemini/antigravity-cli" +mkdir -p "$AGY_DIR/brain/test-conv/.system_generated/logs" "$TEST_HOME/bin" + +# Mock secret-tool to simulate authenticated state +cat >"$TEST_HOME/bin/secret-tool" <<'EOF' +#!/bin/bash +if [[ "$*" == *"lookup service gemini username antigravity"* ]]; then + echo '{"auth_method":"consumer","id_token":"header.eyJlbWFpbCI6InRlc3RAZ21haWwuY29tIn0.signature"}' + exit 0 +fi +exit 1 +EOF +chmod +x "$TEST_HOME/bin/secret-tool" + +# Write mock history and transcript +now_date=$(date +%Y-%m-%d) +timestamp="$(date -u +%Y-%m-%dT%H:%M:%SZ)" + +cat >"$AGY_DIR/history.jsonl" <"$AGY_DIR/brain/test-conv/.system_generated/logs/transcript.jsonl" <"$TEST_HOME/.config/omarchy/agents/antigravity.json" <<'EOF' +{ + "tier": "Enterprise" +} +EOF + +cfg_tier_result=$(HOME="$TEST_HOME" AGY_DIR="$AGY_DIR" XDG_CACHE_HOME="$TEST_HOME/.cache" PATH="$TEST_HOME/bin:$PATH" \ + "$ROOT/bin/omarchy-agent-usage-antigravity" --force) + +[[ $(jq -r '.tierLabel' <<<"$cfg_tier_result") == "Enterprise" ]] || + fail "Antigravity collector honors config file tier override" "$cfg_tier_result" +pass "Antigravity collector honors config file tier override" + +# Test unauthenticated state +cat >"$TEST_HOME/bin/secret-tool" <<'EOF' +#!/bin/bash +exit 1 +EOF + +unauth_result=$(HOME="$TEST_HOME" AGY_DIR="$AGY_DIR" XDG_CACHE_HOME="$TEST_HOME/.cache" PATH="$TEST_HOME/bin:$PATH" \ + "$ROOT/bin/omarchy-agent-usage-antigravity" --force) + +[[ $(jq -r '.tierLabel' <<<"$unauth_result") == "" ]] || + fail "Antigravity collector has empty tier when unauthenticated" "$unauth_result" +pass "Antigravity collector has empty tier when unauthenticated" + +[[ $(jq -r '.usageStatusText' <<<"$unauth_result") == "Sign in to see limits" ]] || + fail "Antigravity collector indicates sign in status when unauthenticated" "$unauth_result" +pass "Antigravity collector indicates sign in status when unauthenticated" diff --git a/test/shell.d/agents-panel-test.sh b/test/shell.d/agents-panel-test.sh index c09bab8793a..17405c068fe 100755 --- a/test/shell.d/agents-panel-test.sh +++ b/test/shell.d/agents-panel-test.sh @@ -14,4 +14,6 @@ assert(/if \(buttonCode === Qt\.RightButton\) root\.launchAgent\(\)/.test(panelS assert(/else if \(buttonCode === Qt\.MiddleButton\) root\.selectProvider\(root\.providerIndex \+ 1\)/.test(panelSource), 'agents middle click still advances the subscription') assert(/else root\.toggle\(\)/.test(panelSource), 'agents left click still toggles the panel') assert(!/if \(buttonCode === Qt\.RightButton\) root\.refreshNow\(\)/.test(panelSource), 'agents right click no longer refreshes') +assert(/normalizeAgentId/.test(panelSource), 'agents panel normalizes default agent name') +assert(/omarchy\/defaults\/agent/.test(panelSource), 'agents panel watches system default agent setting') JS From a38ac1d1e2e654f7232e78c2cbf6dff71e8ad379 Mon Sep 17 00:00:00 2001 From: steph4n <284007938+steph4n-gh@users.noreply.github.com> Date: Thu, 1 Oct 2026 22:55:27 -0400 Subject: [PATCH 2/7] Read verified Antigravity quotas and report failures honestly --- bin/omarchy-agent-usage-antigravity | 973 ++++++------------ manual/17-ai.md | 2 +- shell/plugins/agents/README.md | 13 +- .../agent-usage-antigravity-limits-test.sh | 237 ++--- .../agent-usage-antigravity-scanner-test.sh | 287 ++++-- test/shell.d/fixtures/antigravity/README.md | 5 + test/shell.d/fixtures/antigravity/usage.json | 25 + 7 files changed, 629 insertions(+), 913 deletions(-) create mode 100644 test/shell.d/fixtures/antigravity/README.md create mode 100644 test/shell.d/fixtures/antigravity/usage.json diff --git a/bin/omarchy-agent-usage-antigravity b/bin/omarchy-agent-usage-antigravity index 4f205e1dd76..34325a857ca 100755 --- a/bin/omarchy-agent-usage-antigravity +++ b/bin/omarchy-agent-usage-antigravity @@ -2,735 +2,366 @@ # omarchy:summary=Print the Antigravity usage record as JSON # omarchy:args=[--force] [--limits-only] # omarchy:hidden=true -"""Collect Google Antigravity usage into one display-ready JSON record. +"""Read Antigravity CLI quotas and local prompt/session history. -Local stats come from Antigravity conversation transcripts and history logs -under ~/.gemini/antigravity-cli/. Limits, tier, and session window calculations -come from the Google Cloud Code API (v1internal:loadCodeAssist and -v1internal:retrieveUserQuotaSummary) with local rolling calculations as fallback. -The agents panel only ever reads the JSON this prints. +agy 1.1.11 introduced read-only /usage JSON output. Let the CLI handle +credentials, refresh and quota requests; older versions must never receive +this prompt because they may send it to a model instead. Transcripts do not +record token counts or model IDs, so those fields are deliberately empty. """ -from __future__ import annotations - import argparse -import base64 import fcntl +import hashlib import json +import math import os +import re import shutil +import signal import sqlite3 import subprocess -import sys import tempfile import time -import urllib.error -import urllib.request +from collections import Counter from datetime import datetime, timedelta, timezone from pathlib import Path -from typing import Any - -AGENT_ID = "antigravity" -AGENT_NAME = "Antigravity" -AUTH_HELP = "Run `agy` to authenticate." SCAN_REUSE_SECONDS = 20 LIMITS_ONLY_REUSE_SECONDS = 900 -PROBE_MIN_INTERVAL_SECONDS = 15 +PROBE_REUSE_SECONDS = 60 +CACHE_VERSION = 2 +AUTH_HELP = "Run `agy` to check your sign-in, then refresh the panel." -def app_root() -> Path: - return Path(os.environ.get("AGY_DIR", Path.home() / ".gemini/antigravity-cli")).resolve() +def app_root(): + return Path(os.environ.get("AGY_DIR") or Path.home() / ".gemini/antigravity-cli").expanduser().resolve() -def cache_root() -> Path: - root = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) / "omarchy" / "agent-usage" - root.mkdir(parents=True, exist_ok=True) - return root +def cache_root(): + return Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache") / "omarchy/agent-usage" -def scan_cache_paths() -> tuple[Path, Path]: - root = cache_root() - return root / f"{AGENT_ID}-stats-cache.json", root / f"{AGENT_ID}-stats-cache.lock" +def read_json(path): + try: + data = json.loads(path.read_text(encoding="utf-8")) + return data if isinstance(data, dict) else {} + except (OSError, ValueError): + return {} -def read_fresh_json(path: Path, max_age_seconds: float) -> Any: +def write_json(path, payload): + path.parent.mkdir(parents=True, exist_ok=True) + fd, name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) try: - stat = path.stat() - except OSError: - return None - age = time.time() - stat.st_mtime - if age > max_age_seconds: - return None + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(payload, handle, separators=(",", ":"), allow_nan=False) + handle.write("\n") + os.replace(name, path) + finally: + Path(name).unlink(missing_ok=True) + + +def parse_time(value): try: - with path.open("r", encoding="utf-8") as f: - return json.load(f) - except Exception: + if isinstance(value, (int, float)) and not isinstance(value, bool): + return datetime.fromtimestamp(value / 1000, timezone.utc) + parsed = datetime.fromisoformat(str(value or "").replace("Z", "+00:00")) + return parsed if parsed.tzinfo else parsed.replace(tzinfo=timezone.utc) + except (ValueError, TypeError, OverflowError, OSError): return None -def write_json(path: Path, payload: Any) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - fd, tmp_path = tempfile.mkstemp(prefix=f".{path.name}.", dir=str(path.parent)) - tmp = Path(tmp_path) - try: - with os.fdopen(fd, "w", encoding="utf-8") as f: - json.dump(payload, f, separators=(",", ":")) - f.write("\n") - tmp.chmod(0o644) - tmp.replace(path) - except BaseException: - tmp.unlink(missing_ok=True) - raise - - -def check_auth() -> tuple[bool, str, str, str, int]: - """Check authentication state via Secret Service keyring.""" +def json_lines(path): try: - raw = subprocess.check_output( - ["secret-tool", "lookup", "service", "gemini", "username", "antigravity"], - text=True, - stderr=subprocess.DEVNULL, - timeout=3, - ) - if raw.strip(): - data = json.loads(raw) - auth_method = data.get("auth_method", "consumer") - token_dict = data.get("token") or {} - access_token = str(token_dict.get("access_token") or data.get("access_token") or "") - expiry_raw = str(token_dict.get("expiry") or data.get("expiry") or "") - expires_at_ms = 0 - if expiry_raw: + with path.open(encoding="utf-8", errors="replace") as handle: + for line in handle: try: - if expiry_raw.isdigit(): - expires_at_ms = int(expiry_raw) - if expires_at_ms < 1e12: - expires_at_ms *= 1000 - else: - dt = datetime.fromisoformat(expiry_raw.replace("Z", "+00:00")) - expires_at_ms = int(dt.timestamp() * 1000) - except Exception: - expires_at_ms = 0 - - id_token = data.get("id_token") - email = "" - if id_token: - parts = id_token.split(".") - if len(parts) >= 2: - padded = parts[1] + "=" * ((4 - len(parts[1]) % 4) % 4) - claims = json.loads(base64.urlsafe_b64decode(padded)) - email = claims.get("email", "") - return True, email, auth_method, access_token, expires_at_ms - except Exception: - pass - return False, "", "consumer", "", 0 - - -def normalize_tier_label(tier_name: str, tier_id: str = "") -> str: - tier_id_lower = tier_id.lower() - tier_name_lower = tier_name.lower() - if "ultra" in tier_id_lower or "ultra" in tier_name_lower: - return "Ultra" - if "pro" in tier_id_lower or "pro" in tier_name_lower: - return "Pro" - if "free" in tier_id_lower or "free" in tier_name_lower: - return "Free" - if tier_name.lower().startswith("google ai "): - return tier_name[10:].strip() - if tier_name: - return tier_name - if tier_id: - return tier_id.replace("-tier", "").capitalize() - return "Free" - - -def extract_tier(data: dict[str, Any], auth_method: str = "consumer") -> str: - paid_tier = data.get("paidTier") - if isinstance(paid_tier, dict): - t_name = str(paid_tier.get("name") or "").strip() - t_id = str(paid_tier.get("id") or "").strip() - if t_name or t_id: - return normalize_tier_label(t_name, t_id) - - current_tier = data.get("currentTier") - if isinstance(current_tier, dict): - c_id = str(current_tier.get("id") or "").strip() - c_name = str(current_tier.get("name") or "").strip() - if c_id or c_name: - return normalize_tier_label(c_name, c_id) - - if auth_method == "enterprise": - return "Enterprise" - if auth_method in ("cloud", "vertex"): - return "Cloud" - return "Free" - - -def normalize_reset_at(value: Any) -> str: - if value is None: - return "" - raw = str(value).strip() - if raw == "": - return "" - if raw.isdigit(): - ts = int(raw) - if ts < 1e12: - ts *= 1000 - try: - return datetime.fromtimestamp(ts / 1000, timezone.utc).isoformat() - except Exception: - return raw - try: - parsed = datetime.fromisoformat(raw.replace("Z", "+00:00")) - return parsed.isoformat() - except Exception: - return raw - - -def extract_limits(data: dict[str, Any]) -> list[dict[str, Any]]: - groups = data.get("groups") - if not isinstance(groups, list): - return [] - - gemini_group = None - other_groups = [] - for g in groups: - if not isinstance(g, dict): + entry = json.loads(line) + if isinstance(entry, dict): + yield entry + except ValueError: + continue + except OSError: + return + + +def collect_local_stats(root, now): + today = now.astimezone().date().isoformat() + days = [(now.astimezone().date() - timedelta(days=i)).isoformat() for i in range(6, -1, -1)] + history_prompts = Counter() + transcript_prompts = Counter() + active_dates = set() + sessions = set() + today_sessions = set() + + def activity(value, session=None): + parsed = parse_time(value) + if parsed is None or parsed > now: + return None + day = parsed.astimezone().date().isoformat() + active_dates.add(day) + if isinstance(session, str) and session: + sessions.add(session) + if day == today: + today_sessions.add(session) + return day + + for entry in json_lines(root / "history.jsonl"): + # History includes slash commands; those are not prompts to an agent. + display = entry.get("display", entry.get("prompt", "")) + if isinstance(display, str) and display.lstrip().startswith("/"): continue - name = str(g.get("displayName") or "").lower() - if "gemini" in name or gemini_group is None: - gemini_group = g - else: - other_groups.append(g) - - limits = [] - # Primary group: gemini models - if gemini_group: - buckets = gemini_group.get("buckets") or [] - bucket_5h = None - bucket_weekly = None - for b in buckets: - if not isinstance(b, dict): + day = activity(entry.get("timestamp"), entry.get("conversationId") or entry.get("conversation_id")) + if day: + history_prompts[day] += 1 + + for transcript in (root / "brain").glob("*/.system_generated/logs/transcript.jsonl"): + session = transcript.parents[2].name + # An interrupted append can leave repeated steps. Count each dated user + # input once, without examining or estimating the text's token length. + seen = set() + for entry in json_lines(transcript): + if not isinstance(entry.get("created_at"), str) or not isinstance(entry.get("step_index"), int): continue - w = str(b.get("window") or "").lower() - bid = str(b.get("bucketId") or "").lower() - dname = str(b.get("displayName") or "").lower() - if w == "5h" or "5h" in bid or "five" in dname: - bucket_5h = b - elif w == "weekly" or "weekly" in bid or "week" in dname: - bucket_weekly = b - - # Enforce Session (5-hour) first, then Weekly (7-day) - if bucket_5h: - rem = float(bucket_5h.get("remainingFraction", 1.0)) - used = round(max(0.0, min(1.0, 1.0 - rem)), 4) - limits.append({ - "label": "Session (5-hour)", - "title": "Session", - "percent": used, - "resetsAt": normalize_reset_at(bucket_5h.get("resetTime")), - }) - if bucket_weekly: - rem = float(bucket_weekly.get("remainingFraction", 1.0)) - used = round(max(0.0, min(1.0, 1.0 - rem)), 4) - limits.append({ - "label": "Weekly (7-day)", - "title": "Weekly", - "percent": used, - "resetsAt": normalize_reset_at(bucket_weekly.get("resetTime")), - }) - - # Scoped limits for other groups (e.g. Claude and GPT models) - for g in other_groups: - g_name = str(g.get("displayName") or "").strip() - short_name = "Claude/GPT" if "claude" in g_name.lower() else (g_name or "3P") - for b in (g.get("buckets") or []): - if not isinstance(b, dict): + key = (entry.get("step_index"), entry.get("created_at")) + if key in seen: continue - rem = float(b.get("remainingFraction", 1.0)) - used = round(max(0.0, min(1.0, 1.0 - rem)), 4) - # Only show secondary model limits if there has been usage - if used > 0: - w = str(b.get("window") or "").lower() - w_title = "Weekly" if "week" in w else "Session" - title = f"{short_name} {w_title}" - limits.append({ - "label": title, - "title": title, - "percent": used, - "resetsAt": normalize_reset_at(b.get("resetTime")), - }) - - return limits - - -def probe_antigravity(access_token: str, endpoint_base: str = "") -> dict[str, Any]: - if not endpoint_base: - endpoint_base = os.environ.get("CLOUD_CODE_URL", "https://daily-cloudcode-pa.googleapis.com").rstrip("/") - - headers = { - "Authorization": f"Bearer {access_token}", - "Content-Type": "application/json", - "User-Agent": "antigravity", - } - - result: dict[str, Any] = { - "ok": False, - "tierLabel": "", - "limits": [], - "helpText": "", - "rateLimited": False, - "transport": False, - } - - tier = "" - try: - req = urllib.request.Request(f"{endpoint_base}/v1internal:loadCodeAssist", data=b"{}", headers=headers) - with urllib.request.urlopen(req, timeout=8) as resp: - payload = json.loads(resp.read().decode("utf-8", errors="replace")) - tier = extract_tier(payload) - except urllib.error.HTTPError as err: - if err.code == 429: - result["rateLimited"] = True - elif err.code == 401: - result["helpText"] = "Antigravity authentication expired. Run `agy` to authenticate." - return result - except Exception: - result["transport"] = True - result["helpText"] = "Could not reach Antigravity service." - - limits = [] - try: - req = urllib.request.Request(f"{endpoint_base}/v1internal:retrieveUserQuotaSummary", data=b"{}", headers=headers) - with urllib.request.urlopen(req, timeout=8) as resp: - payload = json.loads(resp.read().decode("utf-8", errors="replace")) - limits = extract_limits(payload) - except urllib.error.HTTPError as err: - if err.code == 429: - result["rateLimited"] = True - elif err.code == 401: - result["helpText"] = "Antigravity authentication expired. Run `agy` to authenticate." - return result - except Exception: - result["transport"] = True - - result["ok"] = bool(tier or limits) - result["tierLabel"] = tier - result["limits"] = limits - return result - - -def limit_window_open(entry: dict[str, Any], now: datetime) -> bool: - raw = str(entry.get("resetsAt") or "") - if raw == "": - return True - try: - resets_at = datetime.fromisoformat(raw.replace("Z", "+00:00")) - except Exception: - return True - if resets_at.tzinfo is None: - resets_at = resets_at.replace(tzinfo=timezone.utc) - return resets_at > now - - -def usable_cached_limits(cached: dict[str, Any]) -> list[dict[str, Any]]: - entries = cached.get("limits") - if not isinstance(entries, list): - return [] - now = datetime.now(timezone.utc) - return [entry for entry in entries if isinstance(entry, dict) and limit_window_open(entry, now)] - - -def detect_tier_label( - root_dir: Path, - auth_ok: bool, - email: str, - auth_method: str = "consumer", - probed_tier: str = "", - cached_tier: str = "", -) -> str: - if not auth_ok: - return "" - - if os.environ.get("AGY_TIER"): - return os.environ["AGY_TIER"] - - cfg_file = Path.home() / ".config" / "omarchy" / "agents" / "antigravity.json" - if cfg_file.exists(): + seen.add(key) + day = activity(entry.get("created_at"), session) + if day and entry.get("type") == "USER_INPUT": + transcript_prompts[day] += 1 + + # Summaries also know about conversations whose transcript is unavailable. + database = root / "conversation_summaries.db" + if database.exists(): try: - with cfg_file.open("r", encoding="utf-8") as f: - data = json.load(f) - if data.get("tier"): - return str(data["tier"]) - except Exception: + connection = sqlite3.connect(database.as_uri() + "?mode=ro", uri=True) + try: + for session, modified in connection.execute("SELECT conversation_id, last_modified_time FROM conversation_summaries"): + if session: + sessions.add(session) + activity(modified, session) + finally: + connection.close() + except sqlite3.Error: pass - if probed_tier: - return probed_tier - - if cached_tier: - return cached_tier - - if auth_method == "enterprise": - return "Enterprise" - if auth_method in ("cloud", "vertex"): - return "Cloud" - - return "Pro" - - -def collect_local_stats(root_dir: Path, now: datetime) -> dict[str, Any]: - today_str = now.strftime("%Y-%m-%d") - recent_dates = [(now - timedelta(days=i)).strftime("%Y-%m-%d") for i in range(6, -1, -1)] - recent = {d: {"date": d, "messageCount": 0} for d in recent_dates} - - history_file = root_dir / "history.jsonl" - brain_dir = root_dir / "brain" + # The same prompt lives in both logs. Per-day maxima use history as a + # fallback without adding it on top of transcript counts. + prompts = {day: max(history_prompts[day], transcript_prompts[day]) for day in active_dates} + return { + "todayPrompts": prompts.get(today, 0), + "todaySessions": len(today_sessions), + "todayTotalTokens": 0, + "todayTokensByModel": {}, + "recentDays": [{"date": day, "messageCount": 0} for day in days], + "totalPrompts": sum(prompts.values()), + "totalSessions": len(sessions), + "activeDays": len(active_dates), + "activeDates": sorted(active_dates), + "modelUsage": {}, + } - today_prompts = 0 - total_prompts = 0 - active_dates: set[str] = set() - today_sessions: set[str] = set() - total_sessions: set[str] = set() - if history_file.exists(): +def installed_cli(): + path = shutil.which("agy") + if path: + executable = Path(path) + # Omarchy's stock wrapper installs agy on first invocation. A passive + # usage scan must not install a CLI, including through a mise shim. try: - with history_file.open("r", encoding="utf-8", errors="ignore") as f: - for line in f: - line = line.strip() - if not line: - continue - try: - entry = json.loads(line) - ts = entry.get("timestamp") - if not ts: - continue - try: - if isinstance(ts, (int, float)): - dt = datetime.fromtimestamp(ts / 1000.0, timezone.utc) - else: - dt = datetime.fromisoformat(str(ts).replace("Z", "+00:00")) - except Exception: - continue - d_str = dt.strftime("%Y-%m-%d") - active_dates.add(d_str) - total_prompts += 1 - conv_id = entry.get("conversationId") or entry.get("conversation_id") - if conv_id: - total_sessions.add(conv_id) - if d_str == today_str: - today_prompts += 1 - if conv_id: - today_sessions.add(conv_id) - except Exception: - continue - except Exception: + wrapper = b"mise use" in executable.read_bytes()[:1024] if executable.stat().st_size < 8192 else False + if executable.resolve().name != "mise" and not wrapper: + return str(executable) + except OSError: pass - - # Fallback to conversation summaries if history.jsonl is sparse - summaries_db = root_dir / "conversation_summaries.db" - if summaries_db.exists(): + mise = shutil.which("mise") + if mise: try: - conn = sqlite3.connect(f"file:{summaries_db}?mode=ro", uri=True) - cur = conn.cursor() - for row in cur.execute("SELECT conversation_id, last_modified_time FROM conversation_summaries"): - cid, lmt = row[0], str(row[1] or "") - if cid: - total_sessions.add(cid) - if lmt.startswith(today_str): - today_sessions.add(cid) - conn.close() - except Exception: + found = subprocess.run([mise, "which", "agy"], stdin=subprocess.DEVNULL, capture_output=True, text=True, timeout=3) + executable = Path(found.stdout.strip()) + if found.returncode == 0 and executable.is_file() and os.access(executable, os.X_OK): + return str(executable) + except (OSError, subprocess.TimeoutExpired): pass + return "" - # Compute token counts across conversation transcripts - model_usage: dict[str, dict[str, int]] = {} - today_tokens = 0 - today_tokens_by_model: dict[str, int] = {} - if brain_dir.exists(): +def run_cli(cli, *args, timeout=20): + # A timed-out CLI must not leave descendants holding our output pipes. + with subprocess.Popen([cli, *args], stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, + stderr=subprocess.PIPE, text=True, start_new_session=True) as process: try: - for conv_folder in brain_dir.iterdir(): - if not conv_folder.is_dir(): - continue - tr_file = conv_folder / ".system_generated/logs/transcript.jsonl" - if not tr_file.exists(): - continue - try: - with tr_file.open("r", encoding="utf-8", errors="ignore") as f: - for line in f: - line = line.strip() - if not line: - continue - try: - step = json.loads(line) - m_name = step.get("model") or "gemini-3.8-flash-high" - if m_name not in model_usage: - model_usage[m_name] = { - "inputTokens": 0, - "outputTokens": 0, - "cacheReadInputTokens": 0, - "cacheCreationInputTokens": 0, - } - - step_usage = step.get("usage") - is_model = step.get("source") == "MODEL" or step.get("type") == "PLANNER_RESPONSE" - if isinstance(step_usage, dict): - inp = int(step_usage.get("prompt_token_count") or step_usage.get("input_tokens") or 0) - out = int(step_usage.get("candidates_token_count") or step_usage.get("output_tokens") or 0) - tok = inp + out - model_usage[m_name]["inputTokens"] += inp - model_usage[m_name]["outputTokens"] += out - else: - c_text = step.get("content") or "" - t_text = step.get("thinking") or "" - tok = max(1, (len(c_text) + len(t_text)) // 4) - if is_model: - model_usage[m_name]["outputTokens"] += tok - else: - model_usage[m_name]["inputTokens"] += tok - - created = step.get("created_at") - is_today = False - if created: - step_d = created[:10] - if step_d in recent: - recent[step_d]["messageCount"] += tok - if step_d == today_str: - is_today = True - else: - is_today = True - - if is_today: - today_tokens += tok - today_tokens_by_model[m_name] = today_tokens_by_model.get(m_name, 0) + tok - except Exception: - continue - except Exception: - continue - except Exception: - pass - - return { - "todayPrompts": today_prompts, - "todaySessions": len(today_sessions), - "todayTotalTokens": today_tokens, - "todayTokensByModel": today_tokens_by_model, - "recentDays": [recent[d] for d in recent_dates], - "totalPrompts": total_prompts, - "totalSessions": len(total_sessions), - "activeDays": len(active_dates), - "activeDates": sorted(list(active_dates)), - "modelUsage": model_usage, - } + stdout, stderr = process.communicate(timeout=timeout) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGKILL) + process.communicate() + raise + return process.returncode, stdout, stderr -def calculate_limits(root_dir: Path, now: datetime) -> list[dict[str, Any]]: - """Fallback calculation of rolling limits from local history files.""" - history_file = root_dir / "history.jsonl" - five_hours_ago_ts = int((now - timedelta(hours=5)).timestamp() * 1000) +def supported_version(value): + match = re.fullmatch(r"(?:agy\s+|Antigravity CLI\s+)?v?(\d+)\.(\d+)\.(\d+)", value.strip(), re.I) + return bool(match and tuple(map(int, match.groups())) >= (1, 1, 11)) - prompts_5h = 0 - oldest_5h_ts = None - prompts_7d = 0 - seven_days_ago_ts = int((now - timedelta(days=7)).timestamp() * 1000) +def extract_limits(data): + groups = data.get("groups") + if not isinstance(groups, list): + return [] + groups = [group for group in groups if isinstance(group, dict)] + primary = next((group for group in groups if "gemini" in str(group.get("name", "")).lower()), None) + ordered = ([primary] if primary is not None else []) + [group for group in groups if group is not primary] + limits = [] + for group in ordered: + name = str(group.get("name") or "Models") + scoped = "Claude/GPT" if "claude" in name.lower() else name + buckets = group.get("buckets") + if not isinstance(buckets, list): + continue + for window, title, label in (("5h", "Session", "Session (5-hour)"), ("weekly", "Weekly", "Weekly (7-day)")): + for bucket in buckets: + if not isinstance(bucket, dict) or bucket.get("window") != window: + continue + remaining = bucket.get("remaining_fraction") + if isinstance(remaining, bool) or not isinstance(remaining, (int, float)): + continue + if not math.isfinite(remaining) or not 0 <= remaining <= 1: + continue + used = round(1 - remaining, 4) + if group is not primary and used == 0: + continue + reset = parse_time(bucket.get("reset_time")) + limits.append({ + "label": label if group is primary else f"{scoped} {label}", + "title": title if group is primary else f"{scoped} {title}", + "percent": used, + "resetsAt": reset.isoformat() if reset else "", + }) + return limits - if history_file.exists(): - try: - with history_file.open("r", encoding="utf-8", errors="ignore") as f: - for line in f: - line = line.strip() - if not line: - continue - try: - entry = json.loads(line) - ts = entry.get("timestamp") - if not ts: - continue - if ts >= seven_days_ago_ts: - prompts_7d += 1 - if ts >= five_hours_ago_ts: - prompts_5h += 1 - if oldest_5h_ts is None or ts < oldest_5h_ts: - oldest_5h_ts = ts - except Exception: - continue - except Exception: - pass - # Rolling 5h reset - if oldest_5h_ts: - reset_5h = datetime.fromtimestamp(oldest_5h_ts / 1000.0, timezone.utc) + timedelta(hours=5) +def probe_failure(message): + # Do not publish raw stderr: it can contain account or credential details. + message = message.lower() + status = "Antigravity limits unavailable" + help_text = "Antigravity could not read quotas. Run `agy` and check `/usage`, then refresh." + retry = False + if re.search(r"\b429\b|resource.?exhausted|rate.?limit", message): + status = "Antigravity limits rate limited" + help_text = "Antigravity is rate limiting quota checks. Try refreshing later." + elif re.search(r"\b401\b|\b403\b|unauthenticated|permission.?denied", message): + status = "Antigravity sign-in required" + help_text = AUTH_HELP + elif re.search(r"network|connection|dial tcp|no such host|timed? ?out|timeout|deadline exceeded|unreachable", message): + help_text = "Could not reach Antigravity. Retrying shortly; local history is still shown." + retry = True + elif re.search(r"sign.?in|log.?in|auth|credential", message): + status = "Antigravity sign-in required" + help_text = AUTH_HELP + return {"usageStatusText": status, "authHelpText": help_text, "retryAdvised": retry} + + +def probe_limits(cli): + try: + code, version, _ = run_cli(cli, "--version", timeout=5) + if code != 0 or not supported_version(version): + return {"usageStatusText": "Antigravity update required", "authHelpText": "Quota checks require Antigravity CLI 1.1.11 or newer. Update `agy`, then refresh."} + code, stdout, stderr = run_cli(cli, "-p", "/usage", "--output-format", "json") + if code != 0: + return probe_failure(stderr + "\n" + stdout) + payload = json.loads(stdout) + command = payload.get("command") if isinstance(payload, dict) else None + if not isinstance(command, dict) or command.get("name") not in ("usage", "quota") or payload.get("status") != "SUCCESS": + return probe_failure("invalid quota response") + data = command.get("data") + limits = extract_limits(data) if isinstance(data, dict) else [] + if not limits: + return {"usageStatusText": "Antigravity limits unavailable", "authHelpText": "Antigravity returned no quota windows for this sign-in. Check `/usage` in `agy`."} + return {"limits": limits, "usageStatusText": "", "authHelpText": ""} + except subprocess.TimeoutExpired: + return probe_failure("timeout") + except (OSError, ValueError): + return probe_failure("could not read CLI output") + + +def collect_limits(cli, path, force): + now = time.time() + cached = read_json(path) + entries = cached.get("limits") + fetched = cached.get("fetchedAtMs", 0) + if not isinstance(fetched, (int, float)) or not math.isfinite(fetched): + fetched = 0 + # Unknown reset times cannot establish that an old window is still open. + fallback = [entry for entry in entries if isinstance(entry, dict) + and isinstance(entry.get("percent"), (int, float)) and 0 <= entry["percent"] <= 1 + and (reset := parse_time(entry.get("resetsAt"))) and reset.timestamp() > now] if isinstance(entries, list) else [] + result = {"limits": fallback, "limitsStale": bool(fallback), "limitsFetchedAt": fetched, + "usageStatusText": "", "authHelpText": "", "retryAdvised": False} + checked = cached.get("checkedAtMs", 0) + reuse = 15 if cached.get("retryAdvised") else PROBE_REUSE_SECONDS + if not cli: + result.update(usageStatusText="Antigravity CLI unavailable", authHelpText="Install Antigravity CLI and run `agy` to sign in.") + elif not force and isinstance(checked, (int, float)) and 0 <= now - checked / 1000 < reuse: + result.update({key: cached[key] for key in ("usageStatusText", "authHelpText", "retryAdvised") if key in cached}) + result["limitsStale"] = bool(fallback) and cached.get("limitsStale", False) + if not fallback and not result["usageStatusText"]: + # A window rolled over during the reuse interval; get its new value. + return collect_limits(cli, path, True) + return result else: - reset_5h = now + timedelta(hours=5) - - # Check if a 429 / resource exhausted error happened recently - cli_log = root_dir / "cli.log" - rate_limited = False - if cli_log.exists(): - try: - stat = cli_log.stat() - if (time.time() - stat.st_mtime) < 18000: - with cli_log.open("r", encoding="utf-8", errors="ignore") as f: - content = f.read()[-8192:] - if "RESOURCE_EXHAUSTED" in content or "rate limit exceeded" in content.lower(): - rate_limited = True - except Exception: - pass - - session_pct = 1.0 if rate_limited else round(min(1.0, prompts_5h / 50.0), 2) - weekly_pct = round(min(1.0, prompts_7d / 250.0), 2) - - return [ - { - "label": "Session (5-hour)", - "title": "Session", - "percent": session_pct, - "resetsAt": reset_5h.strftime("%Y-%m-%dT%H:%M:%S+00:00"), - }, - { - "label": "Weekly (7-day)", - "title": "Weekly", - "percent": weekly_pct, - "resetsAt": (now + timedelta(days=7)).strftime("%Y-%m-%dT%H:%M:%S+00:00"), - }, - ] - - -def collect_limits( - root_dir: Path, - now: datetime, - auth_ok: bool, - access_token: str, - expires_at_ms: int, - auth_method: str, - force: bool, -) -> tuple[list[dict[str, Any]], str, str, str]: - """Collect limits via API probe, falling back to cached or rolling local stats.""" - limits_cache = cache_root() / f"{AGENT_ID}-limits.json" - cached = read_fresh_json(limits_cache, float("inf")) or {} - fallback = usable_cached_limits(cached) - cached_tier = str(cached.get("tierLabel") or "") - - if not auth_ok: - return [], "", "Sign in to see limits", AUTH_HELP - - # If access token is empty, we cannot probe the remote endpoint - if not access_token: - limits = fallback if fallback else calculate_limits(root_dir, now) - return limits, cached_tier, "", AUTH_HELP - - # Check reuse window for limits probe (PROBE_MIN_INTERVAL_SECONDS) - fetched_at = float(cached.get("fetchedAtMs") or 0) / 1000.0 - if fallback and not force and (time.time() - fetched_at) < PROBE_MIN_INTERVAL_SECONDS: - return fallback, cached_tier, "", AUTH_HELP - - # Perform live probe - probe = probe_antigravity(access_token) - if probe["ok"]: - probed_limits = probe["limits"] - probed_tier = probe["tierLabel"] - write_json( - limits_cache, - { - "fetchedAtMs": round(time.time() * 1000), - "limits": probed_limits, - "tierLabel": probed_tier, - }, - ) - return probed_limits, probed_tier, "", AUTH_HELP - - # Probe failed - help_text = probe["helpText"] if probe.get("helpText") else AUTH_HELP - usage_status = "" - if probe.get("rateLimited"): - usage_status = "Antigravity limits rate limited" - elif probe.get("transport"): - usage_status = "" - - if fallback: - return fallback, cached_tier, usage_status, help_text - - # Fallback to local rolling calculation - local_limits = calculate_limits(root_dir, now) - return local_limits, cached_tier, usage_status, help_text - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--force", action="store_true", help="Bypass local caches and force a fresh scan") - parser.add_argument("--limits-only", action="store_true", help="Recompute only rate limits if cache is fresh") - args = parser.parse_args() - - root_dir = app_root() - now = datetime.now(timezone.utc) - today_str = now.strftime("%Y-%m-%d") - - auth_ok, email, auth_method, access_token, expires_at_ms = check_auth() + probe = probe_limits(cli) + result.update(probe) + if probe.get("limits"): + result.update(limitsStale=False, limitsFetchedAt=round(time.time() * 1000)) + elif result["usageStatusText"] == "Antigravity sign-in required": + # The CLI may have signed out or switched accounts. Do not attribute a + # previous account's numbers to the next successful sign-in. + result.update(limits=[], limitsStale=False, limitsFetchedAt=0) + if result["limitsStale"]: + result["authHelpText"] += " Showing the last known limits until their windows reset." + write_json(path, {**result, "fetchedAtMs": result["limitsFetchedAt"], "checkedAtMs": round(time.time() * 1000)}) + return result - cache_file, lock_file = scan_cache_paths() - stats = None - if not args.force: - max_age = LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS - cached = read_fresh_json(cache_file, max_age) - if isinstance(cached, dict) and cached.get("scanDate") == today_str: - stats = cached.get("stats") +def tier_label(): + # /usage does not report a plan. Only show a label the user supplied. + config_home = Path(os.environ.get("XDG_CONFIG_HOME") or Path.home() / ".config") + value = os.environ.get("AGY_TIER") or read_json(config_home / "omarchy/agents/antigravity.json").get("tier") + return value.strip() if isinstance(value, str) else "" - if stats is None: - try: - with lock_file.open("w") as lock: - fcntl.flock(lock, fcntl.LOCK_EX) - stats = collect_local_stats(root_dir, now) - write_json(cache_file, {"schemaVersion": 1, "scanDate": today_str, "stats": stats}) - except Exception as exc: - print(f"omarchy-agent-usage-antigravity: cache lock failed ({exc}), scanning directly", file=sys.stderr) - stats = collect_local_stats(root_dir, now) - - limits, probed_tier, usage_status_text, auth_help_text = collect_limits( - root_dir, - now, - auth_ok, - access_token, - expires_at_ms, - auth_method, - args.force, - ) - - cached_limits_file = cache_root() / f"{AGENT_ID}-limits.json" - cached_payload = read_fresh_json(cached_limits_file, float("inf")) or {} - cached_tier = str(cached_payload.get("tierLabel") or "") - tier_label = detect_tier_label(root_dir, auth_ok, email, auth_method, probed_tier, cached_tier) - - total_prompts = stats.get("totalPrompts", 0) - ready = auth_ok or total_prompts > 0 or len(limits) > 0 - has_local_stats = total_prompts > 0 - - record = { - "schemaVersion": 1, - "id": AGENT_ID, - "name": AGENT_NAME, - "updatedAt": now.strftime("%Y-%m-%dT%H:%M:%S.%f+00:00"), - "ready": ready, - "hasLocalStats": has_local_stats, - "todayPrompts": stats.get("todayPrompts", 0), - "todaySessions": stats.get("todaySessions", 0), - "todayTotalTokens": stats.get("todayTotalTokens", 0), - "todayTokensByModel": stats.get("todayTokensByModel", {}), - "recentDays": stats.get("recentDays", []), - "totalPrompts": total_prompts, - "totalSessions": stats.get("totalSessions", 0), - "activeDays": stats.get("activeDays", 0), - "activeDates": stats.get("activeDates", []), - "modelUsage": stats.get("modelUsage", {}), - "limits": limits, - "tierLabel": tier_label, - "usageStatusText": usage_status_text, - "authHelpText": auth_help_text, - } - print(json.dumps(record, separators=(",", ":"))) +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--force", action="store_true", help="Bypass local caches and read fresh quotas") + parser.add_argument("--limits-only", action="store_true", help="Reuse a recent history scan") + args = parser.parse_args() + root = app_root() + now = datetime.now().astimezone() + # Isolate custom app roots and invalidate the old estimated-token cache. + key = hashlib.sha256(str(root).encode()).hexdigest()[:16] + prefix = cache_root() / f"antigravity-v{CACHE_VERSION}-{key}" + prefix.parent.mkdir(parents=True, exist_ok=True) + with prefix.with_suffix(".lock").open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + stats_path = prefix.with_suffix(".stats.json") + cached = read_json(stats_path) + age = LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS + fetched = cached.get("fetchedAt", 0) + stats = cached.get("stats") + if (args.force or not isinstance(stats, dict) + or not all(isinstance(stats.get(key), int) for key in ("totalPrompts", "totalSessions", "todayPrompts", "todaySessions")) + or cached.get("scanDate") != now.date().isoformat() + or not isinstance(fetched, (int, float)) or not 0 <= time.time() - fetched < age): + stats = collect_local_stats(root, now) + write_json(stats_path, {"scanDate": now.date().isoformat(), "fetchedAt": time.time(), "stats": stats}) + limits = collect_limits(installed_cli(), prefix.with_suffix(".limits.json"), args.force) + has_stats = stats["totalPrompts"] > 0 or stats["totalSessions"] > 0 + record = {"schemaVersion": 1, "id": "antigravity", "name": "Antigravity", + "updatedAt": datetime.now(timezone.utc).isoformat(), "ready": has_stats or bool(limits["limits"]), + "hasLocalStats": has_stats, "tierLabel": tier_label(), **stats, **limits} + print(json.dumps(record, separators=(",", ":"), allow_nan=False)) if __name__ == "__main__": diff --git a/manual/17-ai.md b/manual/17-ai.md index ea8a8d65eaf..30a5fc53982 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -38,7 +38,7 @@ There are terminal shortcuts too: `a` runs the default agent inline in the curre ### The agents panel -The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan and the percentage used of its 5-hour session and weekly limits (or the remaining prepaid balance), under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Claude Code, Codex, Grok, and Fireworks are covered out of the box. +The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan and the percentage used of its 5-hour session and weekly limits (or the remaining prepaid balance), under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Antigravity, Claude Code, Codex, Grok, and Fireworks are covered out of the box. Antigravity needs CLI 1.1.11 or newer and a sign-in through `agy`; its section shows subscription limits and contributes local prompt/session counts, but its saved transcripts do not provide token counts. Left-click the bar icon for the panel, right-click to launch your default agent. The usage records behind it are regenerated every 15 minutes by `omarchy agent usage-update`, and the panel can even merge usage from your other machines via a synced folder. See the README under `$OMARCHY_PATH/shell/plugins/agents/` for the full settings. diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 10aa6472a9d..7f792ae4a06 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -64,13 +64,13 @@ record that lands in the directory regardless of who wrote it. Adding an agent therefore never touches this plugin: ship a collector that prints the record contract (see the `claude` and `codex` collectors in -`bin/`), and the panel gains a tab. An `assets/.svg` mark is optional — +`bin/`), and the panel gains a section. An `assets/.svg` mark is optional — with an `assets/-light.svg` twin if the mark needs a dark variant for light surfaces — and the bar glyph stands in when there is none. | Collector | Limits | Local stats | |---|---|---| -| `antigravity` | Google Cloud Code API | Antigravity CLI transcripts and history | +| `antigravity` | The Antigravity CLI’s read-only `/usage` command (1.1.11 or newer) | Prompt/session counts from transcripts, history and conversation summaries under `~/.gemini/antigravity-cli/` (`AGY_DIR`); token counts are unavailable | | `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) | | `grok` | The credits endpoint behind Grok's `/usage` view (the billing period's included usage) | Each session's `usage.json` (the ledger `grok usage` prints: tokens by model per finished turn), plus `summary.json` for sessions | @@ -102,6 +102,14 @@ again. Fireworks reads opencode stores in `~/.local/share/opencode/auth.json` when Fireworks is signed in there. +### Antigravity quotas and history + +The collector runs `agy -p "/usage" --output-format json` after checking that the installed CLI is at least 1.1.11. That version introduced read-only slash-command output without an agent turn or quota spend; an older or unrecognized version is never sent `/usage`. It resolves an existing mise installation without invoking Omarchy’s installer wrapper. Antigravity itself handles credential storage, token refresh and the account’s quota request. + +Session and weekly meters use the returned quota windows, including used Claude/GPT allowances. Missing percentages are unknown, not zero usage. Failed checks show a status and help text; network failures request an early retry. Last known windows remain visible and dimmed only until they reset, and an authentication failure clears them. Refresh bypasses the short reuse interval. `AGY_DIR` selects another app data directory with its own caches. + +The CLI’s transcripts do not contain token usage or model IDs. Local history contributes prompt and session counts on local calendar dates, without character-based token estimates or invented model rows. History and transcript prompt counts are combined by taking the larger count for each day, since the history may not carry conversation IDs. The `/usage` response does not include a subscription tier, so the plan stays unknown unless you supply `AGY_TIER` or `{ "tier": "Ultra" }` in `~/.config/omarchy/agents/antigravity.json` (honoring `XDG_CONFIG_HOME`). Sign in through `agy`; the panel’s built-in account switching remains for Claude, Codex and Grok. + ### Fireworks balance The collector first asks the account's `:getBalance` endpoint for the real @@ -178,6 +186,7 @@ edit `shell.json` directly): ```bash omarchy bar set omarchy.agents providers '{ + "antigravity": { "enabled": true }, "claude": { "enabled": true }, "codex": { "enabled": false }, "fireworks": { "enabled": true } diff --git a/test/shell.d/agent-usage-antigravity-limits-test.sh b/test/shell.d/agent-usage-antigravity-limits-test.sh index 89c0dbdc945..54852c3eac1 100755 --- a/test/shell.d/agent-usage-antigravity-limits-test.sh +++ b/test/shell.d/agent-usage-antigravity-limits-test.sh @@ -1,136 +1,117 @@ #!/bin/bash -source "$(dirname "$0")/base-test.sh" - -require_command jq +set -euo pipefail +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" require_command python3 -COLLECTOR="$ROOT/bin/omarchy-agent-usage-antigravity" - -# probe_antigravity reaches Google's Cloud Code endpoints, so the readers that -# interpret the answers are exercised with recorded payloads standing in. - -test_tier() { - COLLECTOR="$COLLECTOR" PAYLOAD="$1" python3 - <<'PY' -import importlib.machinery, importlib.util, 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) - -payload = json.loads(os.environ["PAYLOAD"]) -print(collector.extract_tier(payload)) -PY -} - -ultra_tier=$(test_tier '{ - "currentTier": { "id": "free-tier", "name": "Antigravity" }, - "paidTier": { "id": "g1-ultra-tier", "name": "Google AI Ultra" } -}') -[[ $ultra_tier == "Ultra" ]] || - fail "Antigravity collector extracts Ultra tier from paidTier" "$ultra_tier" -pass "Antigravity collector extracts Ultra tier from paidTier" - -pro_tier=$(test_tier '{ - "currentTier": { "id": "free-tier", "name": "Antigravity" }, - "paidTier": { "id": "g1-pro-tier", "name": "Google AI Pro" } -}') -[[ $pro_tier == "Pro" ]] || - fail "Antigravity collector extracts Pro tier from paidTier" "$pro_tier" -pass "Antigravity collector extracts Pro tier from paidTier" - -free_tier=$(test_tier '{ - "currentTier": { "id": "free-tier", "name": "Antigravity" } -}') -[[ $free_tier == "Free" ]] || - fail "Antigravity collector extracts Free tier when no paidTier" "$free_tier" -pass "Antigravity collector extracts Free tier when no paidTier" - -# Test quota extraction -test_limits() { - COLLECTOR="$COLLECTOR" PAYLOAD="$1" python3 - <<'PY' -import importlib.machinery, importlib.util, json, os - -loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) +python3 - <<'PY' +import importlib.machinery +import importlib.util +import json +import os +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +root = Path(os.environ['ROOT']) +loader = importlib.machinery.SourceFileLoader('collector', str(root / 'bin/omarchy-agent-usage-antigravity')) spec = importlib.util.spec_from_loader(loader.name, loader) collector = importlib.util.module_from_spec(spec) loader.exec_module(collector) - -payload = json.loads(os.environ["PAYLOAD"]) -print(json.dumps(collector.extract_limits(payload))) +fixture = json.loads((root / 'test/shell.d/fixtures/antigravity/usage.json').read_text()) + +class Limits(unittest.TestCase): + def test_real_cli_shape_and_reversed_groups(self): + data = fixture['command']['data'] + limits = collector.extract_limits(data) + self.assertEqual([v['title'] for v in limits], ['Session', 'Weekly', 'Claude/GPT Weekly']) + self.assertEqual([v['percent'] for v in limits], [0, .27, .03]) + reversed_data = {'groups': list(reversed(data['groups']))} + self.assertEqual(limits, collector.extract_limits(reversed_data)) + self.assertEqual(limits[0]['resetsAt'], '2099-10-02T07:42:57+00:00') + + def test_missing_invalid_and_exhausted_buckets(self): + for remaining in [None, 'bad', True, float('nan'), float('inf'), -1, 2]: + data = {'groups': [{'name': 'Gemini Models', 'buckets': [{'window': '5h', 'remaining_fraction': remaining}]}]} + self.assertEqual(collector.extract_limits(data), [], remaining) + data['groups'][0]['buckets'][0]['remaining_fraction'] = 0 + self.assertEqual(collector.extract_limits(data)[0]['percent'], 1) + del data['groups'][0]['buckets'][0]['remaining_fraction'] + self.assertEqual(collector.extract_limits(data), []) + for data in [{}, {'groups': None}, {'groups': [None, {'buckets': 7}]}]: + self.assertEqual(collector.extract_limits(data), []) + + def test_only_scoped_group_keeps_its_identity(self): + data = {'groups': [fixture['command']['data']['groups'][0]]} + self.assertEqual(collector.extract_limits(data)[0]['title'], 'Claude/GPT Weekly') + + def test_version_gate_never_prompts_an_old_or_unknown_cli(self): + for version in ['1.1.10', '1.0.99', 'unknown', '1.1.11-preview', '']: + with patch.object(collector, 'run_cli', return_value=(0, version, '')) as run: + result = collector.probe_limits('agy') + self.assertEqual(run.call_count, 1) + self.assertEqual(result['usageStatusText'], 'Antigravity update required') + for version in ['1.1.11', '1.2.14', '2.0.0', 'agy 1.2.14']: + self.assertTrue(collector.supported_version(version)) + + def test_failures_and_malformed_success(self): + cases = [('401 unauthenticated', 'Antigravity sign-in required', False), + ('403 PermissionDenied', 'Antigravity sign-in required', False), + ('429 RESOURCE_EXHAUSTED', 'Antigravity limits rate limited', False), + ('dial tcp: no such host', 'Antigravity limits unavailable', True), + ('refresh credentials: dial tcp: connection refused', 'Antigravity limits unavailable', True), + ('500 internal server error', 'Antigravity limits unavailable', False)] + for error, status, retry in cases: + with patch.object(collector, 'run_cli', side_effect=[(0, '1.2.14', ''), (1, '', error)]): + result = collector.probe_limits('agy') + self.assertEqual(result['usageStatusText'], status) + self.assertEqual(result['retryAdvised'], retry) + self.assertNotIn(error, result['authHelpText']) + for output in ['invalid', '[]', '{}', '{"status":"SUCCESS","response":"pretend usage"}', + '{"status":"SUCCESS","command":{"name":"usage","data":{}}}']: + with patch.object(collector, 'run_cli', side_effect=[(0, '1.2.14', ''), (0, output, '')]): + self.assertTrue(collector.probe_limits('agy')['usageStatusText']) + with patch.object(collector, 'run_cli', side_effect=collector.subprocess.TimeoutExpired('agy', 20)): + self.assertTrue(collector.probe_limits('agy')['retryAdvised']) + + def test_stale_cache_expiry_auth_clear_and_force(self): + good = {'limits': collector.extract_limits(fixture['command']['data']), 'usageStatusText': '', 'authHelpText': ''} + with tempfile.TemporaryDirectory() as tmp: + cache = Path(tmp) / 'limits.json' + with patch.object(collector, 'probe_limits', return_value=good) as probe: + live = collector.collect_limits('agy', cache, False) + self.assertFalse(live['limitsStale']) + self.assertGreater(live['limitsFetchedAt'], 0) + self.assertEqual(collector.collect_limits('agy', cache, False)['limits'], live['limits']) + self.assertEqual(probe.call_count, 1) + collector.collect_limits('agy', cache, True) + self.assertEqual(probe.call_count, 2) + failure = collector.probe_failure('network unavailable') + with patch.object(collector, 'probe_limits', return_value=failure) as probe: + stale = collector.collect_limits('agy', cache, True) + self.assertEqual(stale['limits'], live['limits']) + self.assertTrue(stale['limitsStale']) + self.assertTrue(stale['retryAdvised']) + self.assertIn('last known limits', stale['authHelpText']) + reused = collector.collect_limits('agy', cache, False) + self.assertTrue(reused['limitsStale']) + self.assertEqual(probe.call_count, 1) + cached = json.loads(cache.read_text()) + cached['limits'][0]['resetsAt'] = '2000-01-01T00:00:00Z' + cached['limits'][1]['resetsAt'] = '' + collector.write_json(cache, cached) + remaining = collector.collect_limits('agy', cache, True) + self.assertEqual(len(remaining['limits']), 1) + with patch.object(collector, 'probe_limits', return_value=collector.probe_failure('401')): + signed_out = collector.collect_limits('agy', cache, True) + self.assertEqual(signed_out['limits'], []) + self.assertEqual(json.loads(cache.read_text())['limits'], []) + for invalid in ['[]', '{bad', '{"limits": 5, "fetchedAtMs": "bad"}']: + cache.write_text(invalid) + with patch.object(collector, 'probe_limits', return_value=failure): + self.assertEqual(collector.collect_limits('agy', cache, True)['limits'], []) + +unittest.main(verbosity=2) PY -} - -quota_payload='{ - "groups": [ - { - "displayName": "Gemini Models", - "buckets": [ - { - "bucketId": "gemini-weekly", - "displayName": "Weekly Limit Remaining", - "window": "weekly", - "resetTime": "2026-09-17T18:19:42Z", - "remainingFraction": 0.95 - }, - { - "bucketId": "gemini-5h", - "displayName": "Five Hour Limit Remaining", - "window": "5h", - "resetTime": "2026-09-15T19:31:46Z", - "remainingFraction": 0.90 - } - ] - }, - { - "displayName": "Claude and GPT models", - "buckets": [ - { - "bucketId": "3p-weekly", - "displayName": "Weekly Limit Remaining", - "window": "weekly", - "resetTime": "2026-09-18T12:56:14Z", - "remainingFraction": 0.98 - }, - { - "bucketId": "3p-5h", - "displayName": "Five Hour Limit Remaining", - "window": "5h", - "resetTime": "2026-09-15T22:16:22Z", - "remainingFraction": 1.0 - } - ] - } - ] -}' - -limits=$(test_limits "$quota_payload") - -[[ $(jq -r '.[0].label' <<<"$limits") == "Session (5-hour)" ]] || - fail "Antigravity collector places Session window first" "$limits" -pass "Antigravity collector places Session window first" - -[[ $(jq -r '.[0].percent' <<<"$limits") == "0.1" ]] || - fail "Antigravity collector calculates used percentage from remainingFraction" "$limits" -pass "Antigravity collector calculates used percentage from remainingFraction" - -[[ $(jq -r '.[1].label' <<<"$limits") == "Weekly (7-day)" ]] || - fail "Antigravity collector places Weekly window second" "$limits" -pass "Antigravity collector places Weekly window second" - -[[ $(jq -r '.[1].percent' <<<"$limits") == "0.05" ]] || - fail "Antigravity collector calculates weekly used percentage" "$limits" -pass "Antigravity collector calculates weekly used percentage" - -[[ $(jq -r '.[2].title' <<<"$limits") == "Claude/GPT Weekly" ]] || - fail "Antigravity collector includes used scoped 3P limits" "$limits" -pass "Antigravity collector includes used scoped 3P limits" - -[[ $(jq -r '.[2].percent' <<<"$limits") == "0.02" ]] || - fail "Antigravity collector calculates 3P used percentage" "$limits" -pass "Antigravity collector calculates 3P used percentage" - -[[ $(jq -r 'length' <<<"$limits") == "3" ]] || - fail "Antigravity collector omits unused 3P session limit" "$limits" -pass "Antigravity collector omits unused 3P session limit" +pass "Antigravity quota parsing, version safety, failures and cache freshness" diff --git a/test/shell.d/agent-usage-antigravity-scanner-test.sh b/test/shell.d/agent-usage-antigravity-scanner-test.sh index 48f6fed5119..17fffc921af 100755 --- a/test/shell.d/agent-usage-antigravity-scanner-test.sh +++ b/test/shell.d/agent-usage-antigravity-scanner-test.sh @@ -1,115 +1,180 @@ #!/bin/bash -source "$(dirname "$0")/base-test.sh" - -require_command jq +set -euo pipefail +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" require_command python3 -TEST_HOME=$(mktemp -d) -trap 'rm -rf "$TEST_HOME"' EXIT - -AGY_DIR="$TEST_HOME/.gemini/antigravity-cli" -mkdir -p "$AGY_DIR/brain/test-conv/.system_generated/logs" "$TEST_HOME/bin" - -# Mock secret-tool to simulate authenticated state -cat >"$TEST_HOME/bin/secret-tool" <<'EOF' -#!/bin/bash -if [[ "$*" == *"lookup service gemini username antigravity"* ]]; then - echo '{"auth_method":"consumer","id_token":"header.eyJlbWFpbCI6InRlc3RAZ21haWwuY29tIn0.signature"}' - exit 0 -fi -exit 1 -EOF -chmod +x "$TEST_HOME/bin/secret-tool" - -# Write mock history and transcript -now_date=$(date +%Y-%m-%d) -timestamp="$(date -u +%Y-%m-%dT%H:%M:%SZ)" - -cat >"$AGY_DIR/history.jsonl" <"$AGY_DIR/brain/test-conv/.system_generated/logs/transcript.jsonl" <"$TEST_HOME/.config/omarchy/agents/antigravity.json" <<'EOF' -{ - "tier": "Enterprise" -} -EOF - -cfg_tier_result=$(HOME="$TEST_HOME" AGY_DIR="$AGY_DIR" XDG_CACHE_HOME="$TEST_HOME/.cache" PATH="$TEST_HOME/bin:$PATH" \ - "$ROOT/bin/omarchy-agent-usage-antigravity" --force) - -[[ $(jq -r '.tierLabel' <<<"$cfg_tier_result") == "Enterprise" ]] || - fail "Antigravity collector honors config file tier override" "$cfg_tier_result" -pass "Antigravity collector honors config file tier override" - -# Test unauthenticated state -cat >"$TEST_HOME/bin/secret-tool" <<'EOF' -#!/bin/bash -exit 1 -EOF - -unauth_result=$(HOME="$TEST_HOME" AGY_DIR="$AGY_DIR" XDG_CACHE_HOME="$TEST_HOME/.cache" PATH="$TEST_HOME/bin:$PATH" \ - "$ROOT/bin/omarchy-agent-usage-antigravity" --force) - -[[ $(jq -r '.tierLabel' <<<"$unauth_result") == "" ]] || - fail "Antigravity collector has empty tier when unauthenticated" "$unauth_result" -pass "Antigravity collector has empty tier when unauthenticated" - -[[ $(jq -r '.usageStatusText' <<<"$unauth_result") == "Sign in to see limits" ]] || - fail "Antigravity collector indicates sign in status when unauthenticated" "$unauth_result" -pass "Antigravity collector indicates sign in status when unauthenticated" +python3 - <<'PY' +import importlib.machinery +import importlib.util +import json +import os +import sqlite3 +import subprocess +import tempfile +import time +import unittest +from datetime import datetime, timezone +from pathlib import Path +from unittest.mock import patch + +root = Path(os.environ['ROOT']) +script = root / 'bin/omarchy-agent-usage-antigravity' +loader = importlib.machinery.SourceFileLoader('collector', str(script)) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) + +class Scanner(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + self.home = Path(self.tmp.name) + self.app = self.home / 'app' + self.app.mkdir() + self.bin = self.home / 'bin' + self.bin.mkdir() + self.log = self.home / 'calls' + self.cli = self.bin / 'agy' + self.cli.write_text('''#!/usr/bin/python3 +import json, os, sys +from pathlib import Path +with open(os.environ['AGY_TEST_LOG'], 'a') as f: + f.write(json.dumps(sys.argv[1:]) + '\\n') +if sys.argv[1:] == ['--version']: + print(os.environ.get('AGY_TEST_VERSION', '1.2.14')) + sys.exit(0) +assert sys.argv[1:] == ['-p', '/usage', '--output-format', 'json'] +error = os.environ.get('AGY_TEST_ERROR') +if error: + print(error, file=sys.stderr) + sys.exit(1) +print(Path(os.environ['AGY_TEST_FIXTURE']).read_text()) +''') + self.cli.chmod(0o755) + self.env = {'HOME': str(self.home), 'PATH': str(self.bin) + ':/usr/bin', 'TZ': 'Etc/GMT-14', + 'AGY_DIR': str(self.app), 'XDG_CACHE_HOME': str(self.home / 'cache'), + 'AGY_TEST_LOG': str(self.log), 'AGY_TEST_FIXTURE': str(root / 'test/shell.d/fixtures/antigravity/usage.json')} + + def run_collector(self, *args, **extra): + result = subprocess.run([str(script), *args], env={**self.env, **extra}, text=True, capture_output=True, check=True) + self.assertEqual(result.stderr, '') + return json.loads(result.stdout) + + def transcript(self, name, entries): + path = self.app / 'brain' / name / '.system_generated/logs/transcript.jsonl' + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text('\n'.join(json.dumps(e) for e in entries) + '\n{partial') + + def test_local_dates_real_shapes_and_no_invented_tokens(self): + now = datetime(2026, 9, 28, 1, tzinfo=timezone.utc) + prompt = {'step_index': 0, 'source': 'USER', 'type': 'USER_INPUT', 'status': 'DONE', + 'created_at': '2026-09-27T23:00:00Z', 'content': 'Hello'} + self.transcript('session', [prompt, prompt, + {'step_index': 1, 'source': 'MODEL', 'type': 'PLANNER_RESPONSE', 'created_at': '2026-09-27T23:01:00Z', 'content': 'Answer', 'thinking': 'Thinking'}, + {'step_index': 2, 'source': 'USER', 'type': 'USER_INPUT', 'content': 'Missing date'}, + {'step_index': 3, 'type': 'USER_INPUT', 'created_at': '2099-01-01T00:00:00Z'}, + {'step_index': 4, 'created_at': [], 'type': 'USER_INPUT'}]) + history = [ + {'display': 'Hello', 'timestamp': 1759014000000, 'workspace': '/project'}, + {'display': 'Hello', 'timestamp': '2026-09-27T23:00:00Z', 'workspace': '/project'}, + {'display': '/usage', 'timestamp': '2026-09-27T23:00:00Z'}, + {'display': 'old', 'timestamp': '2026-09-27T01:00:00Z'}, + {'timestamp': []}, {'timestamp': 10**30}, []] + # The first record exercises epoch milliseconds for the same instant. + history[0]['timestamp'] = int(datetime(2026, 9, 27, 23, tzinfo=timezone.utc).timestamp() * 1000) + (self.app / 'history.jsonl').write_text('\n'.join(json.dumps(e) for e in history)) + with sqlite3.connect(self.app / 'conversation_summaries.db') as db: + db.execute('CREATE TABLE conversation_summaries (conversation_id TEXT, last_modified_time TEXT)') + db.executemany('INSERT INTO conversation_summaries VALUES (?, ?)', [('session', '2026-09-27 23:02:00+00:00'), ('summary-only', '2026-09-27 23:03:00+00:00')]) + with patch.dict(os.environ, {'TZ': 'Etc/GMT-14'}): + time.tzset() + stats = collector.collect_local_stats(self.app, now) + time.tzset() + self.assertEqual(stats['todayPrompts'], 2) + self.assertEqual(stats['totalPrompts'], 3) + self.assertEqual(stats['todaySessions'], 2) + self.assertEqual(stats['totalSessions'], 2) + self.assertEqual(stats['activeDates'], ['2026-09-27', '2026-09-28']) + self.assertEqual(stats['recentDays'][-1], {'date': '2026-09-28', 'messageCount': 0}) + self.assertEqual(stats['todayTotalTokens'], 0) + self.assertEqual(stats['modelUsage'], {}) + self.assertEqual(stats['todayTokensByModel'], {}) + + def test_negative_timezone(self): + self.transcript('session', [{'step_index': 0, 'type': 'USER_INPUT', 'created_at': '2026-09-28T01:00:00Z'}]) + with patch.dict(os.environ, {'TZ': 'America/New_York'}): + time.tzset() + stats = collector.collect_local_stats(self.app, datetime(2026, 9, 28, 2, tzinfo=timezone.utc)) + time.tzset() + self.assertEqual(stats['activeDates'], ['2026-09-27']) + self.assertEqual(stats['recentDays'][-1]['date'], '2026-09-27') + self.assertEqual(stats['todayPrompts'], 1) + + def test_executable_record_caches_custom_root_and_tier(self): + stamp = datetime.now(timezone.utc).isoformat() + (self.app / 'history.jsonl').write_text(json.dumps({'display': 'hello', 'timestamp': stamp}) + '\n') + record = self.run_collector('--force') + self.assertEqual(record['id'], 'antigravity') + self.assertTrue(record['ready']) + self.assertTrue(record['hasLocalStats']) + self.assertEqual(record['tierLabel'], '') + self.assertEqual(record['usageStatusText'], '') + self.assertEqual([e['percent'] for e in record['limits']], [0, .27, .03]) + self.assertEqual(record['todayPrompts'], 1) + self.assertEqual(record['todayTotalTokens'], 0) + self.assertEqual(record['modelUsage'], {}) + calls = self.log.read_text() + (self.app / 'history.jsonl').write_text((json.dumps({'display': 'hello', 'timestamp': stamp}) + '\n') * 3) + self.assertEqual(self.run_collector('--limits-only')['todayPrompts'], 1) + self.assertEqual(self.log.read_text(), calls) + self.assertEqual(self.run_collector('--force')['todayPrompts'], 3) + stats_path = next((self.home / 'cache/omarchy/agent-usage').glob('*.stats.json')) + cached = json.loads(stats_path.read_text()) + cached['stats'] = {} + stats_path.write_text(json.dumps(cached)) + self.assertEqual(self.run_collector('--limits-only')['todayPrompts'], 3) + other = self.home / 'other' + other.mkdir() + self.assertEqual(self.run_collector(AGY_DIR=str(other))['totalPrompts'], 0) + self.assertEqual(self.run_collector(AGY_TIER='Ultra')['tierLabel'], 'Ultra') + config = self.home / 'config/omarchy/agents' + config.mkdir(parents=True) + (config / 'antigravity.json').write_text('{"tier":"Enterprise"}') + self.assertEqual(self.run_collector(XDG_CONFIG_HOME=str(self.home / 'config'))['tierLabel'], 'Enterprise') + + def test_real_command_errors_and_safety_gate(self): + for error, status, retry in [('401 expired token', 'Antigravity sign-in required', False), + ('403 permission denied', 'Antigravity sign-in required', False), + ('429 rate limit', 'Antigravity limits rate limited', False), + ('network unreachable', 'Antigravity limits unavailable', True), + ('500 server failure', 'Antigravity limits unavailable', False)]: + record = self.run_collector('--force', AGY_TEST_ERROR=error) + self.assertEqual(record['limits'], []) + self.assertEqual(record['tierLabel'], '') + self.assertEqual(record['usageStatusText'], status) + self.assertEqual(record['retryAdvised'], retry) + self.assertTrue(record['authHelpText']) + self.log.write_text('') + record = self.run_collector('--force', AGY_TEST_VERSION='1.1.10') + self.assertEqual(record['usageStatusText'], 'Antigravity update required') + self.assertEqual(self.log.read_text().splitlines(), ['["--version"]']) + self.run_collector('--force') + stale = self.run_collector('--force', AGY_TEST_ERROR='network unavailable') + self.assertTrue(stale['limitsStale']) + self.assertEqual(len(stale['limits']), 3) + self.assertTrue(stale['retryAdvised']) + + def test_stock_wrapper_does_not_install(self): + self.cli.write_text('#!/bin/bash\nmise use -g antigravity-cli\nexit 99\n') + mise = self.bin / 'mise' + mise.write_text('#!/bin/bash\n[[ $* == "which agy" ]] || exit 99\nexit 1\n') + mise.chmod(0o755) + record = self.run_collector('--force') + self.assertEqual(record['usageStatusText'], 'Antigravity CLI unavailable') + self.assertFalse(record['ready']) + self.assertFalse(self.log.exists()) + +unittest.main(verbosity=2) +PY +pass "Antigravity real CLI integration, local history, timezones and cache isolation" diff --git a/test/shell.d/fixtures/antigravity/README.md b/test/shell.d/fixtures/antigravity/README.md new file mode 100644 index 00000000000..f67de585444 --- /dev/null +++ b/test/shell.d/fixtures/antigravity/README.md @@ -0,0 +1,5 @@ +# Antigravity fixtures + +`usage.json` preserves the structure of `agy -p "/usage" --output-format json` from Antigravity CLI 1.2.14, checked on October 1, 2026. Descriptions and redundant text output are omitted, percentages are rounded, reset dates are moved into the future, and the groups are reversed to exercise order independence. The command reports zero agent turns. Google's [1.1.11 release notes](https://github.com/google-antigravity/antigravity-cli/blob/main/CHANGELOG.md#1111) introduce this read-only interface. + +Local scanner fixtures use the observed 1.2.x transcript fields (`step_index`, `source`, `type`, `status`, `created_at`, `content`) and history fields (`display`, integer-millisecond `timestamp`, `workspace`). The installed CLI's transcripts have no model IDs or token usage fields. Tests must not assume they do or estimate tokens from content. SQLite fixtures reproduce only the summary columns the collector reads; conversation databases contain protobuf blobs and are not decoded. diff --git a/test/shell.d/fixtures/antigravity/usage.json b/test/shell.d/fixtures/antigravity/usage.json new file mode 100644 index 00000000000..4e9c294d911 --- /dev/null +++ b/test/shell.d/fixtures/antigravity/usage.json @@ -0,0 +1,25 @@ +{ + "status": "SUCCESS", + "num_turns": 0, + "command": { + "name": "usage", + "data": { + "groups": [ + { + "name": "Claude and GPT models", + "buckets": [ + {"id": "3p-weekly", "name": "Weekly Limit Remaining", "window": "weekly", "remaining_fraction": 0.97, "reset_time": "2099-10-06T02:43:46Z"}, + {"id": "3p-5h", "name": "Five Hour Limit Remaining", "window": "5h", "remaining_fraction": 1, "reset_time": "2099-10-02T07:42:57Z"} + ] + }, + { + "name": "Gemini Models", + "buckets": [ + {"id": "gemini-weekly", "name": "Weekly Limit Remaining", "window": "weekly", "remaining_fraction": 0.73, "reset_time": "2099-10-02T15:41:56Z"}, + {"id": "gemini-5h", "name": "Five Hour Limit Remaining", "window": "5h", "remaining_fraction": 1, "reset_time": "2099-10-02T07:42:57Z"} + ] + } + ] + } + } +} From 6d7a51f76de4d9472bb95fdb2e84384ab8e725a3 Mon Sep 17 00:00:00 2001 From: steph4n <284007938+steph4n-gh@users.noreply.github.com> Date: Thu, 1 Oct 2026 23:05:06 -0400 Subject: [PATCH 3/7] Keep Antigravity quotas current across sign-ins and merge prompt events --- bin/omarchy-agent-usage-antigravity | 79 +++++++++---------- shell/plugins/agents/README.md | 4 +- .../agent-usage-antigravity-limits-test.sh | 65 ++++++--------- .../agent-usage-antigravity-scanner-test.sh | 37 ++++++++- test/shell.d/fixtures/antigravity/README.md | 2 + 5 files changed, 98 insertions(+), 89 deletions(-) diff --git a/bin/omarchy-agent-usage-antigravity b/bin/omarchy-agent-usage-antigravity index 34325a857ca..76a91ac92e6 100755 --- a/bin/omarchy-agent-usage-antigravity +++ b/bin/omarchy-agent-usage-antigravity @@ -23,14 +23,13 @@ import sqlite3 import subprocess import tempfile import time -from collections import Counter +from collections import Counter, defaultdict from datetime import datetime, timedelta, timezone from pathlib import Path SCAN_REUSE_SECONDS = 20 LIMITS_ONLY_REUSE_SECONDS = 900 -PROBE_REUSE_SECONDS = 60 -CACHE_VERSION = 2 +CACHE_VERSION = 3 AUTH_HELP = "Run `agy` to check your sign-in, then refresh the panel." @@ -91,6 +90,7 @@ def collect_local_stats(root, now): days = [(now.astimezone().date() - timedelta(days=i)).isoformat() for i in range(6, -1, -1)] history_prompts = Counter() transcript_prompts = Counter() + history_events = defaultdict(list) active_dates = set() sessions = set() today_sessions = set() @@ -115,6 +115,9 @@ def collect_local_stats(root, now): day = activity(entry.get("timestamp"), entry.get("conversationId") or entry.get("conversation_id")) if day: history_prompts[day] += 1 + stamp = parse_time(entry["timestamp"]).timestamp() + history_events[int(stamp)].append({"time": stamp, "text": display, "matched": False, + "session": entry.get("conversationId") or entry.get("conversation_id")}) for transcript in (root / "brain").glob("*/.system_generated/logs/transcript.jsonl"): session = transcript.parents[2].name @@ -130,7 +133,25 @@ def collect_local_stats(root, now): seen.add(key) day = activity(entry.get("created_at"), session) if day and entry.get("type") == "USER_INPUT": - transcript_prompts[day] += 1 + stamp = parse_time(entry["created_at"]).timestamp() + content = entry.get("content") + candidates = [] + # CLI history timestamps are milliseconds; transcript creation is + # slightly later and may wrap the prompt in context. + # Match one event at a time so distinct prompts from either source + # survive, including repeated text and partially retained logs. + for second in range(int(stamp) - 1, int(stamp) + 2): + for event in history_events.get(second, []): + text = event["text"] + if (not event["matched"] and abs(event["time"] - stamp) <= 1 + and (not event["session"] or event["session"] == session) + and isinstance(text, str) and text.strip() and isinstance(content, str) + and (content.strip() == text.strip() or (content.startswith("") and text.strip() in content))): + candidates.append(event) + if candidates: + min(candidates, key=lambda event: abs(event["time"] - stamp))["matched"] = True + else: + transcript_prompts[day] += 1 # Summaries also know about conversations whose transcript is unavailable. database = root / "conversation_summaries.db" @@ -147,9 +168,7 @@ def collect_local_stats(root, now): except sqlite3.Error: pass - # The same prompt lives in both logs. Per-day maxima use history as a - # fallback without adding it on top of transcript counts. - prompts = {day: max(history_prompts[day], transcript_prompts[day]) for day in active_dates} + prompts = history_prompts + transcript_prompts return { "todayPrompts": prompts.get(today, 0), "todaySessions": len(today_sessions), @@ -230,8 +249,6 @@ def extract_limits(data): if not math.isfinite(remaining) or not 0 <= remaining <= 1: continue used = round(1 - remaining, 4) - if group is not primary and used == 0: - continue reset = parse_time(bucket.get("reset_time")) limits.append({ "label": label if group is primary else f"{scoped} {label}", @@ -286,42 +303,18 @@ def probe_limits(cli): return probe_failure("could not read CLI output") -def collect_limits(cli, path, force): - now = time.time() - cached = read_json(path) - entries = cached.get("limits") - fetched = cached.get("fetchedAtMs", 0) - if not isinstance(fetched, (int, float)) or not math.isfinite(fetched): - fetched = 0 - # Unknown reset times cannot establish that an old window is still open. - fallback = [entry for entry in entries if isinstance(entry, dict) - and isinstance(entry.get("percent"), (int, float)) and 0 <= entry["percent"] <= 1 - and (reset := parse_time(entry.get("resetsAt"))) and reset.timestamp() > now] if isinstance(entries, list) else [] - result = {"limits": fallback, "limitsStale": bool(fallback), "limitsFetchedAt": fetched, +def collect_limits(cli): + # /usage provides no account identifier. Reusing quotas or errors across + # calls could attribute a previous sign-in's result to the current account. + # Always ask the CLI and leave the meters empty when it cannot answer. + result = {"limits": [], "limitsStale": False, "limitsFetchedAt": 0, "usageStatusText": "", "authHelpText": "", "retryAdvised": False} - checked = cached.get("checkedAtMs", 0) - reuse = 15 if cached.get("retryAdvised") else PROBE_REUSE_SECONDS if not cli: result.update(usageStatusText="Antigravity CLI unavailable", authHelpText="Install Antigravity CLI and run `agy` to sign in.") - elif not force and isinstance(checked, (int, float)) and 0 <= now - checked / 1000 < reuse: - result.update({key: cached[key] for key in ("usageStatusText", "authHelpText", "retryAdvised") if key in cached}) - result["limitsStale"] = bool(fallback) and cached.get("limitsStale", False) - if not fallback and not result["usageStatusText"]: - # A window rolled over during the reuse interval; get its new value. - return collect_limits(cli, path, True) - return result else: - probe = probe_limits(cli) - result.update(probe) - if probe.get("limits"): - result.update(limitsStale=False, limitsFetchedAt=round(time.time() * 1000)) - elif result["usageStatusText"] == "Antigravity sign-in required": - # The CLI may have signed out or switched accounts. Do not attribute a - # previous account's numbers to the next successful sign-in. - result.update(limits=[], limitsStale=False, limitsFetchedAt=0) - if result["limitsStale"]: - result["authHelpText"] += " Showing the last known limits until their windows reset." - write_json(path, {**result, "fetchedAtMs": result["limitsFetchedAt"], "checkedAtMs": round(time.time() * 1000)}) + result.update(probe_limits(cli)) + if result["limits"]: + result["limitsFetchedAt"] = round(time.time() * 1000) return result @@ -334,7 +327,7 @@ def tier_label(): def main(): parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--force", action="store_true", help="Bypass local caches and read fresh quotas") + parser.add_argument("--force", action="store_true", help="Bypass the history cache; quotas are always read fresh") parser.add_argument("--limits-only", action="store_true", help="Reuse a recent history scan") args = parser.parse_args() root = app_root() @@ -356,7 +349,7 @@ def main(): or not isinstance(fetched, (int, float)) or not 0 <= time.time() - fetched < age): stats = collect_local_stats(root, now) write_json(stats_path, {"scanDate": now.date().isoformat(), "fetchedAt": time.time(), "stats": stats}) - limits = collect_limits(installed_cli(), prefix.with_suffix(".limits.json"), args.force) + limits = collect_limits(installed_cli()) has_stats = stats["totalPrompts"] > 0 or stats["totalSessions"] > 0 record = {"schemaVersion": 1, "id": "antigravity", "name": "Antigravity", "updatedAt": datetime.now(timezone.utc).isoformat(), "ready": has_stats or bool(limits["limits"]), diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 7f792ae4a06..45e2de82541 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -106,9 +106,9 @@ signed in there. The collector runs `agy -p "/usage" --output-format json` after checking that the installed CLI is at least 1.1.11. That version introduced read-only slash-command output without an agent turn or quota spend; an older or unrecognized version is never sent `/usage`. It resolves an existing mise installation without invoking Omarchy’s installer wrapper. Antigravity itself handles credential storage, token refresh and the account’s quota request. -Session and weekly meters use the returned quota windows, including used Claude/GPT allowances. Missing percentages are unknown, not zero usage. Failed checks show a status and help text; network failures request an early retry. Last known windows remain visible and dimmed only until they reset, and an authentication failure clears them. Refresh bypasses the short reuse interval. `AGY_DIR` selects another app data directory with its own caches. +Session and weekly meters use the returned quota windows, including unused Claude/GPT allowances. Missing percentages are unknown, not zero usage. Failed checks show a status and help text; network failures request an early retry. Every refresh asks the CLI again. Because the response provides no account identity, quota results and failures are never cached: switching accounts or signing in again takes effect on the next check, and a failed check leaves the meters empty. `AGY_DIR` selects another app data directory with its own history cache. -The CLI’s transcripts do not contain token usage or model IDs. Local history contributes prompt and session counts on local calendar dates, without character-based token estimates or invented model rows. History and transcript prompt counts are combined by taking the larger count for each day, since the history may not carry conversation IDs. The `/usage` response does not include a subscription tier, so the plan stays unknown unless you supply `AGY_TIER` or `{ "tier": "Ultra" }` in `~/.config/omarchy/agents/antigravity.json` (honoring `XDG_CONFIG_HOME`). Sign in through `agy`; the panel’s built-in account switching remains for Claude, Codex and Grok. +The CLI’s transcripts do not contain token usage or model IDs. Local history contributes prompt and session counts on local calendar dates, without character-based token estimates or invented model rows. History and transcript prompt events are combined individually, matching text and timestamps within one second (and conversation IDs when present) to avoid counting the same prompt twice. Transcript `` wrappers are recognized; unmatched events from either source and repeated prompts at different times are retained. The `/usage` response does not include a subscription tier, so the plan stays unknown unless you supply `AGY_TIER` or `{ "tier": "Ultra" }` in `~/.config/omarchy/agents/antigravity.json` (honoring `XDG_CONFIG_HOME`). Sign in through `agy`; the panel’s built-in account switching remains for Claude, Codex and Grok. ### Fireworks balance diff --git a/test/shell.d/agent-usage-antigravity-limits-test.sh b/test/shell.d/agent-usage-antigravity-limits-test.sh index 54852c3eac1..871457f6632 100755 --- a/test/shell.d/agent-usage-antigravity-limits-test.sh +++ b/test/shell.d/agent-usage-antigravity-limits-test.sh @@ -9,7 +9,6 @@ import importlib.machinery import importlib.util import json import os -import tempfile import unittest from pathlib import Path from unittest.mock import patch @@ -25,8 +24,8 @@ class Limits(unittest.TestCase): def test_real_cli_shape_and_reversed_groups(self): data = fixture['command']['data'] limits = collector.extract_limits(data) - self.assertEqual([v['title'] for v in limits], ['Session', 'Weekly', 'Claude/GPT Weekly']) - self.assertEqual([v['percent'] for v in limits], [0, .27, .03]) + self.assertEqual([v['title'] for v in limits], ['Session', 'Weekly', 'Claude/GPT Session', 'Claude/GPT Weekly']) + self.assertEqual([v['percent'] for v in limits], [0, .27, 0, .03]) reversed_data = {'groups': list(reversed(data['groups']))} self.assertEqual(limits, collector.extract_limits(reversed_data)) self.assertEqual(limits[0]['resetsAt'], '2099-10-02T07:42:57+00:00') @@ -44,7 +43,8 @@ class Limits(unittest.TestCase): def test_only_scoped_group_keeps_its_identity(self): data = {'groups': [fixture['command']['data']['groups'][0]]} - self.assertEqual(collector.extract_limits(data)[0]['title'], 'Claude/GPT Weekly') + self.assertEqual([limit['title'] for limit in collector.extract_limits(data)], ['Claude/GPT Session', 'Claude/GPT Weekly']) + self.assertEqual(collector.extract_limits(data)[0]['percent'], 0) def test_version_gate_never_prompts_an_old_or_unknown_cli(self): for version in ['1.1.10', '1.0.99', 'unknown', '1.1.11-preview', '']: @@ -75,43 +75,28 @@ class Limits(unittest.TestCase): with patch.object(collector, 'run_cli', side_effect=collector.subprocess.TimeoutExpired('agy', 20)): self.assertTrue(collector.probe_limits('agy')['retryAdvised']) - def test_stale_cache_expiry_auth_clear_and_force(self): + def test_every_check_reads_the_current_sign_in(self): good = {'limits': collector.extract_limits(fixture['command']['data']), 'usageStatusText': '', 'authHelpText': ''} - with tempfile.TemporaryDirectory() as tmp: - cache = Path(tmp) / 'limits.json' - with patch.object(collector, 'probe_limits', return_value=good) as probe: - live = collector.collect_limits('agy', cache, False) - self.assertFalse(live['limitsStale']) - self.assertGreater(live['limitsFetchedAt'], 0) - self.assertEqual(collector.collect_limits('agy', cache, False)['limits'], live['limits']) - self.assertEqual(probe.call_count, 1) - collector.collect_limits('agy', cache, True) - self.assertEqual(probe.call_count, 2) - failure = collector.probe_failure('network unavailable') - with patch.object(collector, 'probe_limits', return_value=failure) as probe: - stale = collector.collect_limits('agy', cache, True) - self.assertEqual(stale['limits'], live['limits']) - self.assertTrue(stale['limitsStale']) - self.assertTrue(stale['retryAdvised']) - self.assertIn('last known limits', stale['authHelpText']) - reused = collector.collect_limits('agy', cache, False) - self.assertTrue(reused['limitsStale']) - self.assertEqual(probe.call_count, 1) - cached = json.loads(cache.read_text()) - cached['limits'][0]['resetsAt'] = '2000-01-01T00:00:00Z' - cached['limits'][1]['resetsAt'] = '' - collector.write_json(cache, cached) - remaining = collector.collect_limits('agy', cache, True) - self.assertEqual(len(remaining['limits']), 1) - with patch.object(collector, 'probe_limits', return_value=collector.probe_failure('401')): - signed_out = collector.collect_limits('agy', cache, True) - self.assertEqual(signed_out['limits'], []) - self.assertEqual(json.loads(cache.read_text())['limits'], []) - for invalid in ['[]', '{bad', '{"limits": 5, "fetchedAtMs": "bad"}']: - cache.write_text(invalid) - with patch.object(collector, 'probe_limits', return_value=failure): - self.assertEqual(collector.collect_limits('agy', cache, True)['limits'], []) + other = {'limits': [{'title': 'Session', 'percent': .8, 'resetsAt': '2099-01-01T00:00:00Z'}], 'usageStatusText': '', 'authHelpText': ''} + failure = collector.probe_failure('network unavailable') + auth = collector.probe_failure('401') + with patch.object(collector, 'probe_limits', side_effect=[good, other, failure, auth, good]) as probe: + first = collector.collect_limits('agy') + second = collector.collect_limits('agy') + self.assertEqual(first['limits'], good['limits']) + self.assertEqual(second['limits'], other['limits']) + self.assertGreater(second['limitsFetchedAt'], 0) + offline = collector.collect_limits('agy') + self.assertEqual(offline['limits'], []) + self.assertTrue(offline['retryAdvised']) + self.assertEqual(offline['limitsFetchedAt'], 0) + signed_out = collector.collect_limits('agy') + self.assertEqual(signed_out['limits'], []) + signed_in = collector.collect_limits('agy') + self.assertEqual(signed_in['usageStatusText'], '') + self.assertEqual(signed_in['limits'], good['limits']) + self.assertEqual(probe.call_count, 5) unittest.main(verbosity=2) PY -pass "Antigravity quota parsing, version safety, failures and cache freshness" +pass "Antigravity quota parsing, version safety, failures and account changes" diff --git a/test/shell.d/agent-usage-antigravity-scanner-test.sh b/test/shell.d/agent-usage-antigravity-scanner-test.sh index 17fffc921af..d698de956f9 100755 --- a/test/shell.d/agent-usage-antigravity-scanner-test.sh +++ b/test/shell.d/agent-usage-antigravity-scanner-test.sh @@ -120,14 +120,14 @@ print(Path(os.environ['AGY_TEST_FIXTURE']).read_text()) self.assertTrue(record['hasLocalStats']) self.assertEqual(record['tierLabel'], '') self.assertEqual(record['usageStatusText'], '') - self.assertEqual([e['percent'] for e in record['limits']], [0, .27, .03]) + self.assertEqual([e['percent'] for e in record['limits']], [0, .27, 0, .03]) self.assertEqual(record['todayPrompts'], 1) self.assertEqual(record['todayTotalTokens'], 0) self.assertEqual(record['modelUsage'], {}) calls = self.log.read_text() (self.app / 'history.jsonl').write_text((json.dumps({'display': 'hello', 'timestamp': stamp}) + '\n') * 3) self.assertEqual(self.run_collector('--limits-only')['todayPrompts'], 1) - self.assertEqual(self.log.read_text(), calls) + self.assertEqual(len(self.log.read_text().splitlines()), len(calls.splitlines()) + 2) self.assertEqual(self.run_collector('--force')['todayPrompts'], 3) stats_path = next((self.home / 'cache/omarchy/agent-usage').glob('*.stats.json')) cached = json.loads(stats_path.read_text()) @@ -161,10 +161,39 @@ print(Path(os.environ['AGY_TEST_FIXTURE']).read_text()) self.assertEqual(self.log.read_text().splitlines(), ['["--version"]']) self.run_collector('--force') stale = self.run_collector('--force', AGY_TEST_ERROR='network unavailable') - self.assertTrue(stale['limitsStale']) - self.assertEqual(len(stale['limits']), 3) + self.assertFalse(stale['limitsStale']) + self.assertEqual(stale['limits'], []) self.assertTrue(stale['retryAdvised']) + def test_partial_and_disjoint_history_preserve_distinct_prompts(self): + history = [{'display': 'shared', 'timestamp': '2026-09-27T12:00:00Z'}, + {'display': 'history only', 'timestamp': '2026-09-27T12:01:00Z'}] + (self.app / 'history.jsonl').write_text('\n'.join(json.dumps(e) for e in history)) + for shared, expected in [('shared', 4), ('different prompt', 5)]: + self.transcript('session', [ + {'step_index': 0, 'type': 'USER_INPUT', 'created_at': '2026-09-27T12:00:00.400Z', 'content': '\n' + shared + '\n'}, + {'step_index': 1, 'type': 'USER_INPUT', 'created_at': '2026-09-27T12:02:00Z', 'content': 'transcript only'}, + {'step_index': 2, 'type': 'USER_INPUT', 'created_at': '2026-09-27T12:03:00Z', 'content': 'shared'}]) + with patch.dict(os.environ, {'TZ': 'UTC'}): + time.tzset() + stats = collector.collect_local_stats(self.app, datetime(2026, 9, 27, 14, tzinfo=timezone.utc)) + time.tzset() + self.assertEqual(stats['todayPrompts'], expected) + self.assertEqual(stats['totalPrompts'], expected) + + def test_account_switch_and_new_sign_in_are_visible_immediately(self): + first = self.run_collector() + other = json.loads(Path(self.env['AGY_TEST_FIXTURE']).read_text()) + other['command']['data']['groups'][1]['buckets'][1]['remaining_fraction'] = .2 + other_file = self.home / 'other-account.json' + other_file.write_text(json.dumps(other)) + second = self.run_collector(AGY_TEST_FIXTURE=str(other_file)) + self.assertEqual(first['limits'][0]['percent'], 0) + self.assertEqual(second['limits'][0]['percent'], .8) + self.assertEqual(self.run_collector(AGY_TEST_ERROR='401 signed out')['limits'], []) + self.assertEqual(self.run_collector()['usageStatusText'], '') + self.assertFalse(list((self.home / 'cache').rglob('*.limits.json'))) + def test_stock_wrapper_does_not_install(self): self.cli.write_text('#!/bin/bash\nmise use -g antigravity-cli\nexit 99\n') mise = self.bin / 'mise' diff --git a/test/shell.d/fixtures/antigravity/README.md b/test/shell.d/fixtures/antigravity/README.md index f67de585444..c9452c05df9 100644 --- a/test/shell.d/fixtures/antigravity/README.md +++ b/test/shell.d/fixtures/antigravity/README.md @@ -3,3 +3,5 @@ `usage.json` preserves the structure of `agy -p "/usage" --output-format json` from Antigravity CLI 1.2.14, checked on October 1, 2026. Descriptions and redundant text output are omitted, percentages are rounded, reset dates are moved into the future, and the groups are reversed to exercise order independence. The command reports zero agent turns. Google's [1.1.11 release notes](https://github.com/google-antigravity/antigravity-cli/blob/main/CHANGELOG.md#1111) introduce this read-only interface. Local scanner fixtures use the observed 1.2.x transcript fields (`step_index`, `source`, `type`, `status`, `created_at`, `content`) and history fields (`display`, integer-millisecond `timestamp`, `workspace`). The installed CLI's transcripts have no model IDs or token usage fields. Tests must not assume they do or estimate tokens from content. SQLite fixtures reproduce only the summary columns the collector reads; conversation databases contain protobuf blobs and are not decoded. + +Prompt matching also covers `` wrappers observed in local transcripts. Matched history/transcript events in the inspected logs were less than one second apart; prompt text is only compared in memory and is never written to the usage cache. From 74c6361ea80100db6b6f8f42a5454a5ff0679441 Mon Sep 17 00:00:00 2001 From: steph4n <284007938+steph4n-gh@users.noreply.github.com> Date: Thu, 1 Oct 2026 23:11:18 -0400 Subject: [PATCH 4/7] Match complete Antigravity prompts and release scan locks before quotas --- bin/omarchy-agent-usage-antigravity | 17 +++++++--- shell/plugins/agents/README.md | 2 +- .../agent-usage-antigravity-scanner-test.sh | 34 ++++++++++++++++++- 3 files changed, 47 insertions(+), 6 deletions(-) diff --git a/bin/omarchy-agent-usage-antigravity b/bin/omarchy-agent-usage-antigravity index 76a91ac92e6..65f2cc3e1b8 100755 --- a/bin/omarchy-agent-usage-antigravity +++ b/bin/omarchy-agent-usage-antigravity @@ -85,6 +85,16 @@ def json_lines(path): return +def prompt_text(content): + if not isinstance(content, str): + return "" + content = content.strip() + # The request is followed by metadata that can quote unrelated prompts. + # Compare the whole request, never a substring of that surrounding text. + wrapped = re.match(r"\s*(.*?)\s*", content, re.S) + return wrapped[1].strip() if wrapped else content + + def collect_local_stats(root, now): today = now.astimezone().date().isoformat() days = [(now.astimezone().date() - timedelta(days=i)).isoformat() for i in range(6, -1, -1)] @@ -134,7 +144,7 @@ def collect_local_stats(root, now): day = activity(entry.get("created_at"), session) if day and entry.get("type") == "USER_INPUT": stamp = parse_time(entry["created_at"]).timestamp() - content = entry.get("content") + content = prompt_text(entry.get("content")) candidates = [] # CLI history timestamps are milliseconds; transcript creation is # slightly later and may wrap the prompt in context. @@ -145,8 +155,7 @@ def collect_local_stats(root, now): text = event["text"] if (not event["matched"] and abs(event["time"] - stamp) <= 1 and (not event["session"] or event["session"] == session) - and isinstance(text, str) and text.strip() and isinstance(content, str) - and (content.strip() == text.strip() or (content.startswith("") and text.strip() in content))): + and isinstance(text, str) and text.strip() and content == text.strip()): candidates.append(event) if candidates: min(candidates, key=lambda event: abs(event["time"] - stamp))["matched"] = True @@ -349,7 +358,7 @@ def main(): or not isinstance(fetched, (int, float)) or not 0 <= time.time() - fetched < age): stats = collect_local_stats(root, now) write_json(stats_path, {"scanDate": now.date().isoformat(), "fetchedAt": time.time(), "stats": stats}) - limits = collect_limits(installed_cli()) + limits = collect_limits(installed_cli()) has_stats = stats["totalPrompts"] > 0 or stats["totalSessions"] > 0 record = {"schemaVersion": 1, "id": "antigravity", "name": "Antigravity", "updatedAt": datetime.now(timezone.utc).isoformat(), "ready": has_stats or bool(limits["limits"]), diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 45e2de82541..27d35cda832 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -108,7 +108,7 @@ The collector runs `agy -p "/usage" --output-format json` after checking that th Session and weekly meters use the returned quota windows, including unused Claude/GPT allowances. Missing percentages are unknown, not zero usage. Failed checks show a status and help text; network failures request an early retry. Every refresh asks the CLI again. Because the response provides no account identity, quota results and failures are never cached: switching accounts or signing in again takes effect on the next check, and a failed check leaves the meters empty. `AGY_DIR` selects another app data directory with its own history cache. -The CLI’s transcripts do not contain token usage or model IDs. Local history contributes prompt and session counts on local calendar dates, without character-based token estimates or invented model rows. History and transcript prompt events are combined individually, matching text and timestamps within one second (and conversation IDs when present) to avoid counting the same prompt twice. Transcript `` wrappers are recognized; unmatched events from either source and repeated prompts at different times are retained. The `/usage` response does not include a subscription tier, so the plan stays unknown unless you supply `AGY_TIER` or `{ "tier": "Ultra" }` in `~/.config/omarchy/agents/antigravity.json` (honoring `XDG_CONFIG_HOME`). Sign in through `agy`; the panel’s built-in account switching remains for Claude, Codex and Grok. +The CLI’s transcripts do not contain token usage or model IDs. Local history contributes prompt and session counts on local calendar dates, without character-based token estimates or invented model rows. History and transcript prompt events are combined individually, matching the complete prompt text and timestamps within one second (and conversation IDs when present) to avoid counting the same prompt twice. Transcript `` wrappers are recognized and surrounding metadata is excluded; unmatched events from either source and repeated prompts at different times are retained. The `/usage` response does not include a subscription tier, so the plan stays unknown unless you supply `AGY_TIER` or `{ "tier": "Ultra" }` in `~/.config/omarchy/agents/antigravity.json` (honoring `XDG_CONFIG_HOME`). Sign in through `agy`; the panel’s built-in account switching remains for Claude, Codex and Grok. ### Fireworks balance diff --git a/test/shell.d/agent-usage-antigravity-scanner-test.sh b/test/shell.d/agent-usage-antigravity-scanner-test.sh index d698de956f9..0bfe4f9b764 100755 --- a/test/shell.d/agent-usage-antigravity-scanner-test.sh +++ b/test/shell.d/agent-usage-antigravity-scanner-test.sh @@ -5,6 +5,7 @@ source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" require_command python3 python3 - <<'PY' +import fcntl import importlib.machinery import importlib.util import json @@ -37,7 +38,7 @@ class Scanner(unittest.TestCase): self.log = self.home / 'calls' self.cli = self.bin / 'agy' self.cli.write_text('''#!/usr/bin/python3 -import json, os, sys +import json, os, sys, time from pathlib import Path with open(os.environ['AGY_TEST_LOG'], 'a') as f: f.write(json.dumps(sys.argv[1:]) + '\\n') @@ -45,6 +46,9 @@ if sys.argv[1:] == ['--version']: print(os.environ.get('AGY_TEST_VERSION', '1.2.14')) sys.exit(0) assert sys.argv[1:] == ['-p', '/usage', '--output-format', 'json'] +if os.environ.get('AGY_TEST_STARTED'): + Path(os.environ['AGY_TEST_STARTED']).touch() + time.sleep(2) error = os.environ.get('AGY_TEST_ERROR') if error: print(error, file=sys.stderr) @@ -194,6 +198,34 @@ print(Path(os.environ['AGY_TEST_FIXTURE']).read_text()) self.assertEqual(self.run_collector()['usageStatusText'], '') self.assertFalse(list((self.home / 'cache').rglob('*.limits.json'))) + def test_prompt_matching_ignores_substrings_and_metadata(self): + (self.app / 'history.jsonl').write_text(json.dumps({'display': 'hello', 'timestamp': '2026-09-27T12:00:00Z'})) + for content, expected in [ + ('hellocontext', 1), + ('hello worldcontext', 2), + ('differenthello', 2)]: + self.transcript('session', [{'step_index': 0, 'type': 'USER_INPUT', 'created_at': '2026-09-27T12:00:00.400Z', 'content': content}]) + stats = collector.collect_local_stats(self.app, datetime(2026, 9, 27, 14, tzinfo=timezone.utc)) + self.assertEqual(stats['totalPrompts'], expected) + + def test_slow_quota_query_does_not_hold_the_history_lock(self): + started = self.home / 'quota-started' + process = subprocess.Popen([str(script)], env={**self.env, 'AGY_TEST_STARTED': str(started)}, + stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + try: + deadline = time.monotonic() + 5 + while not started.exists() and process.poll() is None and time.monotonic() < deadline: + time.sleep(.02) + self.assertTrue(started.exists(), 'quota query started') + lock_path = next((self.home / 'cache').rglob('*.lock')) + with lock_path.open('w') as lock: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + fcntl.flock(lock, fcntl.LOCK_UN) + finally: + stdout, stderr = process.communicate(timeout=5) + self.assertEqual(process.returncode, 0, stderr) + self.assertEqual(json.loads(stdout)['usageStatusText'], '') + def test_stock_wrapper_does_not_install(self): self.cli.write_text('#!/bin/bash\nmise use -g antigravity-cli\nexit 99\n') mise = self.bin / 'mise' From f46ada89c1088e2793363a0204c8dcee31a5c05e Mon Sep 17 00:00:00 2001 From: steph4n <284007938+steph4n-gh@users.noreply.github.com> Date: Thu, 1 Oct 2026 23:19:34 -0400 Subject: [PATCH 5/7] Prevent stale usage publication and match literal prompt delimiters --- bin/omarchy-agent-usage-antigravity | 24 +++++---- bin/omarchy-agent-usage-update | 29 ++++++++--- shell/plugins/agents/README.md | 5 +- .../agent-usage-antigravity-scanner-test.sh | 9 ++++ test/shell.d/agent-usage-update-test.sh | 49 +++++++++++++++++++ 5 files changed, 99 insertions(+), 17 deletions(-) diff --git a/bin/omarchy-agent-usage-antigravity b/bin/omarchy-agent-usage-antigravity index 65f2cc3e1b8..c8cb156fb91 100755 --- a/bin/omarchy-agent-usage-antigravity +++ b/bin/omarchy-agent-usage-antigravity @@ -85,14 +85,18 @@ def json_lines(path): return -def prompt_text(content): - if not isinstance(content, str): - return "" - content = content.strip() - # The request is followed by metadata that can quote unrelated prompts. - # Compare the whole request, never a substring of that surrounding text. - wrapped = re.match(r"\s*(.*?)\s*", content, re.S) - return wrapped[1].strip() if wrapped else content +def prompt_matches(text, content): + if not isinstance(text, str) or not text.strip() or not isinstance(content, str): + return False + text, content = text.strip(), content.strip() + if text == content: + return True + # Match the complete candidate literally between the wrapper delimiters. + # A prompt can itself contain closing-tag text; splitting at the first + # closing tag would truncate it. Metadata may quote unrelated prompts. + return re.fullmatch(r"\s*" + re.escape(text) + + r"\s*(?:\s*.*)?", + content, re.S) is not None def collect_local_stats(root, now): @@ -144,7 +148,7 @@ def collect_local_stats(root, now): day = activity(entry.get("created_at"), session) if day and entry.get("type") == "USER_INPUT": stamp = parse_time(entry["created_at"]).timestamp() - content = prompt_text(entry.get("content")) + content = entry.get("content") candidates = [] # CLI history timestamps are milliseconds; transcript creation is # slightly later and may wrap the prompt in context. @@ -155,7 +159,7 @@ def collect_local_stats(root, now): text = event["text"] if (not event["matched"] and abs(event["time"] - stamp) <= 1 and (not event["session"] or event["session"] == session) - and isinstance(text, str) and text.strip() and content == text.strip()): + and prompt_matches(text, content)): candidates.append(event) if candidates: min(candidates, key=lambda event: abs(event["time"] - stamp))["matched"] = True diff --git a/bin/omarchy-agent-usage-update b/bin/omarchy-agent-usage-update index e759dc72459..e5dbbf39661 100755 --- a/bin/omarchy-agent-usage-update +++ b/bin/omarchy-agent-usage-update @@ -39,17 +39,34 @@ wanted() { return 1 } -collect() { +collect() ( local collector="$1" agent="$2" - local record tmp + local record tmp pending="$USAGE_DIR/.$agent.pending" lock="$USAGE_DIR/.$agent.lock" + tmp=$(mktemp "$USAGE_DIR/.$agent.XXXXXX") || return 1 + trap 'rm -f "$tmp"' EXIT + + # Overlapping refreshes may finish in either order, including across + # account changes. Reserve publication before probing; only the latest + # started check may replace the record. Hold the lock only for these short + # file operations, never while waiting for the collector's network calls. + { + flock 9 || return 1 + printf '%s\n' "$tmp" >"$pending" + } 9>"$lock" || return 1 + if ! record=$("$collector" "${flags[@]}") || [[ -z $record ]] || ! jq -e . >/dev/null 2>&1 <<<"$record"; then echo "omarchy-agent-usage-update: $agent collector failed" >&2 return 1 fi - tmp=$(mktemp "$USAGE_DIR/.$agent.XXXXXX") - printf '%s\n' "$record" >"$tmp" - mv "$tmp" "$USAGE_DIR/$agent.json" -} + printf '%s\n' "$record" >"$tmp" || return 1 + { + flock 9 || return 1 + if [[ -f $pending && $(<"$pending") == "$tmp" ]]; then + mv "$tmp" "$USAGE_DIR/$agent.json" || return 1 + rm -f "$pending" + fi + } 9>"$lock" +) # Accounts can be signed in to someone else since they were added; the # records name them from the registry, so bring it up to date first. diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 27d35cda832..76c23db1af4 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -60,7 +60,10 @@ Each agent is one JSON record in `~/.local/state/omarchy/agents/usage/`, written by `omarchy-agent-usage-update`. That command runs one `omarchy-agent-usage-` collector per agent; the widget invokes it on its refresh timer and whenever you ask for a refresh, and picks up any -record that lands in the directory regardless of who wrote it. +record that lands in the directory regardless of who wrote it. Overlapping +updates reserve publication before collecting: only the latest started check +for each agent can replace its record, so a slow older result cannot overwrite +a newer account’s limits. The publication lock is released during collection. Adding an agent therefore never touches this plugin: ship a collector that prints the record contract (see the `claude` and `codex` collectors in diff --git a/test/shell.d/agent-usage-antigravity-scanner-test.sh b/test/shell.d/agent-usage-antigravity-scanner-test.sh index 0bfe4f9b764..0d3f268169e 100755 --- a/test/shell.d/agent-usage-antigravity-scanner-test.sh +++ b/test/shell.d/agent-usage-antigravity-scanner-test.sh @@ -226,6 +226,15 @@ print(Path(os.environ['AGY_TEST_FIXTURE']).read_text()) self.assertEqual(process.returncode, 0, stderr) self.assertEqual(json.loads(stdout)['usageStatusText'], '') + def test_literal_wrapper_tags_in_prompt_count_once(self): + for text in ['Explain please', 'nested', + 'Document literally']: + (self.app / 'history.jsonl').write_text(json.dumps({'display': text, 'timestamp': '2026-09-27T12:00:00Z'})) + wrapped = '\n' + text + '\n\ncontext' + self.transcript('session', [{'step_index': 0, 'type': 'USER_INPUT', 'created_at': '2026-09-27T12:00:00.400Z', 'content': wrapped}]) + stats = collector.collect_local_stats(self.app, datetime(2026, 9, 27, 14, tzinfo=timezone.utc)) + self.assertEqual(stats['totalPrompts'], 1) + def test_stock_wrapper_does_not_install(self): self.cli.write_text('#!/bin/bash\nmise use -g antigravity-cli\nexit 99\n') mise = self.bin / 'mise' diff --git a/test/shell.d/agent-usage-update-test.sh b/test/shell.d/agent-usage-update-test.sh index ee3aaf0112e..65ed21ffd05 100755 --- a/test/shell.d/agent-usage-update-test.sh +++ b/test/shell.d/agent-usage-update-test.sh @@ -63,3 +63,52 @@ pass "update succeeds when the requested collectors all pass" [[ -e $usage_dir/skipped.json && ! -e $usage_dir/noisy.json ]] || fail "update with agent arguments only runs the named collectors" pass "update with agent arguments only runs the named collectors" + +require_command python3 +require_command flock + +TEST_HOME="$TEST_HOME" FAKE_OMARCHY="$FAKE_OMARCHY" python3 - <<'PY' +import json +import os +import subprocess +import time +from pathlib import Path + +home = Path(os.environ['TEST_HOME']) +fake = Path(os.environ['FAKE_OMARCHY']) +updater = Path(os.environ['ROOT']) / 'bin/omarchy-agent-usage-update' +collector = fake / 'bin/omarchy-agent-usage-racy' +collector.write_text('''#!/bin/bash +touch "$STARTED" +if [[ -n $WAIT_FOR ]]; then + while [[ ! -e $WAIT_FOR ]]; do sleep .02; done +fi +printf '%s\\n' "$RECORD" +''') +collector.chmod(0o755) +env = {**os.environ, 'HOME': str(home), 'OMARCHY_PATH': str(fake), 'XDG_STATE_HOME': ''} +usage_dir = home / '.local/state/omarchy/agents/usage' +# An old successful answer and an old error must both lose to a newer check. +for index, old in enumerate([{'id': 'racy', 'limits': [{'percent': .1}]}, + {'id': 'racy', 'limits': [], 'usageStatusText': 'offline'}]): + started = home / f'started-{index}' + release = home / f'release-{index}' + older = subprocess.Popen([str(updater), 'racy'], env={**env, 'STARTED': str(started), 'WAIT_FOR': str(release), 'RECORD': json.dumps(old)}, + stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + try: + deadline = time.monotonic() + 5 + while not started.exists() and older.poll() is None and time.monotonic() < deadline: + time.sleep(.02) + assert started.exists(), 'older collector started' + new = {'id': 'racy', 'limits': [{'percent': .8}], 'usageStatusText': ''} + subprocess.run([str(updater), 'racy'], env={**env, 'STARTED': str(home / 'newer-started'), 'WAIT_FOR': '', 'RECORD': json.dumps(new)}, + capture_output=True, text=True, check=True, timeout=5) + assert json.loads((usage_dir / 'racy.json').read_text()) == new + finally: + release.touch() + stdout, stderr = older.communicate(timeout=5) + assert older.returncode == 0, stderr + assert json.loads((usage_dir / 'racy.json').read_text()) == new, 'older completion replaced newer data' + assert sorted(p.name for p in usage_dir.glob('.racy.*')) == ['.racy.lock'], 'temporary records are cleaned up' +print('ok - overlapping updates publish only the latest started check without blocking its probe') +PY From 8d3ae7586f59a996cd4dedf7f93f61ccdeecea1b Mon Sep 17 00:00:00 2001 From: steph4n <284007938+steph4n-gh@users.noreply.github.com> Date: Thu, 1 Oct 2026 23:29:16 -0400 Subject: [PATCH 6/7] Handle appended Antigravity settings metadata --- bin/omarchy-agent-usage-antigravity | 3 ++- test/shell.d/agent-usage-antigravity-scanner-test.sh | 1 + test/shell.d/fixtures/antigravity/README.md | 2 ++ 3 files changed, 5 insertions(+), 1 deletion(-) diff --git a/bin/omarchy-agent-usage-antigravity b/bin/omarchy-agent-usage-antigravity index c8cb156fb91..f66826a7b2e 100755 --- a/bin/omarchy-agent-usage-antigravity +++ b/bin/omarchy-agent-usage-antigravity @@ -95,7 +95,8 @@ def prompt_matches(text, content): # A prompt can itself contain closing-tag text; splitting at the first # closing tag would truncate it. Metadata may quote unrelated prompts. return re.fullmatch(r"\s*" + re.escape(text) - + r"\s*(?:\s*.*)?", + + r"\s*(?:\s*.*)?" + + r"(?:\s*.*)?", content, re.S) is not None diff --git a/test/shell.d/agent-usage-antigravity-scanner-test.sh b/test/shell.d/agent-usage-antigravity-scanner-test.sh index 0d3f268169e..a0e6d8ffec9 100755 --- a/test/shell.d/agent-usage-antigravity-scanner-test.sh +++ b/test/shell.d/agent-usage-antigravity-scanner-test.sh @@ -202,6 +202,7 @@ print(Path(os.environ['AGY_TEST_FIXTURE']).read_text()) (self.app / 'history.jsonl').write_text(json.dumps({'display': 'hello', 'timestamp': '2026-09-27T12:00:00Z'})) for content, expected in [ ('hellocontext', 1), + ('hellocontext\nmodel changed', 1), ('hello worldcontext', 2), ('differenthello', 2)]: self.transcript('session', [{'step_index': 0, 'type': 'USER_INPUT', 'created_at': '2026-09-27T12:00:00.400Z', 'content': content}]) diff --git a/test/shell.d/fixtures/antigravity/README.md b/test/shell.d/fixtures/antigravity/README.md index c9452c05df9..d26d0550f28 100644 --- a/test/shell.d/fixtures/antigravity/README.md +++ b/test/shell.d/fixtures/antigravity/README.md @@ -5,3 +5,5 @@ Local scanner fixtures use the observed 1.2.x transcript fields (`step_index`, `source`, `type`, `status`, `created_at`, `content`) and history fields (`display`, integer-millisecond `timestamp`, `workspace`). The installed CLI's transcripts have no model IDs or token usage fields. Tests must not assume they do or estimate tokens from content. SQLite fixtures reproduce only the summary columns the collector reads; conversation databases contain protobuf blobs and are not decoded. Prompt matching also covers `` wrappers observed in local transcripts. Matched history/transcript events in the inspected logs were less than one second apart; prompt text is only compared in memory and is never written to the usage cache. + +After the closing request tag, transcripts can append both an `` block and a separate `` block. Both are outside the user prompt and must be excluded from the comparison. From f008cf624907c1e6d5c3b7be9e435c866367a62b Mon Sep 17 00:00:00 2001 From: steph4n <284007938+steph4n-gh@users.noreply.github.com> Date: Thu, 1 Oct 2026 23:29:17 -0400 Subject: [PATCH 7/7] Publish valid usage despite overlapping collector failures --- bin/omarchy-agent-usage-update | 41 +++++++++++++++------ shell/plugins/agents/README.md | 9 +++-- test/shell.d/agent-usage-update-test.sh | 48 ++++++++++++++++++++----- 3 files changed, 77 insertions(+), 21 deletions(-) diff --git a/bin/omarchy-agent-usage-update b/bin/omarchy-agent-usage-update index e5dbbf39661..4bd1ed3a6ea 100755 --- a/bin/omarchy-agent-usage-update +++ b/bin/omarchy-agent-usage-update @@ -39,31 +39,52 @@ wanted() { return 1 } +published_sequence() { + local value + value=$(jq -r '._collectionSequence // 0' "$1" 2>/dev/null) || value=0 + if [[ $value =~ ^[0-9]{1,15}$ ]]; then + printf '%s\n' "$((10#$value))" + else + printf '0\n' + fi +} + collect() ( local collector="$1" agent="$2" - local record tmp pending="$USAGE_DIR/.$agent.pending" lock="$USAGE_DIR/.$agent.lock" + local record tmp sequence_tmp="" generation=0 published + local sequence="$USAGE_DIR/.$agent.sequence" lock="$USAGE_DIR/.$agent.lock" tmp=$(mktemp "$USAGE_DIR/.$agent.XXXXXX") || return 1 - trap 'rm -f "$tmp"' EXIT + trap 'rm -f "$tmp" "$sequence_tmp"' EXIT + sequence_tmp=$(mktemp "$USAGE_DIR/.$agent.sequence.XXXXXX") || return 1 # Overlapping refreshes may finish in either order, including across - # account changes. Reserve publication before probing; only the latest - # started check may replace the record. Hold the lock only for these short - # file operations, never while waiting for the collector's network calls. + # account changes. Assign request order before probing, then publish only + # if no later request has already published a valid record. A collector + # that crashes cannot prevent another valid result from landing. Hold the + # lock only for file operations, never through the network calls. { flock 9 || return 1 - printf '%s\n' "$tmp" >"$pending" + [[ -f $sequence ]] && read -r generation <"$sequence" + [[ $generation =~ ^[0-9]{1,15}$ ]] || generation=0 + generation=$((10#$generation)) + published=$(published_sequence "$USAGE_DIR/$agent.json") + (( published > generation )) && generation=$published + (( ++generation )) + printf '%s\n' "$generation" >"$sequence_tmp" && mv "$sequence_tmp" "$sequence" } 9>"$lock" || return 1 - if ! record=$("$collector" "${flags[@]}") || [[ -z $record ]] || ! jq -e . >/dev/null 2>&1 <<<"$record"; then + if ! record=$("$collector" "${flags[@]}") || [[ -z $record ]] || ! jq -es 'length == 1 and (.[0] | type == "object")' >/dev/null 2>&1 <<<"$record"; then echo "omarchy-agent-usage-update: $agent collector failed" >&2 return 1 fi - printf '%s\n' "$record" >"$tmp" || return 1 + # Stamp the sequence into the same file as the data so the final rename + # publishes both atomically, even if the updater exits just after it. + jq --argjson sequence "$generation" '._collectionSequence = $sequence' <<<"$record" >"$tmp" || return 1 { flock 9 || return 1 - if [[ -f $pending && $(<"$pending") == "$tmp" ]]; then + published=$(published_sequence "$USAGE_DIR/$agent.json") + if (( generation > published )); then mv "$tmp" "$USAGE_DIR/$agent.json" || return 1 - rm -f "$pending" fi } 9>"$lock" ) diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 76c23db1af4..e2ef7efad38 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -61,9 +61,12 @@ written by `omarchy-agent-usage-update`. That command runs one `omarchy-agent-usage-` collector per agent; the widget invokes it on its refresh timer and whenever you ask for a refresh, and picks up any record that lands in the directory regardless of who wrote it. Overlapping -updates reserve publication before collecting: only the latest started check -for each agent can replace its record, so a slow older result cannot overwrite -a newer account’s limits. The publication lock is released during collection. +updates reserve publication before collecting: a sequence assigned before collection +prevents a slow older result from replacing a newer published record. A crashed +collector or malformed response does not cancel a valid in-flight result; a +newer valid error record, such as a sign-out, does. The sequence and record +are published together atomically. The publication lock is released during +collection. Adding an agent therefore never touches this plugin: ship a collector that prints the record contract (see the `claude` and `codex` collectors in diff --git a/test/shell.d/agent-usage-update-test.sh b/test/shell.d/agent-usage-update-test.sh index 65ed21ffd05..db9f598c8f9 100755 --- a/test/shell.d/agent-usage-update-test.sh +++ b/test/shell.d/agent-usage-update-test.sh @@ -84,13 +84,23 @@ if [[ -n $WAIT_FOR ]]; then while [[ ! -e $WAIT_FOR ]]; do sleep .02; done fi printf '%s\\n' "$RECORD" +exit "${COLLECT_EXIT:-0}" ''') collector.chmod(0o755) env = {**os.environ, 'HOME': str(home), 'OMARCHY_PATH': str(fake), 'XDG_STATE_HOME': ''} usage_dir = home / '.local/state/omarchy/agents/usage' -# An old successful answer and an old error must both lose to a newer check. -for index, old in enumerate([{'id': 'racy', 'limits': [{'percent': .1}]}, - {'id': 'racy', 'limits': [], 'usageStatusText': 'offline'}]): +def published_record(): + record = json.loads((usage_dir / 'racy.json').read_text()) + sequence = record.pop('_collectionSequence') + assert isinstance(sequence, int) and sequence > 0 + return record + +# A published service error is authoritative too: an older success must not +# undo a newer sign-out. Only a hard failure to emit a record permits fallback. +good_old = {'id': 'racy', 'limits': [{'percent': .1}]} +good_new = {'id': 'racy', 'limits': [{'percent': .8}], 'usageStatusText': ''} +error = {'id': 'racy', 'limits': [], 'usageStatusText': 'sign-in required'} +for index, (old, new) in enumerate([(good_old, good_new), (error, good_new), (good_old, error)]): started = home / f'started-{index}' release = home / f'release-{index}' older = subprocess.Popen([str(updater), 'racy'], env={**env, 'STARTED': str(started), 'WAIT_FOR': str(release), 'RECORD': json.dumps(old)}, @@ -100,15 +110,37 @@ for index, old in enumerate([{'id': 'racy', 'limits': [{'percent': .1}]}, while not started.exists() and older.poll() is None and time.monotonic() < deadline: time.sleep(.02) assert started.exists(), 'older collector started' - new = {'id': 'racy', 'limits': [{'percent': .8}], 'usageStatusText': ''} subprocess.run([str(updater), 'racy'], env={**env, 'STARTED': str(home / 'newer-started'), 'WAIT_FOR': '', 'RECORD': json.dumps(new)}, capture_output=True, text=True, check=True, timeout=5) - assert json.loads((usage_dir / 'racy.json').read_text()) == new + assert published_record() == new finally: release.touch() stdout, stderr = older.communicate(timeout=5) assert older.returncode == 0, stderr - assert json.loads((usage_dir / 'racy.json').read_text()) == new, 'older completion replaced newer data' - assert sorted(p.name for p in usage_dir.glob('.racy.*')) == ['.racy.lock'], 'temporary records are cleaned up' -print('ok - overlapping updates publish only the latest started check without blocking its probe') + assert published_record() == new, 'older completion replaced newer data' + assert sorted(p.name for p in usage_dir.glob('.racy.*')) == ['.racy.lock', '.racy.sequence'], 'temporary records are cleaned up' +print('ok - older updates cannot overwrite a newer valid record, including sign-out errors') + +for index, failed in enumerate([{'RECORD': 'not json', 'COLLECT_EXIT': '0'}, + {'RECORD': '{}\n{}', 'COLLECT_EXIT': '0'}, + {'RECORD': '{}', 'COLLECT_EXIT': '1'}]): + started = home / f'failure-started-{index}' + release = home / f'failure-release-{index}' + good = {'id': 'racy', 'limits': [{'percent': .4}]} + older = subprocess.Popen([str(updater), 'racy'], env={**env, 'STARTED': str(started), 'WAIT_FOR': str(release), 'RECORD': json.dumps(good)}, + stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + try: + deadline = time.monotonic() + 5 + while not started.exists() and older.poll() is None and time.monotonic() < deadline: + time.sleep(.02) + assert started.exists() + result = subprocess.run([str(updater), 'racy'], env={**env, 'STARTED': str(home / 'failure-newer'), 'WAIT_FOR': '', **failed}, + capture_output=True, text=True, timeout=5) + assert result.returncode != 0, 'hard collector failures are reported' + finally: + release.touch() + stdout, stderr = older.communicate(timeout=5) + assert older.returncode == 0, stderr + assert published_record() == good, 'a failed newer request suppressed a valid in-flight result' +print('ok - failed collectors do not cancel valid in-flight publications') PY