Google Antigravity

#!/usr/bin/env python3
"""Antigravity usage in the title bar: how much of the Gemini weekly limit and the five-hour limit
you have used, as Antigravity's own Settings, Models & Usage page shows them. Reads the running
Antigravity hub in this workspace (agy --hub: its port and CSRF token are on its command line) and
asks it for RetrieveUserQuotaSummary, the call that page makes. No model runs, so no quota is spent.
Asks at most once every three minutes, pauses while the window is hidden, retries at 5, 15, 45 s on
errors. Click the slot to cycle the views."""
import json, os, re, subprocess, time, urllib.request
from datetime import datetime
from importlib.machinery import SourceFileLoader
HERE = os.path.dirname(os.path.realpath(__file__))
marks = SourceFileLoader('agmarks', os.path.join(HERE, 'hydrogen-widget-marks.py')).load_module()
lib = SourceFileLoader('hwlib', os.path.join(HERE, 'hydrogen-widget-lib.py')).load_module()

def hubs():
    out = subprocess.run(['ps', '-eo', 'args'], capture_output=True, text=True).stdout
    found = []
    for line in out.splitlines():
        if re.search(r'(^|/)agy\s.*--hub', line):
            p = re.search(r'--hub-port[= ](\d+)', line); c = re.search(r'--csrf_token[= ](\S+)', line)
            if p and c: found.append((int(p.group(1)), c.group(1)))
    return found

AGY = os.path.expanduser('~/.gemini/bin/agy')

def own_hub():
    """No Antigravity hub running (no Antigravity tab open, or the editor just restarted): start a
    private one just for this reading, the same way the extension does (AGY_ENABLE_HUB and
    ANTIGRAVITY_VSCODE_HOST make agy run headless; without them it wants a terminal). John
    2026-10-19: "i should NEVER have to go open an antigravity tab just to get this widget to work".
    Returns (port, csrf, process) once it answers, or None. The caller stops it after the reading:
    a hub holds about 200 MB, too much to leave running for a title-bar number."""
    import shutil, socket, uuid
    exe = AGY if os.path.exists(AGY) else shutil.which('agy')
    if not exe: return None
    s = socket.socket(); s.bind(('127.0.0.1', 0)); port = s.getsockname()[1]; s.close()
    csrf = str(uuid.uuid4())
    env = dict(os.environ, AGY_ENABLE_HUB='1', ANTIGRAVITY_VSCODE_HOST='1')
    p = subprocess.Popen([exe, '--hub', f'--hub-port={port}', '--app_data_dir=antigravity', f'--csrf_token={csrf}'],
                         env=env, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True)
    for _ in range(60):
        if p.poll() is not None: return None
        try:
            rpc(port, csrf, 'GetUserStatus'); return port, csrf, p
        except Exception: time.sleep(0.5)
    stop(p); return None

def stop(p):
    try:
        p.terminate(); p.wait(timeout=5)
    except Exception:
        try: p.kill()
        except Exception: pass

def rpc(port, csrf, method):
    req = urllib.request.Request(f'http://127.0.0.1:{port}/exa.language_server_pb.LanguageServerService/{method}', data=b'{}',
                                 headers={'Content-Type': 'application/json', 'x-codeium-csrf-token': csrf}, method='POST')
    with urllib.request.urlopen(req, timeout=15) as r:
        j = json.loads(r.read().decode())
        return j.get('response') or j   # most calls wrap the answer in 'response'; GetUserStatus does not

def epoch(s):
    try: return int(datetime.fromisoformat(s.replace('Z', '+00:00')).timestamp())
    except Exception: return None

def data():
    hs = [] if os.environ.get('HW_AGY_OWN_HUB') else hubs()   # HW_AGY_OWN_HUB: test the private-hub path
    mine = None
    if not hs:
        mine = own_hub()
        if not mine:
            return {'error': 'Antigravity is not installed in this workspace' if not (os.path.exists(AGY) or __import__('shutil').which('agy')) else 'Antigravity did not start; retrying'}
        hs = [mine[:2]]
    port, csrf = hs[0]
    try:
        return reading(port, csrf, hs, mine is not None)
    finally:
        if mine: stop(mine[2])

def reading(port, csrf, hs, private):
    q = rpc(port, csrf, 'RetrieveUserQuotaSummary')
    plan = None
    try: plan = (rpc(port, csrf, 'GetLoadCodeAssist').get('paidTier') or {}).get('name')
    except Exception: pass
    items = []
    for g in q.get('groups') or []:
        gem = 'Gemini' in (g.get('displayName') or '')
        for b in g.get('buckets') or []:
            weekly = b.get('window') == 'weekly'
            used = round((1 - float(b.get('remainingFraction', 1))) * 100)
            items.append({'label': ('week' if weekly else '5h') if gem else ('3p wk' if weekly else '3p 5h'), 'pct': used,
                          'severity': 'critical' if used >= 90 else ('warning' if used >= 75 else 'normal'),
                          'name': ('Gemini ' if gem else 'Claude and GPT ') + ('this week' if weekly else 'five-hour window'),
                          'resets_at': epoch(b.get('resetTime') or ''), 'window_s': 7 * 86400 if weekly else 5 * 3600, 'gemini': gem})
    items.sort(key=lambda i: (not i.pop('gemini'),))
    # Which Google account these numbers belong to (GetUserStatus), so a second sign-in is never a
    # mystery (John 2026-10-18: this showed 0% for his own account while the team account was at 14%).
    account = None
    try: account = (rpc(port, csrf, 'GetUserStatus').get('userStatus') or {}).get('email')
    except Exception: pass
    if account and len(hs) > 1: account += f' (the first of {len(hs)} Antigravity hubs running here)'
    return {'plan': plan, 'account': account, 'items': items}

RENDER = r"""(d.error?`<span class="err">${d.error}</span>`:(d.items||[]).slice(0,2).map(b=>bar(b.label,b.pct,b.severity)).join(''))"""
lib.serve('Antigravity usage', data, RENDER, 180000, mark=marks.MARKS['antigravity'], default_view='big:week')