opencode-antigravity-quota-tmux-script

opencode-antigravity-quota-tmux-script

#!/usr/bin/env bash
set -o pipefail

# =============================================================================
# ANTIGRAVITY CONSTANTS
# =============================================================================
readonly GOOGLE_TOKEN_URL="https://oauth2.googleapis.com/token"
readonly CLOUDCODE_BASE_URL="https://cloudcode-pa.googleapis.com"
readonly CLIENT_ID="1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
readonly CLIENT_SECRET="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"
readonly AG_ERROR_OUTPUT="C:--% • G:--%"

# Antigravity cache settings
readonly AG_CACHE_DIR="${HOME}/.cache/antigravity-quota"
readonly AG_CACHE_FILE="${AG_CACHE_DIR}/quota.json"
readonly DEFAULT_CACHE_TTL=60

# =============================================================================
# CODEX CONSTANTS
# =============================================================================
readonly CODEX_API_URL="https://chatgpt.com/backend-api/wham/usage"
readonly CODEX_AUTH_ENDPOINT="https://auth.openai.com/oauth/token"
readonly CODEX_CLIENT_ID="app_EMoamEEZ73f0CkXaXp7hrann"
readonly CODEX_ERROR_OUTPUT="O:--%"
readonly CODEX_TIMEOUT=5
readonly CODEX_REFRESH_DAYS=8

# Codex cache settings
readonly CODEX_CACHE_DIR="${HOME}/.cache/codex-quota"
readonly CODEX_CACHE_FILE="${CODEX_CACHE_DIR}/usage.json"

# =============================================================================
# FIRMWARE CONSTANTS
# =============================================================================
readonly FIRMWARE_API_URL="https://app.firmware.ai/api/v1/quota"
readonly FIRMWARE_AUTH_FILE="${HOME}/.local/share/opencode/auth.json"
readonly FIRMWARE_ERROR_OUTPUT="F:--%"
readonly FIRMWARE_TIMEOUT=5
readonly FIRMWARE_CACHE_TTL=0  # Always fetch fresh

# Firmware cache settings
readonly FIRMWARE_CACHE_DIR="${HOME}/.cache/firmware-quota"
readonly FIRMWARE_CACHE_FILE="${FIRMWARE_CACHE_DIR}/usage.json"

# Config paths to search
readonly CONFIG_PATHS=(
  "$HOME/.config/opencode/antigravity-accounts.json"
  "$HOME/.local/share/opencode/antigravity-accounts.json"
)

# Find first existing config file
find_config_file() {
  for path in "${CONFIG_PATHS[@]}"; do
    if [[ -f "$path" ]]; then
      echo "$path"
      return 0
    fi
  done
  return 1
}

# Get active account from config
get_active_account() {
  local config_file="$1"
  local active_index
  
  active_index=$(jq -r '.activeIndex // 0' "$config_file" 2>/dev/null)
  jq -r ".accounts[$active_index]" "$config_file" 2>/dev/null
}

# Refresh access token using OAuth
refresh_access_token() {
  local refresh_token="$1"
  local response
  
  # Strip |projectId suffix if present (some auth plugins append this)
  local clean_token="${refresh_token%%|*}"
  
  response=$(curl -s -X POST "$GOOGLE_TOKEN_URL" \
    -H "Content-Type: application/x-www-form-urlencoded" \
    -d "client_id=$CLIENT_ID" \
    -d "client_secret=$CLIENT_SECRET" \
    -d "refresh_token=$clean_token" \
    -d "grant_type=refresh_token")
  
  echo "$response" | jq -r '.access_token // empty' 2>/dev/null
}

# Load code assist to discover project ID
load_code_assist() {
  local access_token="$1"
  
  curl -s -X POST "${CLOUDCODE_BASE_URL}/v1internal:loadCodeAssist" \
    -H "Authorization: Bearer $access_token" \
    -H "Content-Type: application/json" \
    -H "User-Agent: antigravity" \
    -d '{"metadata": {"ideType": "ANTIGRAVITY", "platform": "PLATFORM_UNSPECIFIED", "pluginType": "GEMINI"}}'
}

