diff --git a/bin/omarchy-agent-usage-agy b/bin/omarchy-agent-usage-agy new file mode 100755 index 00000000000..a5e6b15e351 --- /dev/null +++ b/bin/omarchy-agent-usage-agy @@ -0,0 +1,780 @@ +#!/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. + +The id is agy, the name omarchy-default-agent and omarchy-agent know the CLI +by, so the record and the default agent agree on which agent this is. + +Local stats come from Antigravity's history log, conversation summaries, and +conversation transcripts under ~/.gemini/antigravity-cli/, and from the +Antigravity answers pi, omp, OpenClaw, and opencode's Antigravity plugin +recorded. The plan and limits come from the Google Cloud Code API +(v1internal:loadCodeAssist and v1internal:retrieveUserQuotaSummary), asked +with the sign-in agy keeps in the Secret Service keyring, or with the +Antigravity sign-in pi or opencode keeps; omp's come from `omp usage`. The +agents panel only ever reads the JSON this prints. +""" + +import argparse +import base64 +import hashlib +import json +import os +import sqlite3 +import shutil +import subprocess +import tempfile +import time +import urllib.error +import urllib.request +from datetime import datetime, timedelta, timezone +from pathlib import Path + +AGENT_ID = "agy" +AGENT_NAME = "Antigravity" +AUTH_HELP = "Run `agy` to sign in, or sign in to Antigravity in omp, pi, or opencode." +# The panel reads "Waiting for auth" and "Sign-in expired" as "offer to sign +# in here", and its link runs omarchy-agent-account-add, which knows Claude, +# Codex, and Grok only. Antigravity is signed in outside the panel, so its +# record names the state in words and leaves the instruction to authHelpText, +# which the panel prints where the link would have been. +NO_SIGN_IN = "No Antigravity sign-in" +LAPSED = "Antigravity sign-in expired" +# How each source is named when its sign-in needs refreshing. +SOURCE_NAMES = {"agy": "`agy`", "omp": "omp", "pi": "pi", "opencode": "opencode"} +# A panel opened and shut repeatedly reuses an answer for about a minute. +PROBE_REUSE_SECONDS = 60 +# A limits-only refresh needs fresh limits, not a fresh history scan. +LIMITS_ONLY_SCAN_REUSE_SECONDS = 900 +# Where the plan and quota are asked for. The bearer token can now come from +# another harness' store, so a destination named in the environment is taken +# only over https, and anything else falls back to Google's own endpoint. +DEFAULT_CLOUD_CODE_URL = "https://daily-cloudcode-pa.googleapis.com" + + +def app_root(): + return Path(os.environ.get("AGY_DIR") or (Path.home() / ".gemini" / "antigravity-cli")) + + +def cache_root(): + return Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage" + + +def read_json(path): + try: + return json.loads(Path(path).read_text(encoding="utf-8")) + except Exception: + return None + + +# A temporary file of its own per write, so overlapping runs can't rename +# each other's away. +def write_json(path, payload): + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp = tempfile.mkstemp(prefix=path.name + ".", dir=str(path.parent)) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(payload, handle) + os.replace(tmp, path) + except Exception: + Path(tmp).unlink(missing_ok=True) + raise + + +def parse_time(value): + if value is None or value == "" or isinstance(value, bool): + return None + try: + if isinstance(value, (int, float)) or str(value).isdigit(): + stamp = float(value) + parsed = datetime.fromtimestamp(stamp / 1000 if stamp > 1e12 else stamp, timezone.utc) + else: + parsed = datetime.fromisoformat(str(value).replace("Z", "+00:00")) + if not parsed.tzinfo: + parsed = parsed.replace(tzinfo=timezone.utc) + except (ValueError, OverflowError, OSError): + return None + # Go's zero time marks an uninitialized summary, not dated activity. + return None if parsed == datetime.min.replace(tzinfo=timezone.utc) else parsed + + +def local_day(value): + parsed = parse_time(value) + if parsed is None: + return "" + try: + return parsed.astimezone().date().isoformat() + except (ValueError, OverflowError, OSError): + # Even a parsed date can cross years 1/9999 during local conversion. + return "" + + +# The sign-in agy keeps in the keyring: its access token, when that lapses, +# and which kind of account it is. Without a Secret Service, agy keeps the +# same token in a file under its home instead. +def login(): + try: + raw = subprocess.check_output( + ["secret-tool", "lookup", "service", "gemini", "username", "antigravity"], + text=True, stderr=subprocess.DEVNULL, timeout=3, + ) + data = json.loads(raw) if raw.strip() else None + except Exception: + data = None + if not isinstance(data, dict): + data = read_json(app_root() / "antigravity-oauth-token") + if not isinstance(data, dict): + return None + + token = data.get("token") if isinstance(data.get("token"), dict) else {} + expires = parse_time(token.get("expiry") or data.get("expiry")) + return { + "source": "agy", + "accessToken": str(token.get("access_token") or data.get("access_token") or ""), + "expiresAt": expires.timestamp() if expires else 0, + "authMethod": str(data.get("auth_method") or "consumer"), + "email": email_of(data.get("id_token")), + "refresh": str(token.get("refresh_token") or data.get("refresh_token") or ""), + } + + +# omp, pi, and opencode's Antigravity plugin sign in with the same Google +# client agy uses, so the access tokens they keep reach the same plan and +# quota. Each refreshes its token while it runs; these are only ever read. +def harness_login(source, data): + if not isinstance(data, dict) or data.get("type", "oauth") != "oauth" or not data.get("access"): + return None + expires = parse_time(data.get("expires")) + return { + "source": source, + "accessToken": str(data["access"]), + "expiresAt": expires.timestamp() if expires else 0, + "authMethod": "consumer", + "email": str(data.get("email") or ""), + "refresh": str(data.get("refresh") or ""), + } + + +# omp and pi both honor PI_CODING_AGENT_DIR for their agent directory. +def harness_dir(name): + return Path(os.environ.get("PI_CODING_AGENT_DIR") or (Path.home() / name / "agent")) + + +def omp_logins(): + database = harness_dir(".omp") / "agent.db" + if not database.is_file(): + return [] + try: + conn = sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + try: + rows = conn.execute( + "SELECT data FROM auth_credentials WHERE provider = 'google-antigravity' " + "AND credential_type = 'oauth' AND disabled_cause IS NULL" + ).fetchall() + finally: + conn.close() + except sqlite3.Error: + return [] + logins = [] + for (data,) in rows: + try: + entry = harness_login("omp", json.loads(data)) + except (TypeError, ValueError): + continue + if entry: + logins.append(entry) + return logins + + +def pi_logins(): + auth = read_json(harness_dir(".pi") / "auth.json") + entry = harness_login("pi", auth.get("google-antigravity")) if isinstance(auth, dict) else None + return [entry] if entry else [] + + +# opencode files the Antigravity plugin's sign-in under its "google" +# provider, which other Google sign-ins share, so it only counts alongside +# the plugin's own accounts file. +def opencode_logins(): + config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) + if not (config_home / "opencode" / "antigravity-accounts.json").is_file(): + return [] + data_home = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) + auth = read_json(data_home / "opencode" / "auth.json") + entry = harness_login("opencode", auth.get("google")) if isinstance(auth, dict) else None + return [entry] if entry else [] + + +def logins(): + entry = login() + return ([entry] if entry else []) + omp_logins() + pi_logins() + opencode_logins() + + +def expired_help(entry): + name = SOURCE_NAMES[entry["source"]] + return f"The Antigravity sign-in {name} keeps expired. Start {name} to refresh it." + + +# Which Google account a sign-in belongs to, so limits kept from an earlier +# check can be tied to it: the email when the harness records one, otherwise +# a digest of that harness' own token. The refresh token outlives the access +# tokens that rotate; the access token, which changes with the account or a +# refresh, drops the cache rather than let it pass for another's. +def identity_of(entry): + if entry.get("email"): + return str(entry["email"]) + secret = str(entry.get("refresh") or entry.get("accessToken") or "") + return f"{entry['source']}:{hashlib.sha256(secret.encode()).hexdigest()[:16]}" if secret else "" + + +def email_of(id_token): + try: + claims = str(id_token).split(".")[1] + return str(json.loads(base64.urlsafe_b64decode(claims + "=" * (-len(claims) % 4))).get("email") or "") + except Exception: + return "" + + +def normalize_tier_label(name, tier_id=""): + for word in ("ultra", "pro", "free"): + if word in tier_id.lower() or word in name.lower(): + return word.capitalize() + if name.lower().startswith("google ai "): + return name[10:].strip() + return name or tier_id.replace("-tier", "").capitalize() + + +# A paid tier names the plan; without one, the current tier does. +def extract_tier(data): + for key in ("paidTier", "currentTier"): + tier = data.get(key) + if isinstance(tier, dict): + name = str(tier.get("name") or "").strip() + tier_id = str(tier.get("id") or "").strip() + if name or tier_id: + return normalize_tier_label(name, tier_id) + return "" + + +def used(bucket): + return round(max(0.0, min(1.0, 1.0 - float(bucket.get("remainingFraction", 1.0)))), 4) + + +def reset_of(bucket): + parsed = parse_time(bucket.get("resetTime")) + return parsed.isoformat() if parsed else "" + + +# The Gemini group's 5-hour and weekly buckets are the plan's own limits; the +# other groups (Claude and GPT models) are scoped allowances, shown once +# they have been touched. +def extract_limits(data): + groups = [g for g in data.get("groups") or [] if isinstance(g, dict)] + if not groups: + return [] + gemini = next((g for g in groups if "gemini" in str(g.get("displayName") or "").lower()), groups[0]) + + limits = [] + session = weekly = None + for bucket in gemini.get("buckets") or []: + if not isinstance(bucket, dict): + continue + window = str(bucket.get("window") or "").lower() + bucket_id = str(bucket.get("bucketId") or "").lower() + name = str(bucket.get("displayName") or "").lower() + if window == "5h" or "5h" in bucket_id or "five" in name: + session = bucket + elif window == "weekly" or "weekly" in bucket_id or "week" in name: + weekly = bucket + if session: + limits.append({"label": "Session (5-hour)", "title": "Session", "percent": used(session), "resetsAt": reset_of(session)}) + if weekly: + limits.append({"label": "Weekly (7-day)", "title": "Weekly", "percent": used(weekly), "resetsAt": reset_of(weekly)}) + + for group in groups: + if group is gemini: + continue + group_name = str(group.get("displayName") or "").strip() + short = "Claude/GPT" if "claude" in group_name.lower() else (group_name or "3P") + for bucket in group.get("buckets") or []: + if isinstance(bucket, dict) and used(bucket) > 0: + title = f"{short} {'Weekly' if 'week' in str(bucket.get('window') or '').lower() else 'Session'}" + limits.append({"label": title, "title": title, "percent": used(bucket), "resetsAt": reset_of(bucket)}) + return limits + + +def post(endpoint, token): + base = (os.environ.get("CLOUD_CODE_URL") or DEFAULT_CLOUD_CODE_URL).rstrip("/") + if not base.startswith("https://"): + base = DEFAULT_CLOUD_CODE_URL + request = urllib.request.Request( + f"{base}/v1internal:{endpoint}", data=b"{}", + headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json", "User-Agent": "antigravity"}, + ) + with urllib.request.urlopen(request, timeout=8) as response: + return json.loads(response.read().decode("utf-8", errors="replace")) + + +# A window kept from an earlier check is dropped once it has reset, since +# what it said no longer holds. A damaged cache entry is dropped with it +# rather than taken down the collector. +def open_windows(limits): + now = datetime.now(timezone.utc) + return [ + limit for limit in (limits if isinstance(limits, list) else []) + if isinstance(limit, dict) and (not parse_time(limit.get("resetsAt")) or parse_time(limit.get("resetsAt")) > now) + ] + + +# omp refreshes its own Antigravity sign-in when asked for usage, so its +# limits come from `omp usage` even after the token it keeps has lapsed. The +# report is folded back into the shape Google's quota summary has. +def omp_usage(omp): + try: + output = subprocess.run( + [omp, "usage", "--json", "--provider", "google-antigravity"], + capture_output=True, text=True, timeout=30, check=True, + ).stdout + reports = [r for r in json.loads(output).get("reports") or [] if isinstance(r, dict) and r.get("provider") == "google-antigravity"] + except (OSError, ValueError, AttributeError, subprocess.SubprocessError): + return None + if not reports: + return None + groups, seen = {}, set() + for limit in reports[0].get("limits") or []: + if not isinstance(limit, dict): + continue + window = limit.get("window") if isinstance(limit.get("window"), dict) else {} + amount = limit.get("amount") if isinstance(limit.get("amount"), dict) else {} + label = str(limit.get("label") or "") + # A shared allowance is listed once per model family that draws on it. + if (label, window.get("id")) in seen: + continue + seen.add((label, window.get("id"))) + try: + remaining = 1.0 - float(amount.get("usedFraction") or 0) + except (TypeError, ValueError): + continue + groups.setdefault(label, []).append({"window": str(window.get("id") or ""), "remainingFraction": remaining, "resetTime": window.get("resetsAt")}) + metadata = reports[0].get("metadata") if isinstance(reports[0].get("metadata"), dict) else {} + quota = {"groups": [{"displayName": name, "buckets": buckets} for name, buckets in groups.items()]} + return quota, str(metadata.get("email") or "") + + +class Lapsed(Exception): + pass + + +def token_fresh(entry, now): + return bool(entry["accessToken"]) and (not entry["expiresAt"] or entry["expiresAt"] > now) + + +# The plan and limits one sign-in answers with. omp is asked through its CLI, +# and its own token, just refreshed, then only names the plan. +def probe(entry, omp, now): + if entry["source"] == "omp" and omp: + answer = omp_usage(omp) + if answer: + quota, email = answer + refreshed = next((e for e in omp_logins() if not email or e["email"] == email), entry) + try: + tier = extract_tier(post("loadCodeAssist", refreshed["accessToken"])) if token_fresh(refreshed, time.time()) else "" + except Exception: + tier = "" + return tier, extract_limits(quota) or None, identity_of(refreshed) + if not token_fresh(entry, now): + raise Lapsed() + tier = extract_tier(post("loadCodeAssist", entry["accessToken"])) + return tier, extract_limits(post("retrieveUserQuotaSummary", entry["accessToken"])) or None, identity_of(entry) + + +def collect_limits(entries, force): + result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": AUTH_HELP, "live": False, "fetchedAtMs": 0} + if not entries: + result["usageStatusText"] = NO_SIGN_IN + return result + + now = time.time() + omp = shutil.which("omp") + # Every harness refreshes its own token while it runs, so the one lapsing + # last was refreshed most recently. A token of unknown expiry is tried + # last, and only Google can say whether it still holds. + signed_in = sorted((e for e in entries if e["accessToken"] or (e["source"] == "omp" and omp)), key=lambda e: e["expiresAt"], reverse=True) + usable = [e for e in signed_in if token_fresh(e, now) or (e["source"] == "omp" and omp)] + + cache = cache_root() / "agy-limits.json" + cached = read_json(cache) or {} + fetched_at = int(cached.get("fetchedAtMs") or 0) + kept = open_windows(cached.get("limits")) + tier = str(cached.get("tierLabel") or "") + # Limits kept from an earlier check belong to the account that earned them, + # so they stand in only while a sign-in for that account is still here: + # another account's never stand in for it, and without an identity to match + # they are dropped rather than shown. + signed_in_as = {identity_of(e) for e in entries} + if str(cached.get("identity") or "") not in signed_in_as - {""}: + kept, tier, fetched_at = [], "", 0 + lead = (usable or signed_in or [{}])[0] + result.update(fetchedAtMs=fetched_at, tierLabel=tier) + + if not signed_in: + result.update(limits=kept, usageStatusText=NO_SIGN_IN) + return result + + # A token left to lapse can't be used until its harness runs again. + if not usable: + result.update(limits=kept, usageStatusText=LAPSED, authHelpText=expired_help(lead)) + return result + + # A panel opened and shut repeatedly reuses the answer it just got, but + # only for the account that answer describes. + if kept and not force and identity_of(lead) == str(cached.get("identity") or "") and now - fetched_at / 1000 < PROBE_REUSE_SECONDS: + result.update(limits=kept, live=True) + return result + + # A sign-in Google refuses gives way to the next one. + tier_from_probe, limits, identity, rejected = "", None, "", None + for entry in usable: + try: + tier_from_probe, limits, identity = probe(entry, omp, now) + except Lapsed: + rejected = rejected or entry + continue + except urllib.error.HTTPError as error: + if error.code in (401, 403): + rejected = rejected or entry + continue + except Exception: + pass + rejected = None + break + if rejected: + result.update(usageStatusText=LAPSED, authHelpText=expired_help(rejected)) + + # A plan label, and limits, kept from an earlier check describe the account + # that earned them: they stand in for that account, or for a run nothing + # answered at all, but never beside another account's answer. + cached_identity = str(cached.get("identity") or "") + if identity and identity != cached_identity: + result["tierLabel"] = "" + kept = [] + if tier_from_probe: + result["tierLabel"] = tier_from_probe + + if limits: + result.update(limits=limits, live=True, fetchedAtMs=round(time.time() * 1000)) + write_json(cache, {"fetchedAtMs": result["fetchedAtMs"], "limits": limits, "tierLabel": result["tierLabel"], "identity": identity}) + else: + result["limits"] = kept + if not result["usageStatusText"] and not kept: + result["usageStatusText"] = "Antigravity limits unavailable" + return result + + +def model_bucket(models, model): + return models.setdefault(model, {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}) + + +def history_entries(root): + try: + lines = (root / "history.jsonl").read_text(encoding="utf-8", errors="ignore").splitlines() + except OSError: + return + for line in lines: + try: + entry = json.loads(line) + except ValueError: + continue + if isinstance(entry, dict): + yield entry + + +def number(value): + try: + return max(0, int(value or 0)) + except (TypeError, ValueError): + return 0 + + +# `omp --profile=` (and pi's equivalent) moves the whole agent tree +# under /profiles//, so each profile's sessions are scanned too. +# OpenClaw keeps the same transcript format, one directory per agent. +def session_roots(): + roots = [] + if os.environ.get("PI_CODING_AGENT_DIR"): + roots.append(Path(os.environ["PI_CODING_AGENT_DIR"]) / "sessions") + for base in (Path.home() / ".pi", Path.home() / ".omp"): + roots.append(base / "agent" / "sessions") + try: + roots.extend(sorted(child / "agent" / "sessions" for child in (base / "profiles").iterdir() if child.is_dir())) + except OSError: + pass + openclaw = Path(os.environ.get("OPENCLAW_STATE_DIR") or (Path.home() / ".openclaw")) + try: + roots.extend(sorted(child / "sessions" for child in (openclaw / "agents").iterdir() if child.is_dir())) + except OSError: + pass + return list(dict.fromkeys(roots)) + + +# Each answer pi, omp, or OpenClaw got through its google-antigravity +# provider: the session it belongs to, its day, its model, and its input, +# output, cache read, and cache write tokens. +def session_messages(): + rg = shutil.which("rg") + if not rg: + return + seen = set() + for root in session_roots(): + if not root.is_dir(): + continue + try: + proc = subprocess.Popen( + # --no-ignore: a home that is itself a git checkout would otherwise + # hide every session file behind its ignore rules. + [rg, "--json", "--no-ignore", "-e", r'"provider"\s*:\s*"google-antigravity"', str(root)], + stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, errors="replace", + ) + except OSError: + return + for raw in proc.stdout: + try: + event = json.loads(raw) + if event.get("type") != "match": + continue + path = event["data"]["path"]["text"] + entry = json.loads(event["data"]["lines"]["text"]) + except (ValueError, KeyError, TypeError): + continue + if not isinstance(entry, dict) or entry.get("type") != "message": + continue + message = entry.get("message") if isinstance(entry.get("message"), dict) else {} + if message.get("role") != "assistant" or message.get("provider") != "google-antigravity": + continue + # A fork copies its parent's entries, and they keep their id, so each + # answer counts once. + if entry.get("id"): + key = (entry["id"], entry.get("timestamp")) if entry.get("timestamp") else (path, entry["id"]) + if key in seen: + continue + seen.add(key) + usage = message.get("usage") if isinstance(message.get("usage"), dict) else {} + tokens = (number(usage.get("input")), number(usage.get("output")), number(usage.get("cacheRead")), number(usage.get("cacheWrite"))) + if not any(tokens): + tokens = (number(usage.get("totalTokens")), 0, 0, 0) + day = local_day(entry.get("timestamp") or message.get("timestamp")) + if any(tokens) and day: + yield path, day, str(message.get("model") or "unknown"), tokens + try: + proc.wait(timeout=1) + except Exception: + proc.kill() + + +# opencode's Antigravity plugin runs under the "google" provider, with model +# ids it prefixes with antigravity- to keep them apart from Gemini API ones. +def opencode_messages(): + data_home = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) + database = data_home / "opencode" / "opencode.db" + if not database.is_file(): + return + try: + conn = sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + except sqlite3.Error: + return + try: + # The SQL is only acceleration: the Python filter below is the authority + # for which rows are this subscription's, and the guards in front can + # gate out a row it would have accepted. + rows = conn.execute( + "SELECT session_id, data FROM message" + " WHERE data LIKE '%\"role\"%:%\"assistant\"%' AND data LIKE '%antigravity-%'" + " AND CASE WHEN json_valid(data) THEN json_extract(data, '$.role') END = 'assistant'" + ).fetchall() + except sqlite3.Error: + rows = [] + finally: + conn.close() + for session, raw in rows: + try: + entry = json.loads(raw) + provider = str(entry.get("providerID") or "") + model = str(entry.get("modelID") or "") + # antigravity- is what the plugin prefixes its models with; Gemini + # asked for through the same provider id is a different subscription. + if not provider.startswith("google") or not model.startswith("antigravity-"): + continue + tokens = entry.get("tokens") if isinstance(entry.get("tokens"), dict) else {} + cache = tokens.get("cache") if isinstance(tokens.get("cache"), dict) else {} + created = (entry.get("time") or {}).get("created") + except (ValueError, AttributeError): + continue + # opencode keeps thinking tokens out of output; both are generated. + counts = (number(tokens.get("input")), number(tokens.get("output")) + number(tokens.get("reasoning")), number(cache.get("read")), number(cache.get("write"))) + day = local_day(created) + if any(counts) and day: + yield "opencode:" + str(session), day, model, counts + + +# Hermes reaches Antigravity only through plugins that run agy itself, so its +# sign-in and transcripts are agy's own, which local_stats reads already. +def harness_messages(): + yield from session_messages() + yield from opencode_messages() + + +# Every prompt from history.jsonl, every conversation from it and from the +# conversation summaries, and the tokens each model step in a transcript +# reports. Steps that report no usage add no tokens. Antigravity driven from +# another harness adds each answer it recorded as a prompt, the way the Codex +# collector counts them, with every token it reports. +def local_stats(root): + today = datetime.now().astimezone().date() + week = [(today - timedelta(days=offset)).isoformat() for offset in range(6, -1, -1)] + today = today.isoformat() + by_day = {day: 0 for day in week} + prompts = today_prompts = today_tokens = 0 + days, sessions, today_sessions = set(), set(), set() + today_by_model, models = {}, {} + + for entry in history_entries(root): + day = local_day(entry.get("timestamp")) + if not day: + continue + prompts += 1 + days.add(day) + conversation = entry.get("conversationId") or entry.get("conversation_id") + if conversation: + sessions.add(conversation) + if day == today: + today_prompts += 1 + if conversation: + today_sessions.add(conversation) + + summaries = root / "conversation_summaries.db" + if summaries.is_file(): + try: + conn = sqlite3.connect(summaries.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + try: + for conversation, modified in conn.execute("SELECT conversation_id, last_modified_time FROM conversation_summaries"): + if conversation: + sessions.add(conversation) + if local_day(modified) == today: + today_sessions.add(conversation) + finally: + conn.close() + except sqlite3.Error: + pass + + for transcript in (root / "brain").glob("*/.system_generated/logs/transcript.jsonl"): + try: + lines = transcript.read_text(encoding="utf-8", errors="ignore").splitlines() + except OSError: + continue + for line in lines: + try: + step = json.loads(line) + except ValueError: + continue + usage = step.get("usage") if isinstance(step, dict) else None + if not isinstance(usage, dict): + continue + uncached = number(usage.get("prompt_token_count") or usage.get("input_tokens") or 0) + output = number(usage.get("candidates_token_count") or usage.get("output_tokens") or 0) + model = str(step.get("model") or "unknown") + bucket = model_bucket(models, model) + bucket["inputTokens"] += uncached + bucket["outputTokens"] += output + day = local_day(step.get("created_at")) + if day: + days.add(day) + if day in by_day: + by_day[day] += uncached + output + if day == today: + today_tokens += uncached + output + today_by_model[model] = today_by_model.get(model, 0) + uncached + output + + for session, day, model, counts in harness_messages(): + total = sum(counts) + prompts += 1 + days.add(day) + sessions.add(session) + bucket = model_bucket(models, model) + bucket["inputTokens"] += counts[0] + bucket["outputTokens"] += counts[1] + bucket["cacheReadInputTokens"] += counts[2] + bucket["cacheCreationInputTokens"] += counts[3] + if day in by_day: + by_day[day] += total + if day == today: + today_prompts += 1 + today_sessions.add(session) + today_tokens += total + today_by_model[model] = today_by_model.get(model, 0) + total + + return { + "hasLocalStats": prompts > 0 or bool(sessions) or bool(models), + "todayPrompts": today_prompts, + "todaySessions": len(today_sessions), + "todayTotalTokens": today_tokens, + "todayTokensByModel": today_by_model, + "recentDays": [{"date": day, "messageCount": by_day[day]} for day in week], + "totalPrompts": prompts, + "totalSessions": len(sessions), + "activeDays": len(days), + "activeDates": sorted(days), + "modelUsage": models, + } + + +# A scan from before midnight would count yesterday's prompts as today's. +def cached_local_stats(root, max_age): + cache = cache_root() / "agy-stats.json" + cached = read_json(cache) or {} + today = datetime.now().astimezone().date().isoformat() + if cached.get("root") == str(root) and cached.get("day") == today and time.time() - float(cached.get("at") or 0) < max_age: + return cached["stats"] + stats = local_stats(root) + write_json(cache, {"root": str(root), "day": today, "at": time.time(), "stats": stats}) + return stats + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + args = parser.parse_args() + + # Nobody signed in to Antigravity here: no limits and no plan, so without + # local history the panel skips the record. + entries = logins() + limits = collect_limits(entries, args.force) + agy = next((e for e in entries if e["source"] == "agy"), None) + tier = limits["tierLabel"] + if not tier and agy and agy["authMethod"] == "enterprise": + tier = "Enterprise" + elif not tier and agy and agy["authMethod"] in ("cloud", "vertex"): + tier = "Cloud" + + record = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": bool(entries), + "tierLabel": tier, + "usageStatusText": limits["usageStatusText"], + "authHelpText": limits["authHelpText"], + "limits": limits["limits"], + "limitsStale": not limits["live"], + "limitsFetchedAt": limits["fetchedAtMs"], + } + record.update(cached_local_stats(app_root(), LIMITS_ONLY_SCAN_REUSE_SECONDS if args.limits_only and not args.force else 0)) + print(json.dumps(record, separators=(",", ":"))) + + +if __name__ == "__main__": + main() diff --git a/manual/17-ai.md b/manual/17-ai.md index 0eba2b148e5..60860112f2a 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -39,7 +39,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, a meter for each of its 5-hour session and weekly limits with the time until it resets (hover for the exact percentage), 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, Fireworks, and Kimi are covered out of the box. Kimi currently reports account quota limits without local token or prompt statistics. +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, a meter for each of its 5-hour session and weekly limits with the time until it resets (hover for the exact percentage), 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, Fireworks, and Kimi are covered out of the box, as is Antigravity with its 5-hour and weekly limits, whether you signed in to it with `agy` or in omp, pi, or opencode, and its tokens count wherever you used it: `agy`, pi, omp, OpenClaw, or opencode. Kimi currently reports account quota limits without local token or prompt statistics. 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 d05eb8aec1f..fe8e7bf602d 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -72,6 +72,7 @@ Adding an agent needs no QML changes: ship a collector that prints the record co | `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 | | `fireworks` | Estimated prepaid balance: configured funding minus rated account costs | Fireworks billing API, grouped by day and model for the last 30 days | | `kimi` | Kimi Code quotas (5-hour session, weekly, and monthly allowances) | none; quota only | +| `agy` | Google's Cloud Code API (`loadCodeAssist` for the plan, `retrieveUserQuotaSummary` for the 5-hour and weekly windows), or `omp usage` for omp's sign-in | `history.jsonl`, `conversation_summaries.db`, and each conversation's `transcript.jsonl` under `~/.gemini/antigravity-cli/` (`AGY_DIR`), plus pi, omp, and OpenClaw sessions on the `google-antigravity` provider and opencode sessions on its Antigravity plugin's `antigravity-` models | When `~/.local/state/omarchy/agents/accounts/.json` registers more than one account, the `claude`, `codex`, and `grok` records @@ -91,7 +92,24 @@ Codex CLI will front any OpenAI-compatible backend — `--oss`, or a custom `mod 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`, Grok via `GROK_HOME`. Grok's +`CLAUDE_CONFIG_DIR`, Codex via `CODEX_HOME`, Grok via `GROK_HOME`, +Antigravity via `AGY_DIR`. Antigravity's plan and limits are +asked with the sign-in `agy` keeps in the Secret Service keyring (service +`gemini`, username `antigravity`, or `antigravity-oauth-token` under +`AGY_DIR` without a keyring). Without `agy`, the Antigravity sign-in of +another harness stands in: omp's `agent.db`, pi's `auth.json` (both under +`PI_CODING_AGENT_DIR` when set), or opencode's `google` sign-in when the +Antigravity plugin's `antigravity-accounts.json` is there. Of these, the +token refreshed most recently is asked first, and one Google refuses gives +way to the next. omp is asked through `omp usage`, which refreshes its own +sign-in, and its token then only names the plan. Every other harness +refreshes its token while it runs; one left to lapse shows the last limits +until that harness runs again. Because the panel's _Sign-in required_ link +runs a sign-in flow that knows Claude, Codex, and Grok only, the Antigravity +record names that state in words and puts the instruction in `authHelpText` +instead of asking for a link. Hermes reaches Antigravity only through +plugins that run `agy`, so its usage is agy's own. The record's id is +`agy`, the name the default agent knows it by. Grok's plan comes from the settings it caches in its home, and its limit from the credits endpoint its own `/usage` view reads, asked with each account's sign-in; a sign-in left to lapse shows the last credits until Grok runs diff --git a/shell/plugins/agents/assets/agy-light.svg b/shell/plugins/agents/assets/agy-light.svg new file mode 100644 index 00000000000..67ced694468 --- /dev/null +++ b/shell/plugins/agents/assets/agy-light.svg @@ -0,0 +1 @@ +Antigravity diff --git a/shell/plugins/agents/assets/agy.svg b/shell/plugins/agents/assets/agy.svg new file mode 100644 index 00000000000..3743be8edf4 --- /dev/null +++ b/shell/plugins/agents/assets/agy.svg @@ -0,0 +1 @@ +Antigravity diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index 21dc959dcf2..5e6bccbf3c1 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, Grok, Fireworks, and Kimi usage, limits, and pace in a native Omarchy bar panel.", + "description": "Claude Code, Codex, Grok, Fireworks, Kimi, and Antigravity usage, limits, and pace in a native Omarchy bar panel.", "kinds": ["bar-widget"], "activation": "on-demand", "entryPoints": { @@ -23,7 +23,8 @@ "codex": { "enabled": true }, "grok": { "enabled": true }, "fireworks": { "enabled": true }, - "kimi": { "enabled": true } + "kimi": { "enabled": true }, + "agy": { "enabled": true } }, "refreshIntervalSec": 900, "syncMode": "Off", diff --git a/test/shell.d/agent-usage-agy-harness-test.sh b/test/shell.d/agent-usage-agy-harness-test.sh new file mode 100755 index 00000000000..f4496fba8d6 --- /dev/null +++ b/test/shell.d/agent-usage-agy-harness-test.sh @@ -0,0 +1,266 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 +require_command sqlite3 + +COLLECTOR="$ROOT/bin/omarchy-agent-usage-agy" + +# Antigravity signed in through another harness, with agy itself absent: the +# record still carries the plan and limits, asked with that harness' token. +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT +mkdir -p "$test_tmp/bin" +cat >"$test_tmp/bin/secret-tool" <<'EOF' +#!/bin/bash +exit 1 +EOF +# omp answers `omp usage` with the report in omp-report.json, when there is +# one, and fails otherwise. +cat >"$test_tmp/bin/omp" <<'EOF' +#!/bin/bash +[[ "$*" == "usage --json --provider google-antigravity" ]] || exit 2 +cat "$HOME/omp-report.json" 2>/dev/null +EOF +chmod +x "$test_tmp/bin/secret-tool" "$test_tmp/bin/omp" + +in_ms() { + python3 -c "import sys, time; print(round((time.time() + float(sys.argv[1])) * 1000))" "$1" +} + +omp_login() { + mkdir -p "$test_tmp/.omp/agent" + rm -f "$test_tmp/.omp/agent/agent.db" + sqlite3 "$test_tmp/.omp/agent/agent.db" \ + "CREATE TABLE auth_credentials (id INTEGER PRIMARY KEY, provider TEXT, credential_type TEXT, data TEXT, disabled_cause TEXT); + INSERT INTO auth_credentials (provider, credential_type, data) VALUES + ('google-antigravity', 'oauth', '{\"access\":\"$1\",\"refresh\":\"r\",\"expires\":$2,\"email\":\"$3\"}');" +} + +pi_login() { + mkdir -p "$test_tmp/.pi/agent" + printf '{"google-antigravity":{"type":"oauth","access":"%s","refresh":"r","expires":%s,"email":"%s"}}\n' "$1" "$2" "$3" \ + >"$test_tmp/.pi/agent/auth.json" +} + +# Google answers every token in ACCEPTED and refuses the rest; each request's +# token is logged so a test can tell which sign-in was asked. +collect() { + env -u PI_CODING_AGENT_DIR -u OPENCLAW_STATE_DIR -u XDG_CONFIG_HOME -u XDG_DATA_HOME \ + HOME="$test_tmp" XDG_CACHE_HOME="$test_tmp/cache" PATH="$test_tmp/bin:$PATH" \ + CLOUD_CODE_URL="${2:-}" COLLECTOR="$COLLECTOR" ACCEPTED="$1" ASKED="$test_tmp/asked" \ + HOSTS="$test_tmp/hosts" python3 - <<'PY' +import importlib.machinery, importlib.util, io, json, os, sys, urllib.error + +loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) + +answers = { + "loadCodeAssist": {"paidTier": {"id": "g1-pro-tier", "name": "Google AI Pro"}}, + "retrieveUserQuotaSummary": {"groups": [{"displayName": "Gemini Models", "buckets": [ + {"window": "5h", "remainingFraction": 0.75, "resetTime": "2999-01-01T00:00:00Z"} + ]}]}, +} +# An account on a plan with no quota windows to report. +plans_without_quota = { + "no-quota-token": { + "loadCodeAssist": {"paidTier": {"id": "g1-ultra-tier", "name": "Google AI Ultra"}}, + "retrieveUserQuotaSummary": {"groups": []}, + }, +} + +def urlopen(request, timeout=None): + token = request.get_header("Authorization").removeprefix("Bearer ") + with open(os.environ["ASKED"], "a") as asked: + asked.write(token + "\n") + with open(os.environ["HOSTS"], "a") as hosts: + hosts.write(request.full_url.split("/")[2] + "\n") + if token not in os.environ["ACCEPTED"].split(): + raise urllib.error.HTTPError(request.full_url, 403, "Forbidden", {}, io.BytesIO()) + endpoint = request.full_url.rsplit(":", 1)[1] + return io.BytesIO(json.dumps(plans_without_quota.get(token, answers)[endpoint]).encode()) + +collector.urllib.request.urlopen = urlopen +sys.argv = ["omarchy-agent-usage-agy", "--force"] +collector.main() +PY +} + +asked() { + sort -u "$test_tmp/asked" | paste -sd' ' + rm -f "$test_tmp/asked" +} + +hosts() { + sort -u "$test_tmp/hosts" 2>/dev/null | paste -sd' ' + rm -f "$test_tmp/hosts" +} + +omp_login omp-token "$(in_ms 3600)" me@example.com +record=$(collect omp-token) +[[ $(jq -c '{ready, tierLabel, stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == '{"ready":true,"tierLabel":"Pro","stale":false,"percent":0.25}' ]] || + fail "Antigravity collector reports limits from an omp sign-in without agy" "$record" +[[ $(asked) == "omp-token" ]] || fail "Antigravity collector asks Google with omp's token" "$record" +pass "Antigravity collector reports limits from an omp sign-in without agy" + +# The bearer token can come from another harness' store, so a destination +# named in the environment is taken only over https. +record=$(collect omp-token "http://evil.example") +[[ $(hosts) == "daily-cloudcode-pa.googleapis.com" && $(jq '.limits[0].percent' <<<"$record") == "0.25" ]] || + fail "Antigravity collector ignores a plain-http destination" "$record" +record=$(collect omp-token "https://mirror.example") +[[ $(hosts) == "mirror.example" && $(jq '.limits[0].percent' <<<"$record") == "0.25" ]] || + fail "Antigravity collector honors an https destination" "$record" +asked >/dev/null +pass "Antigravity collector asks only an https destination, defaulting to Google's" + +pi_login pi-token "$(in_ms 7200)" me@example.com +record=$(collect "omp-token pi-token") +[[ $(asked) == "pi-token" ]] || + fail "Antigravity collector asks with the most recently refreshed sign-in" "$record" +pass "Antigravity collector asks with the most recently refreshed sign-in" + +record=$(collect omp-token) +[[ $(asked) == "omp-token pi-token" && $(jq -c '{stale: .limitsStale, usageStatusText}' <<<"$record") == '{"stale":false,"usageStatusText":""}' ]] || + fail "Antigravity collector moves on from a sign-in Google refuses" "$record" +pass "Antigravity collector moves on from a sign-in Google refuses" + +# The account that earned the kept limits is still signed in, so they stand +# in when a newer sign-in is refused and the rest cannot answer either. +mkdir -p "$test_tmp/.local/share/opencode" "$test_tmp/.config/opencode" +printf '{"google":{"type":"oauth","access":"newer-token","refresh":"newer|p","expires":%s}}\n' "$(in_ms 9999)" \ + >"$test_tmp/.local/share/opencode/auth.json" +echo '{"version":3,"accounts":[]}' >"$test_tmp/.config/opencode/antigravity-accounts.json" +record=$(collect "") +[[ $(jq -c '{stale: .limitsStale, usageStatusText, percent: .limits[0].percent}' <<<"$record") == '{"stale":true,"usageStatusText":"Antigravity sign-in expired","percent":0.25}' ]] || + fail "Antigravity collector keeps limits whose account is still signed in when every sign-in fails" "$record" +pass "Antigravity collector keeps limits whose account is still signed in when every sign-in fails" +asked >/dev/null +rm -rf "$test_tmp/.local/share/opencode" "$test_tmp/.config/opencode" + +rm -rf "$test_tmp/.pi" +omp_login omp-token "$(in_ms -60)" me@example.com +record=$(collect omp-token) +[[ ! -e $test_tmp/asked ]] || fail "Antigravity collector sends no lapsed token to Google" "$(asked)" +[[ $(jq -c '{ready, stale: .limitsStale, usageStatusText, authHelpText, percent: .limits[0].percent}' <<<"$record") == '{"ready":true,"stale":true,"usageStatusText":"Antigravity sign-in expired","authHelpText":"The Antigravity sign-in omp keeps expired. Start omp to refresh it.","percent":0.25}' ]] || + fail "Antigravity collector keeps the last limits and names the harness whose sign-in lapsed" "$record" +pass "Antigravity collector keeps the last limits and names the harness whose sign-in lapsed" + +omp_login other-token "$(in_ms -60)" other@example.com +record=$(collect "") +[[ $(jq -c '{limits, tierLabel}' <<<"$record") == '{"limits":[],"tierLabel":""}' ]] || + fail "Antigravity collector shows no other Google account's kept limits" "$record" +pass "Antigravity collector shows no other Google account's kept limits" + +# omp refreshes its own sign-in when asked for usage, so a lapsed token in +# its database still gets live limits through `omp usage`. +cat >"$test_tmp/omp-report.json" <<'EOF' +{"reports":[{"provider":"google-antigravity","metadata":{"email":"other@example.com"},"limits":[ + {"label":"Gemini","amount":{"usedFraction":0.6},"window":{"id":"5h","resetsAt":32503680000000}}, + {"label":"Gemini","amount":{"usedFraction":0.3},"window":{"id":"weekly","resetsAt":32503680000000}}, + {"label":"Claude & GPT (shared)","amount":{"usedFraction":0.1},"window":{"id":"weekly","resetsAt":32503680000000}}, + {"label":"Claude & GPT (shared)","amount":{"usedFraction":0.1},"window":{"id":"weekly","resetsAt":32503680000000}} +]}]} +EOF +record=$(collect "") +[[ ! -e $test_tmp/asked ]] || fail "Antigravity collector sends omp's lapsed token nowhere" "$(asked)" +[[ $(jq -c '{stale: .limitsStale, usageStatusText, limits: [.limits[] | {title, percent}]}' <<<"$record") == '{"stale":false,"usageStatusText":"","limits":[{"title":"Session","percent":0.6},{"title":"Weekly","percent":0.3},{"title":"Claude/GPT Weekly","percent":0.1}]}' ]] || + fail "Antigravity collector takes omp's limits from omp usage" "$record" +rm "$test_tmp/omp-report.json" +pass "Antigravity collector takes omp's limits from omp usage" + +# The plan label kept from an earlier check names the account that earned it. +# omp answers for a different account, which reports limits but no plan. +pi_login pi-token "$(in_ms 3600)" me@example.com +record=$(collect pi-token) +asked >/dev/null +pi_login pi-token "$(in_ms -60)" me@example.com +cmp_report='{"reports":[{"provider":"google-antigravity","metadata":{"email":"other@example.com"},"limits":[ + {"label":"Gemini","amount":{"usedFraction":0.6},"window":{"id":"5h","resetsAt":32503680000000}} +]}]}' +omp_login omp-token "$(in_ms -60)" other@example.com +echo "$cmp_report" >"$test_tmp/omp-report.json" +[[ $(jq -r '.tierLabel' <<<"$record") == "Pro" ]] || + fail "Antigravity collector records the plan of the account that answered" "$record" +record=$(collect "") +rm -f "$test_tmp/asked" +[[ $(jq -c '{tierLabel, percent: .limits[0].percent, usageStatusText}' <<<"$record") == '{"tierLabel":"","percent":0.6,"usageStatusText":""}' ]] || + fail "Antigravity collector keeps no other account's plan label" "$record" +pass "Antigravity collector keeps no other account's plan label" + +# A newer account that reports a plan it has no quota windows for must not +# end up wearing the kept account's meters beside that plan. +pi_login pi-token "$(in_ms -60)" me@example.com +omp_login omp-token "$(in_ms -60)" other@example.com +rm -f "$test_tmp/omp-report.json" +pi_login pi-token "$(in_ms 3600)" me@example.com +record=$(collect pi-token) +asked >/dev/null +pi_login pi-token "$(in_ms -60)" me@example.com +mkdir -p "$test_tmp/.local/share/opencode" "$test_tmp/.config/opencode" +printf '{"google":{"type":"oauth","access":"no-quota-token","expires":%s}}\n' "$(in_ms 9999)" \ + >"$test_tmp/.local/share/opencode/auth.json" +echo '{"version":3,"accounts":[]}' >"$test_tmp/.config/opencode/antigravity-accounts.json" +record=$(collect no-quota-token) +[[ $(jq -c '{limits, tierLabel, usageStatusText}' <<<"$record") == '{"limits":[],"tierLabel":"Ultra","usageStatusText":"Antigravity limits unavailable"}' ]] || + fail "Antigravity collector never pairs one account's plan with another's limits" "$record" +pass "Antigravity collector never pairs one account's plan with another's limits" +rm -rf "$test_tmp/.local/share/opencode" "$test_tmp/.config/opencode" "$test_tmp/.omp" "$test_tmp/.pi/agent/auth.json" +rm -f "$test_tmp/asked" + +# opencode files any Google sign-in under "google"; only the Antigravity +# plugin's accounts file makes it an Antigravity one. +rm -rf "$test_tmp/.omp" "$test_tmp/cache" +mkdir -p "$test_tmp/.local/share/opencode" +printf '{"google":{"type":"oauth","access":"opencode-token","refresh":"r|p","expires":%s}}\n' "$(in_ms 3600)" \ + >"$test_tmp/.local/share/opencode/auth.json" +record=$(collect opencode-token) +[[ $(jq -c '{ready, limits}' <<<"$record") == '{"ready":false,"limits":[]}' && ! -e $test_tmp/asked ]] || + fail "Antigravity collector ignores opencode's Google sign-in without the Antigravity plugin" "$record" +mkdir -p "$test_tmp/.config/opencode" +echo '{"version":3,"accounts":[]}' >"$test_tmp/.config/opencode/antigravity-accounts.json" +record=$(collect opencode-token) +[[ $(asked) == "opencode-token" && $(jq '.limits[0].percent' <<<"$record") == "0.25" ]] || + fail "Antigravity collector reads opencode's Antigravity plugin sign-in" "$record" +pass "Antigravity collector reads opencode's Antigravity plugin sign-in only with the plugin" + +# opencode records no email for its sign-in, so a replaced sign-in has no +# email to compare: without one, the kept limits must go rather than pass +# for the new account's. +printf '{"google":{"type":"oauth","access":"switched-token","refresh":"switched|p","expires":%s}}\n' "$(in_ms 3600)" \ + >"$test_tmp/.local/share/opencode/auth.json" +record=$(collect "") +[[ $(jq -c '{limits, tierLabel, usageStatusText}' <<<"$record") == '{"limits":[],"tierLabel":"","usageStatusText":"Antigravity sign-in expired"}' ]] || + fail "Antigravity collector drops kept limits when the account cannot be matched" "$record" +pass "Antigravity collector drops kept limits when the account cannot be matched" + +# Antigravity used through pi, omp, OpenClaw, and opencode's plugin counts in +# the local stats; Gemini through any other provider does not, and a forked +# session's copy of an answer counts once. +if command -v rg >/dev/null; then + now=$(date -u +%Y-%m-%dT%H:%M:%S.000Z) + mkdir -p "$test_tmp/.omp/agent/sessions/project" "$test_tmp/.openclaw/agents/main/sessions" + answer='{"type":"message","id":"a1","timestamp":"'$now'","message":{"role":"assistant","provider":"google-antigravity","model":"gemini-3-pro","usage":{"input":100,"output":20,"cacheRead":5,"cacheWrite":0}}}' + printf '%s\n%s\n' "$answer" \ + '{"type":"message","id":"a2","timestamp":"'$now'","message":{"role":"assistant","provider":"google-gemini-cli","model":"gemini-3-pro","usage":{"input":999,"output":999}}}' \ + >"$test_tmp/.omp/agent/sessions/project/one.jsonl" + printf '%s\n' "$answer" >"$test_tmp/.omp/agent/sessions/project/fork.jsonl" + # OpenClaw's entries carry no id; two alike answers are two answers. + openclaw='{"type":"message","timestamp":"'$now'","message":{"role":"assistant","provider":"google-antigravity","model":"gemini-3-pro","usage":{"input":1,"output":1}}}' + printf '%s\n%s\n' "$openclaw" "$openclaw" >"$test_tmp/.openclaw/agents/main/sessions/s.jsonl" + sqlite3 "$test_tmp/.local/share/opencode/opencode.db" \ + "CREATE TABLE message (session_id TEXT, data TEXT); + INSERT INTO message VALUES + ('s1', '{\"role\":\"assistant\",\"providerID\":\"google\",\"modelID\":\"antigravity-gemini-3-pro\",\"time\":{\"created\":$(in_ms 0)},\"tokens\":{\"input\":10,\"output\":5,\"reasoning\":1,\"cache\":{\"read\":2,\"write\":0}}}'), + ('s2', '{\"role\":\"assistant\",\"providerID\":\"google\",\"modelID\":\"gemini-3-pro\",\"time\":{\"created\":$(in_ms 0)},\"tokens\":{\"input\":999,\"output\":999}}');" + record=$(collect "") + [[ $(jq -c '{hasLocalStats, totalPrompts, totalSessions, todayTotalTokens, models: (.modelUsage | keys)}' <<<"$record") == '{"hasLocalStats":true,"totalPrompts":4,"totalSessions":3,"todayTotalTokens":147,"models":["antigravity-gemini-3-pro","gemini-3-pro"]}' ]] || + fail "Antigravity collector counts usage from pi, omp, OpenClaw, and opencode's plugin" "$record" + pass "Antigravity collector counts usage from pi, omp, OpenClaw, and opencode's plugin" +else + skip "Antigravity collector counts usage from pi, omp, OpenClaw, and opencode's plugin (no rg)" +fi diff --git a/test/shell.d/agent-usage-agy-limits-test.sh b/test/shell.d/agent-usage-agy-limits-test.sh new file mode 100755 index 00000000000..543ca547ccb --- /dev/null +++ b/test/shell.d/agent-usage-agy-limits-test.sh @@ -0,0 +1,206 @@ +#!/bin/bash + +source "$(dirname "$0")/base-test.sh" + +require_command jq +require_command python3 + +COLLECTOR="$ROOT/bin/omarchy-agent-usage-agy" + +# 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" + +reordered=$(test_limits '{"groups": [ + {"displayName": "Claude and GPT models", "buckets": [{"window": "5h", "remainingFraction": 1.0}]}, + {"displayName": "Gemini Models", "buckets": [{"window": "weekly", "remainingFraction": 0.5}]} +]}') +[[ $(jq -c 'map({label, percent})' <<<"$reordered") == '[{"label":"Weekly (7-day)","percent":0.5}]' ]] || + fail "Antigravity collector finds the Gemini group wherever it is listed" "$reordered" +pass "Antigravity collector finds the Gemini group wherever it is listed" + +# The whole record with Google's endpoints and the keyring standing in: a +# good answer is cached, and a failed check after it keeps those limits, +# marked stale, rather than making any up. +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT +mkdir -p "$test_tmp/bin" +cat >"$test_tmp/bin/secret-tool" <<'EOF' +#!/bin/bash +echo '{"auth_method":"consumer","token":{"access_token":"token-main","expiry":"2999-01-01T00:00:00Z"}}' +EOF +chmod +x "$test_tmp/bin/secret-tool" +reset_at=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) + dt.timedelta(hours=3)).isoformat())') + +collect() { + env -u PI_CODING_AGENT_DIR -u OPENCLAW_STATE_DIR -u XDG_CONFIG_HOME -u XDG_DATA_HOME \ + HOME="$test_tmp" XDG_CACHE_HOME="$test_tmp/cache" PATH="$test_tmp/bin:$PATH" COLLECTOR="$COLLECTOR" RESET_AT="$reset_at" python3 - "$@" <<'PY' +import importlib.machinery, importlib.util, io, json, os, sys, urllib.error + +loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) + +answers = { + "loadCodeAssist": {"paidTier": {"id": "g1-pro-tier", "name": "Google AI Pro"}}, + "retrieveUserQuotaSummary": {"groups": [{"displayName": "Gemini Models", "buckets": [ + {"window": "5h", "remainingFraction": 0.75, "resetTime": os.environ["RESET_AT"]} + ]}]}, +} +state = sys.argv[1] + +def urlopen(request, timeout=None): + if state == "down": + raise urllib.error.URLError("offline") + assert request.get_header("Authorization") == "Bearer token-main" + return io.BytesIO(json.dumps(answers[request.full_url.rsplit(":", 1)[1]]).encode()) + +collector.urllib.request.urlopen = urlopen +sys.argv = ["omarchy-agent-usage-agy", "--force"] +collector.main() +PY +} + +record=$(collect up) +[[ $(jq -c '{id, ready, tierLabel, stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == '{"id":"agy","ready":true,"tierLabel":"Pro","stale":false,"percent":0.25}' ]] || + fail "Antigravity collector reports its plan and live limits" "$record" +(( $(jq '.limitsFetchedAt' <<<"$record") > 0 )) || fail "Antigravity collector says when the limits were fetched" "$record" +pass "Antigravity collector reports its plan and live limits" + +fetched=$(jq '.limitsFetchedAt' <<<"$record") +record=$(collect down) +[[ $(jq -c '{tierLabel, stale: .limitsStale, percent: .limits[0].percent}' <<<"$record") == '{"tierLabel":"Pro","stale":true,"percent":0.25}' ]] || + fail "Antigravity collector keeps the last limits, marked stale, after a failed check" "$record" +[[ $(jq '.limitsFetchedAt' <<<"$record") == "$fetched" ]] || fail "stale Antigravity limits keep when they were fetched" "$record" +pass "Antigravity collector keeps the last limits, marked stale, after a failed check" + +rm -rf "$test_tmp/cache" +record=$(collect down) +[[ $(jq -c '{limits, usageStatusText}' <<<"$record") == '{"limits":[],"usageStatusText":"Antigravity limits unavailable"}' ]] || + fail "Antigravity collector invents no limits when it has never reached Google" "$record" +pass "Antigravity collector invents no limits when it has never reached Google" diff --git a/test/shell.d/agent-usage-agy-scanner-test.sh b/test/shell.d/agent-usage-agy-scanner-test.sh new file mode 100755 index 00000000000..d65cb168a47 --- /dev/null +++ b/test/shell.d/agent-usage-agy-scanner-test.sh @@ -0,0 +1,108 @@ +#!/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 a signed-in keyring entry without a token, so +# nothing is asked of Google. +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 +timestamp="$(date -u +%Y-%m-%dT%H:%M:%SZ)" + +cat >"$AGY_DIR/history.jsonl" <"$AGY_DIR/brain/test-conv/.system_generated/logs/transcript.jsonl" <>"$AGY_DIR/brain/test-conv/.system_generated/logs/transcript.jsonl" <>"$AGY_DIR/brain/test-conv/.system_generated/logs/transcript.jsonl" +result=$(collect) +older_day=$(date -d "$older" +%Y-%m-%d) + +[[ $(jq -c '{activeDays, hasOlderDay: (.activeDates | index("'"$older_day"'") != null), olderDay: ([.recentDays[] | select(.date == "'"$older_day"'") | .messageCount] | add)}' <<<"$result") == '{"activeDays":2,"hasOlderDay":true,"olderDay":15}' ]] || + fail "Antigravity collector counts a transcript-only day as an active day" "$result" +pass "Antigravity collector counts a transcript-only day as an active day" + +[[ $(jq -c '{limits, usageStatusText}' <<<"$result") == '{"limits":[],"usageStatusText":"No Antigravity sign-in"}' ]] || + fail "Antigravity collector makes up no limits from prompt counts" "$result" +pass "Antigravity collector makes up no limits from prompt counts" + +# Nobody signed in and no history: a record the panel skips. +cat >"$TEST_HOME/bin/secret-tool" <<'EOF' +#!/bin/bash +exit 1 +EOF +rm -rf "$AGY_DIR" +result=$(collect) + +[[ $(jq -c '{ready, tierLabel, limits, totalPrompts, totalSessions, activeDays, modelUsage}' <<<"$result") == '{"ready":false,"tierLabel":"","limits":[],"totalPrompts":0,"totalSessions":0,"activeDays":0,"modelUsage":{}}' ]] || + fail "Antigravity collector prints an empty record without Antigravity" "$result" +pass "Antigravity collector prints an empty record without Antigravity" diff --git a/test/shell.d/agent-usage-agy-timestamps-test.sh b/test/shell.d/agent-usage-agy-timestamps-test.sh new file mode 100755 index 00000000000..189fa4c8d9f --- /dev/null +++ b/test/shell.d/agent-usage-agy-timestamps-test.sh @@ -0,0 +1,305 @@ +#!/bin/bash + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +require_command python3 + +python3 - "$ROOT/bin/omarchy-agent-usage-agy" <<'PY' +import contextlib +import io +import json +import os +import runpy +import sqlite3 +import subprocess +import sys +import tempfile +import time +import types +import unittest +from datetime import datetime, timezone +from pathlib import Path +from unittest.mock import Mock, patch + +collector_path = Path(sys.argv.pop()) +collector = types.SimpleNamespace(**runpy.run_path(str(collector_path))) +# Functions loaded by runpy share this namespace, rather than the wrapper. +namespace = collector.local_stats.__globals__ +# These cases are about agy's own records; the host's other harnesses stay out. +namespace['harness_messages'] = lambda: iter(()) +NOW = datetime(2026, 10, 2, 12, tzinfo=timezone.utc) +ZONES = ('UTC0', 'EST5', 'JST-9') # Fixed offsets; no host zoneinfo or DST rules. +INVALID = (None, '', 'not a date', '2026-02-30T12:00:00Z', [], {}, True, + '0001-01-01 00:00:00+00:00', '0001-01-01T00:00:00Z', + '0001-01-01T00:00:00', '0001-01-01T01:00:00+01:00', + 10**30, -(10**30), 10**400, '9' * 400, float('inf'), float('nan')) + + +@contextlib.contextmanager +def local_zone(zone): + old = os.environ.get('TZ') + try: + os.environ['TZ'] = zone + time.tzset() + yield + finally: + if old is None: + os.environ.pop('TZ', None) + else: + os.environ['TZ'] = old + time.tzset() + + +class FrozenDatetime(datetime): + @classmethod + def now(cls, tz=None): + return NOW.astimezone(tz) if tz else NOW.astimezone().replace(tzinfo=None) + + +def write_fixture(root, invalid): + history = [dict(timestamp='2026-10-02T02:00:00Z', conversation_id='valid')] + history += [dict(timestamp=value, conversation_id='bad-history') for value in invalid] + history += [{}, dict(timestamp='1969-07-20T12:00:00Z', conversation_id='old'), + dict(timestamp='2026-10-02T12:00:00Z', conversation_id='valid')] + (root / 'history.jsonl').write_text('\n'.join(map(json.dumps, history))) + with contextlib.closing(sqlite3.connect(root / 'conversation_summaries.db')) as db, db: + db.execute('CREATE TABLE conversation_summaries (conversation_id TEXT, last_modified_time TEXT)') + db.executemany('INSERT INTO conversation_summaries VALUES (?, ?)', [ + ('valid', '2026-10-02T12:00:00Z'), ('undated', '0001-01-01 00:00:00+00:00'), + ('undated', None), ('undated', 'not a date'), + ('summary-only', '2026-10-02T12:00:00Z')]) + transcript = root / 'brain/valid/.system_generated/logs/transcript.jsonl' + transcript.parent.mkdir(parents=True) + steps = [dict(created_at='2026-10-02T12:00:00Z', model='dated', + usage=dict(input_tokens=10, output_tokens=5))] + steps += [dict(created_at=value, model='undated', usage=dict(input_tokens=2, output_tokens=1)) + for value in invalid] + steps += [dict(model='undated', usage=dict(input_tokens=2, output_tokens=1)), + dict(created_at='2026-10-02T02:00:00Z', model='dated', + usage=dict(input_tokens=7, output_tokens=3))] + transcript.write_text('\n'.join(map(json.dumps, steps))) + + +class TimestampTests(unittest.TestCase): + def test_zero_and_invalid_times_are_undated_in_every_zone(self): + for zone in ZONES: + with local_zone(zone): + for value in INVALID: + with self.subTest(zone=zone, value=value): + self.assertEqual(collector.local_day(value), '') + self.assertIsNone(collector.parse_time('0001-01-01 00:00:00+00:00')) + + def test_valid_times_keep_existing_formats_and_local_dates(self): + for zone, early_day in zip(ZONES, ('2026-10-02', '2026-10-01', '2026-10-02')): + with local_zone(zone), self.subTest(zone=zone): + for value in ('2026-10-02T02:00:00Z', '2026-10-02 02:00:00', + '2026-10-02T04:00:00+02:00', 1790906400, 1790906400000, + '1790906400000'): + self.assertEqual(collector.local_day(value), early_day) + self.assertEqual(collector.local_day('1969-07-20T12:00:00Z'), '1969-07-20') + self.assertEqual(collector.local_day(0), '1969-12-31' if zone == 'EST5' else '1970-01-01') + self.assertEqual(collector.local_day('2026-10-02T20:00:00Z'), + '2026-10-03' if zone == 'JST-9' else '2026-10-02') + + def test_conversion_boundaries(self): + cases = ( + ('0001-01-01T00:00:01Z', ('0001-01-01', '', '0001-01-01')), + ('0001-01-01T12:00:00Z', ('0001-01-01', '0001-01-01', '0001-01-01')), + ('9999-12-31T23:59:59Z', ('9999-12-31', '9999-12-31', '')), + ('9999-12-31T12:00:00Z', ('9999-12-31', '9999-12-31', '9999-12-31')), + ('0001-01-01T00:00:01+14:00', ('', '', '')), + ('9999-12-31T23:59:59-12:00', ('', '', '')), + ) + for value, expected in cases: + for zone, day in zip(ZONES, expected): + with local_zone(zone), self.subTest(zone=zone, value=value): + self.assertEqual(collector.local_day(value), day) + + def test_mixed_activity_survives_without_inventing_dates(self): + for zone in ZONES: + # These offset timestamps overflow while normalizing to UTC in all zones. + invalid = INVALID + ('0001-01-01T00:00:01+14:00', '9999-12-31T23:59:59-12:00') + if zone == 'EST5': + invalid += ('0001-01-01T00:00:01Z',) + elif zone == 'JST-9': + invalid += ('9999-12-31T23:59:59Z',) + with local_zone(zone), tempfile.TemporaryDirectory() as scratch, self.subTest(zone=zone): + root = Path(scratch) + write_fixture(root, invalid) + with patch.dict(namespace, datetime=FrozenDatetime): + stats = collector.local_stats(root) + self.assertTrue(stats['hasLocalStats']) + self.assertEqual(stats['totalPrompts'], 3) + self.assertEqual(stats['todayPrompts'], 1 if zone == 'EST5' else 2) + # Summary identities remain countable without a date; bad history isn't. + self.assertEqual(stats['totalSessions'], 4) + self.assertEqual(stats['todaySessions'], 2) + dates = ['1969-07-20', '2026-10-02'] + if zone == 'EST5': + dates.insert(1, '2026-10-01') + self.assertEqual(stats['activeDates'], dates) + self.assertEqual(stats['activeDays'], len(dates)) + self.assertEqual(stats['todayTotalTokens'], 15 if zone == 'EST5' else 25) + self.assertEqual(stats['todayTokensByModel'], {'dated': stats['todayTotalTokens']}) + recent = {row['date']: row['messageCount'] for row in stats['recentDays']} + self.assertEqual(recent['2026-10-02'], stats['todayTotalTokens']) + self.assertEqual(sum(recent.values()), 25) + self.assertEqual(stats['modelUsage']['dated']['inputTokens'], 17) + self.assertEqual(stats['modelUsage']['dated']['outputTokens'], 8) + # Undated usage retains lifetime totals, just like a missing timestamp. + self.assertEqual(stats['modelUsage']['undated']['inputTokens'], 2 * (len(invalid) + 1)) + self.assertEqual(stats['modelUsage']['undated']['outputTokens'], len(invalid) + 1) + + def test_login_expiry_and_quota_probe_policy(self): + # Unknown expiry cannot prove expiration or validity: ask the service. + # Known expired tokens should not be sent; a future expiry permits a probe. + cases = [('missing', {}, 0, False)] + cases += [(repr(value), {'expiry': value}, 0, False) for value in INVALID] + cases += [ + ('expired', {'expiry': '2026-10-02T11:00:00Z'}, NOW.timestamp() - 3600, True), + ('expires now', {'expiry': NOW.isoformat()}, NOW.timestamp(), True), + ('future', {'expiry': '2026-10-02T15:00:00+02:00'}, NOW.timestamp() + 3600, False), + ('future milliseconds', {'expiry': (NOW.timestamp() + 3600) * 1000}, + NOW.timestamp() + 3600, False), + ] + answers = [ + {'paidTier': {'name': 'Google AI Pro'}}, + {'groups': [{'displayName': 'Gemini', 'buckets': [ + {'window': '5h', 'remainingFraction': 0.75, 'resetTime': '2026-10-02T13:00:00Z'}]}]}, + ] + for zone in ZONES: + for location in ('token', 'root'): + for name, expiry, expected, expired in cases: + with local_zone(zone), tempfile.TemporaryDirectory() as scratch, self.subTest( + zone=zone, location=location, expiry=name): + payload = {'token': {'access_token': 'synthetic-token'}} + (payload['token'] if location == 'token' else payload).update(expiry) + probe = Mock(side_effect=answers) + with patch.object(collector.subprocess, 'check_output', return_value=json.dumps(payload)), \ + patch.dict(namespace, datetime=FrozenDatetime, cache_root=lambda: Path(scratch), + time=types.SimpleNamespace(time=NOW.timestamp), post=probe): + entry = collector.login() + self.assertEqual(entry['expiresAt'], expected) + self.assertEqual(entry['accessToken'], 'synthetic-token') + result = collector.collect_limits([entry], force=True) + if expired: + probe.assert_not_called() + self.assertEqual(result['usageStatusText'], 'Antigravity sign-in expired') + self.assertFalse(result['live']) + else: + self.assertEqual(probe.call_args_list, [ + unittest.mock.call('loadCodeAssist', 'synthetic-token'), + unittest.mock.call('retrieveUserQuotaSummary', 'synthetic-token')]) + self.assertTrue(result['live']) + self.assertEqual(result['usageStatusText'], '') + self.assertEqual(result['limits'][0]['percent'], 0.25) + if expected == 0: + # An unknown local expiry must still honor an authoritative 401. + probe.reset_mock(side_effect=True) + with contextlib.closing(collector.urllib.error.HTTPError( + 'https://example.invalid', 401, 'Unauthorized', {}, io.BytesIO())) as rejection: + probe.side_effect = rejection + rejected = collector.collect_limits([entry], force=True) + probe.assert_called_once_with('loadCodeAssist', 'synthetic-token') + self.assertEqual(rejected['usageStatusText'], 'Antigravity sign-in expired') + self.assertFalse(rejected['live']) + + def test_quota_reset_display_keeps_unknown_and_valid_windows(self): + # A usable quota remains useful without a trustworthy reset countdown. + cases = [({}, '')] + [({'resetTime': value}, '') for value in INVALID] + cases += [ + ({'resetTime': '2026-10-02T11:00:00Z'}, '2026-10-02T11:00:00+00:00'), + ({'resetTime': '2026-10-02T15:00:00+02:00'}, '2026-10-02T15:00:00+02:00'), + ({'resetTime': '2026-10-02 13:00:00'}, '2026-10-02T13:00:00+00:00'), + ({'resetTime': (NOW.timestamp() + 3600) * 1000}, '2026-10-02T13:00:00+00:00'), + ] + for zone in ZONES: + for reset, expected in cases: + with local_zone(zone), self.subTest(zone=zone, reset=reset): + bucket = dict(reset, window='5h', remainingFraction=0.75) + self.assertEqual(collector.reset_of(bucket), expected) + limits = collector.extract_limits({'groups': [ + {'displayName': 'Gemini', 'buckets': [bucket]}]}) + self.assertEqual(len(limits), 1) + self.assertEqual(limits[0]['resetsAt'], expected) + self.assertEqual(limits[0]['percent'], 0.25) + + def test_cached_windows_prune_only_known_expired_resets(self): + # Unknown reset is no evidence the cached allowance has ended. Retain it + # as stale during an outage, without inventing a reset or a fresh fetch. + unknown = [{'label': 'missing', 'percent': 0.25}] + unknown += [dict(label=f'unknown-{index}', percent=0.25, resetsAt=value) + for index, value in enumerate(INVALID)] + expired = [dict(label='expired', resetsAt='2026-10-02T11:00:00Z'), + dict(label='expires now', resetsAt=NOW.isoformat())] + future = [dict(label='future', percent=0.5, resetsAt='2026-10-02T15:00:00+02:00'), + dict(label='future milliseconds', percent=0.5, + resetsAt=(NOW.timestamp() + 3600) * 1000)] + windows = unknown + expired + future + fetched = (NOW.timestamp() - 120) * 1000 + for zone in ZONES: + with local_zone(zone), tempfile.TemporaryDirectory() as scratch, self.subTest(zone=zone): + probe = Mock(side_effect=collector.urllib.error.URLError('synthetic outage')) + with patch.dict(namespace, datetime=FrozenDatetime, cache_root=lambda: Path(scratch), + time=types.SimpleNamespace(time=NOW.timestamp), post=probe): + self.assertEqual(collector.open_windows(windows), unknown + future) + self.assertEqual(collector.open_windows(None), []) + # A damaged cache entry is dropped, not raised over. + self.assertEqual(collector.open_windows([None, "junk", 5, dict(label='junk')]), + [dict(label='junk')]) + entry = dict(source='agy', accessToken='synthetic-token', expiresAt=NOW.timestamp() + 3600) + collector.write_json(Path(scratch) / 'agy-limits.json', + dict(limits=windows, tierLabel='Pro', fetchedAtMs=fetched, + identity=collector.identity_of(entry))) + result = collector.collect_limits([entry], force=True) + probe.assert_called_once_with('loadCodeAssist', 'synthetic-token') + # JSON round trips create a new NaN, which cannot compare equal to + # itself; compare serialized windows to also cover corrupt caches. + self.assertEqual(json.dumps(result['limits']), json.dumps(unknown + future)) + self.assertFalse(result['live']) + self.assertEqual(result['fetchedAtMs'], fetched) + self.assertEqual(result['tierLabel'], 'Pro') + + def test_unrelated_errors_propagate(self): + with patch.dict(namespace, datetime=types.SimpleNamespace( + fromisoformat=lambda value: (_ for _ in ()).throw(RuntimeError('unexpected parser failure')))): + with self.assertRaisesRegex(RuntimeError, 'unexpected parser failure'): + collector.local_day('2026-10-02T12:00:00Z') + with patch.dict(namespace, parse_time=lambda value: types.SimpleNamespace( + astimezone=lambda: (_ for _ in ()).throw(RuntimeError('unexpected conversion failure')))): + with self.assertRaisesRegex(RuntimeError, 'unexpected conversion failure'): + collector.local_day('2026-10-02T12:00:00Z') + + def test_isolated_cli_smoke_and_cached_record(self): + for zone in ZONES: + with tempfile.TemporaryDirectory() as scratch, self.subTest(zone=zone): + root = Path(scratch) + write_fixture(root, ('0001-01-01 00:00:00+00:00', None, 'not a date')) + stub_dir = root / 'bin' + stub_dir.mkdir() + secret_tool = stub_dir / 'secret-tool' + secret_tool.write_text('#!/bin/bash\nexit 1\n') + secret_tool.chmod(0o755) + env = {key: value for key, value in os.environ.items() + if key not in ('PI_CODING_AGENT_DIR', 'OPENCLAW_STATE_DIR', 'XDG_CONFIG_HOME', 'XDG_DATA_HOME')} + env.update(HOME=scratch, AGY_DIR=scratch, XDG_CACHE_HOME=str(root / 'cache'), + PATH=str(stub_dir) + os.pathsep + os.defpath, TZ=zone) + for flag in ('--force', '--limits-only'): + result = subprocess.run([str(collector_path), flag], env=env, capture_output=True, + text=True, check=True, timeout=10) + record = json.loads(result.stdout) + self.assertEqual(result.stderr, '') + self.assertEqual(record['id'], 'agy') + self.assertFalse(record['ready']) + self.assertEqual(record['limits'], []) + self.assertEqual(record['totalPrompts'], 3) + self.assertEqual(record['totalSessions'], 4) + self.assertEqual(record['modelUsage']['dated']['inputTokens'], 17) + + +unittest.main(verbosity=2) +PY + +pass "Antigravity timestamp boundaries, mixed records, and isolated CLI smoke"