# Get project ID from account or fetch from API
get_project_id() {
  local account="$1"
  local access_token="$2"
  local project_id
  
  # Try account's projectId first
  project_id=$(echo "$account" | jq -r '.projectId // empty' 2>/dev/null)
  
  if [[ -n "$project_id" && "$project_id" != "null" ]]; then
    echo "$project_id"
    return 0
  fi
  
  # Try managedProjectId
  project_id=$(echo "$account" | jq -r '.managedProjectId // empty' 2>/dev/null)
  
  if [[ -n "$project_id" && "$project_id" != "null" ]]; then
    echo "$project_id"
    return 0
  fi
  
  # Fetch from loadCodeAssist API
  local response
  response=$(load_code_assist "$access_token")
  
  # Extract from cloudaicompanionProject (can be string or object with id)
  project_id=$(echo "$response" | jq -r '.cloudaicompanionProject // empty' 2>/dev/null)
  
  # If it's an object, extract the id field
  if [[ "$project_id" == "{"* ]]; then
    project_id=$(echo "$response" | jq -r '.cloudaicompanionProject.id // empty' 2>/dev/null)
  fi
  
  if [[ -n "$project_id" && "$project_id" != "null" ]]; then
    echo "$project_id"
    return 0
  fi
  
  # No project ID found - return empty (caller handles error)
  return 1
}

# Fetch quota from API
fetch_quota() {
  local access_token="$1"
  local project_id="$2"
  
  local payload='{}'
  if [[ -n "$project_id" ]]; then
    payload="{\"project\": \"${project_id}\"}"
  fi
  
  curl -s -X POST "${CLOUDCODE_BASE_URL}/v1internal:fetchAvailableModels" \
    -H "Authorization: Bearer $access_token" \
    -H "Content-Type: application/json" \
    -H "User-Agent: antigravity" \
    -d "$payload"
}

# Ensure cache directory exists
ensure_ag_cache_dir() {
  [[ -d "$AG_CACHE_DIR" ]] || mkdir -p "$AG_CACHE_DIR"
}

ensure_codex_cache_dir() {
  [[ -d "$CODEX_CACHE_DIR" ]] || mkdir -p "$CODEX_CACHE_DIR"
}

# Check if cache is valid (exists and not expired)
is_ag_cache_valid() {
  local ttl="${QUOTA_CACHE_TTL:-$DEFAULT_CACHE_TTL}"
  
  [[ -f "$AG_CACHE_FILE" ]] || return 1
  
  local cache_age now
  now=$(date +%s)
  cache_age=$(stat -f %m "$AG_CACHE_FILE" 2>/dev/null || stat -c %Y "$AG_CACHE_FILE" 2>/dev/null)
  
  (( now - cache_age < ttl ))
}

is_codex_cache_valid() {
  local ttl="${QUOTA_CACHE_TTL:-$DEFAULT_CACHE_TTL}"
  
  [[ -f "$CODEX_CACHE_FILE" ]] || return 1
  
  local cache_age now
  now=$(date +%s)
  cache_age=$(stat -f %m "$CODEX_CACHE_FILE" 2>/dev/null || stat -c %Y "$CODEX_CACHE_FILE" 2>/dev/null)
  
  (( now - cache_age < ttl ))
}

is_firmware_cache_valid() {
  local ttl="${FIRMWARE_CACHE_TTL}"
  
  [[ -f "$FIRMWARE_CACHE_FILE" ]] || return 1
  
  local cache_age now
  now=$(date +%s)
  cache_age=$(stat -f %m "$FIRMWARE_CACHE_FILE" 2>/dev/null || stat -c %Y "$FIRMWARE_CACHE_FILE" 2>/dev/null)
  
  (( now - cache_age < ttl ))
}

# Read cached response
read_ag_cache() {
  cat "$AG_CACHE_FILE" 2>/dev/null
}

read_codex_cache() {
  cat "$CODEX_CACHE_FILE" 2>/dev/null
}

read_firmware_cache() {
  cat "$FIRMWARE_CACHE_FILE" 2>/dev/null
}

# Write response to cache
write_ag_cache() {
  local json="$1"
  ensure_ag_cache_dir
  echo "$json" > "$AG_CACHE_FILE"
}

write_codex_cache() {
  local json="$1"
  ensure_codex_cache_dir
  echo "$json" > "$CODEX_CACHE_FILE"
}

ensure_firmware_cache_dir() {
  [[ -d "$FIRMWARE_CACHE_DIR" ]] || mkdir -p "$FIRMWARE_CACHE_DIR"
}

write_firmware_cache() {
  local json="$1"
  ensure_firmware_cache_dir
  echo "$json" > "$FIRMWARE_CACHE_FILE"
}

# Format quota JSON into tmux-friendly output
format_quota() {
  local json="$1"
  
  # Extract Claude quota (any model containing 'claude')
  local claude_fraction
  claude_fraction=$(echo "$json" | jq -r '
    .models | to_entries[] | 
    select(.key | test("claude"; "i")) | 
    .value.quotaInfo.remainingFraction // empty
  ' 2>/dev/null | head -1)
  
  # Extract Gemini quota (gemini-2.5-flash specifically)
  local gemini_fraction
  gemini_fraction=$(echo "$json" | jq -r '
    .models["gemini-2.5-flash"].quotaInfo.remainingFraction // empty
  ' 2>/dev/null)
  
  # Calculate percentages
  local claude_pct="--"
  local gemini_pct="--"
  
  if [[ -n "$claude_fraction" && "$claude_fraction" != "null" ]]; then
    claude_pct=$(echo "$claude_fraction * 100" | bc | cut -d. -f1)
  fi
  
  if [[ -n "$gemini_fraction" && "$gemini_fraction" != "null" ]]; then
    gemini_pct=$(echo "$gemini_fraction * 100" | bc | cut -d. -f1)
  fi
  
  echo "C:${claude_pct}% • G:${gemini_pct}%"
}

# =============================================================================
# CODEX FUNCTIONS
# =============================================================================

# Find Codex auth file
find_codex_auth_file() {
  local auth_file
  if [[ -n "${CODEX_HOME:-}" ]]; then
    auth_file="${CODEX_HOME}/auth.json"
  else
    auth_file="${HOME}/.codex/auth.json"
  fi
  
  [[ -f "$auth_file" ]] && echo "$auth_file"
}

# Check if Codex token needs refresh (> 8 days since last refresh)
codex_needs_refresh() {
  local auth_file="$1"
  local last_refresh
  
  last_refresh=$(jq -r '.last_refresh // empty' "$auth_file" 2>/dev/null)
  [[ -z "$last_refresh" ]] && return 0
  
  # Parse ISO8601 date and compare to now
  local last_ts now_ts
  last_ts=$(date -j -f "%Y-%m-%dT%H:%M:%S" "${last_refresh%%.*}" +%s 2>/dev/null || \
            date -d "${last_refresh}" +%s 2>/dev/null || echo 0)
  now_ts=$(date +%s)
  
  local days_since=$(( (now_ts - last_ts) / 86400 ))
  (( days_since >= CODEX_REFRESH_DAYS ))
}

# Refresh Codex OAuth token
refresh_codex_token() {
  local auth_file="$1"
  local refresh_token
  
  refresh_token=$(jq -r '.tokens.refresh_token // empty' "$auth_file" 2>/dev/null)
  [[ -z "$refresh_token" ]] && return 1
  
  local response
  response=$(curl -s --max-time "$CODEX_TIMEOUT" -X POST "$CODEX_AUTH_ENDPOINT" \
    -H "Content-Type: application/json" \
    -d "{\"client_id\":\"$CODEX_CLIENT_ID\",\"grant_type\":\"refresh_token\",\"refresh_token\":\"$refresh_token\",\"scope\":\"openid profile email\"}" 2>/dev/null)
  
  # Check for error
  if echo "$response" | jq -e '.error' >/dev/null 2>&1; then
    return 1
  fi
  
  # Extract new tokens
  local new_access new_refresh new_id
  new_access=$(echo "$response" | jq -r '.access_token // empty')
  new_refresh=$(echo "$response" | jq -r '.refresh_token // empty')
  new_id=$(echo "$response" | jq -r '.id_token // empty')
  
  [[ -z "$new_access" ]] && return 1
  
  # Update auth file atomically
  local now_iso
  now_iso=$(date -u +"%Y-%m-%dT%H:%M:%SZ")
  
  local updated
  updated=$(jq --arg at "$new_access" \
               --arg rt "${new_refresh:-$(jq -r '.tokens.refresh_token' "$auth_file")}" \
               --arg it "$new_id" \
               --arg lr "$now_iso" \
    '.tokens.access_token = $at | .tokens.refresh_token = $rt | .tokens.id_token = $it | .last_refresh = $lr' \
    "$auth_file" 2>/dev/null)
  
  [[ -z "$updated" ]] && return 1
  
  echo "$updated" > "$auth_file"
  return 0
}

# Fetch Codex usage from API
fetch_codex_usage() {
  local auth_file="$1"
  local access_token
  
  access_token=$(jq -r '.tokens.access_token // .OPENAI_API_KEY // empty' "$auth_file" 2>/dev/null)
  [[ -z "$access_token" ]] && return 1
  
  curl -s --max-time "$CODEX_TIMEOUT" "$CODEX_API_URL" \
    -H "Authorization: Bearer $access_token" \
    -H "Accept: application/json" 2>/dev/null
}

# Parse Codex remaining percent from API response
parse_codex_remaining() {
  local json="$1"
  local used_pct
  
  used_pct=$(echo "$json" | jq -r '.rate_limit.primary_window.used_percent // empty' 2>/dev/null)
  [[ -z "$used_pct" || "$used_pct" == "null" ]] && return 1
  
  # Calculate remaining and clamp to 0-100
  local remaining
  remaining=$(( 100 - used_pct ))
  (( remaining < 0 )) && remaining=0
  (( remaining > 100 )) && remaining=100
  
  echo "$remaining"
}

# Get Codex quota string
get_codex_quota() {
  # Check cache first
  if is_codex_cache_valid; then
    local cached
    cached=$(read_codex_cache)
    local remaining
    remaining=$(parse_codex_remaining "$cached")
    if [[ -n "$remaining" ]]; then
      echo "O:${remaining}%"
      return 0
    fi
  fi
  
  # Cache miss - fetch from API
  local auth_file
  auth_file=$(find_codex_auth_file)
  [[ -z "$auth_file" ]] && { echo "$CODEX_ERROR_OUTPUT"; return 0; }
  
  # Refresh token if needed
  if codex_needs_refresh "$auth_file"; then
    refresh_codex_token "$auth_file" || true
  fi
  
  # Fetch usage
  local usage_json
  usage_json=$(fetch_codex_usage "$auth_file")
  
  if [[ -z "$usage_json" ]] || echo "$usage_json" | jq -e '.error // .detail' >/dev/null 2>&1; then
    echo "$CODEX_ERROR_OUTPUT"
    return 0
  fi
  
  # Cache valid response
  write_codex_cache "$usage_json"
  
  # Parse and format
  local remaining
  remaining=$(parse_codex_remaining "$usage_json")
  if [[ -n "$remaining" ]]; then
    echo "O:${remaining}%"
  else
    echo "$CODEX_ERROR_OUTPUT"
  fi
}

# =============================================================================
# FIRMWARE FUNCTIONS
# =============================================================================

# Fetch Firmware usage from API
fetch_firmware_usage() {
  [[ ! -f "$FIRMWARE_AUTH_FILE" ]] && return 1
  
  local api_key
  api_key=$(jq -r '.firmware.key // empty' "$FIRMWARE_AUTH_FILE" 2>/dev/null)
  [[ -z "$api_key" ]] && return 1
  
  curl -s --max-time "$FIRMWARE_TIMEOUT" "$FIRMWARE_API_URL" \
    -H "Authorization: Bearer $api_key" 2>/dev/null
}

# Parse Firmware remaining percent from API response
parse_firmware_remaining() {
  local json="$1"
  local used
  
  used=$(echo "$json" | jq -r '.used // empty' 2>/dev/null)
  [[ -z "$used" || "$used" == "null" ]] && return 1
  
  # Convert used ratio (0-1) to remaining percent
  local remaining
  remaining=$(echo "scale=0; (1 - $used) * 100 / 1" | bc)
  (( remaining < 0 )) && remaining=0
  (( remaining > 100 )) && remaining=100
  
  echo "$remaining"
}

# Get Firmware quota string
get_firmware_quota() {
  # Check cache first
  if is_firmware_cache_valid; then
    local cached
    cached=$(read_firmware_cache)
    local remaining
    remaining=$(parse_firmware_remaining "$cached")
    if [[ -n "$remaining" ]]; then
      echo "F:${remaining}%"
      return 0
    fi
  fi
  
  # Cache miss - fetch from API
  local usage_json
  usage_json=$(fetch_firmware_usage)
  
  if [[ -z "$usage_json" ]] || echo "$usage_json" | jq -e '.error // .detail' >/dev/null 2>&1; then
    echo "$FIRMWARE_ERROR_OUTPUT"
    return 0
  fi
  
  # Cache valid response
  write_firmware_cache "$usage_json"
  
  # Parse and format
  local remaining
  remaining=$(parse_firmware_remaining "$usage_json")
  if [[ -n "$remaining" ]]; then
    echo "F:${remaining}%"
  else
    echo "$FIRMWARE_ERROR_OUTPUT"
  fi
}

# Get active username (fast, always reads from config)
get_username() {
  local config_file
  config_file=$(find_config_file) || return 1
  local email
  email=$(jq -r '.accounts[.activeIndex].email // empty' "$config_file" 2>/dev/null)
  echo "${email%%@*}"
}

# Main flow
main() {
  local quota_json
  local username
  local codex_quota
  username=$(get_username)
  
  # Get Codex quota (independent, can run even if antigravity fails)
  codex_quota=$(get_codex_quota)
  
  # Get Firmware quota
  firmware_quota=$(get_firmware_quota)
  
  # Check antigravity cache first
  if is_ag_cache_valid; then
    quota_json=$(read_ag_cache)
    echo "${username} $(format_quota "$quota_json") • ${codex_quota} • ${firmware_quota}"
    return 0
  fi
  
  # Cache miss - fetch from API
  
  # Find config file
  local config_file
  config_file=$(find_config_file)
  if [[ -z "$config_file" ]]; then
    echo "${username:-} $AG_ERROR_OUTPUT • ${codex_quota} • ${firmware_quota}"
    exit 0
  fi
  
  # Get active account
  local account
  account=$(get_active_account "$config_file")
  if [[ -z "$account" || "$account" == "null" ]]; then
    echo "${username:-} $AG_ERROR_OUTPUT • ${codex_quota} • ${firmware_quota}"
    exit 0
  fi
  
  # Get refresh token
  local refresh_token
  refresh_token=$(echo "$account" | jq -r '.refreshToken // empty' 2>/dev/null)
  if [[ -z "$refresh_token" ]]; then
    echo "${username:-} $AG_ERROR_OUTPUT • ${codex_quota} • ${firmware_quota}"
    exit 0
  fi
  
  # Refresh access token
  local access_token
  access_token=$(refresh_access_token "$refresh_token")
  if [[ -z "$access_token" ]]; then
    echo "${username:-} $AG_ERROR_OUTPUT • ${codex_quota} • ${firmware_quota}"
    exit 0
  fi
  
  # Get project ID
  local project_id
  project_id=$(get_project_id "$account" "$access_token")
  # Note: project_id can be empty - API might work without it
  
  # Fetch quota
  quota_json=$(fetch_quota "$access_token" "$project_id")
  
  if [[ -z "$quota_json" ]] || echo "$quota_json" | jq -e '.error' >/dev/null 2>&1; then
    echo "${username:-} $AG_ERROR_OUTPUT • ${codex_quota} • ${firmware_quota}"
    exit 0
  fi
  
  # Cache valid response
  write_ag_cache "$quota_json"
  
  echo "${username} $(format_quota "$quota_json") • ${codex_quota} • ${firmware_quota}"
}

main "$@"