diff --git a/bin/omarchy b/bin/omarchy index 2bace6df..35b8094a 100755 --- a/bin/omarchy +++ b/bin/omarchy @@ -26,7 +26,7 @@ declare -A ROUTE_IS_ALIAS declare -A BINARY_TO_KEY declare -A GROUP_DESCRIPTIONS -GROUP_DESCRIPTIONS[agent]="AI coding agent usage data" +GROUP_DESCRIPTIONS[agent]="AI coding agents, their usage, and subscription accounts" GROUP_DESCRIPTIONS[ascii]="Text drawn as ASCII art" GROUP_DESCRIPTIONS[audio]="Audio input and output controls" GROUP_DESCRIPTIONS[bar]="Omarchy shell bar layout and settings" @@ -371,6 +371,20 @@ command_requires_args() { [[ -n $required ]] } +# A command whose documented arguments are all flags can't take a word, so a +# word after it names something else: `omarchy agent account` is the account +# commands, not an argument to omarchy-agent. +command_takes_only_flags() { + local key="$1" + local args="${COMMAND_ARGS[$key]}" + local token="" + + [[ -n $args ]] || return 1 + for token in $args; do + [[ $token == -* || $token == "[-"* ]] || return 1 + done +} + load_group_commands() { local group="$1" local file="" @@ -956,6 +970,17 @@ dispatch_fast_or_help() { remaining=("${args[@]:DIRECT_RESOLVED_COUNT}") binary_path="$OMARCHY_BIN_DIR/$DIRECT_RESOLVED_BINARY" + if (( ${#remaining[@]} == 1 || ${#remaining[@]} == 2 )) && [[ ${remaining[0]} != -* ]] && + { (( ${#remaining[@]} == 1 )) || remaining_has_help_flag "${remaining[1]}"; } && + group_has_child_commands "${DIRECT_RESOLVED_BINARY#omarchy-}-${remaining[0]}"; then + load_command_by_binary "$DIRECT_RESOLVED_BINARY" + key="${BINARY_TO_KEY[$DIRECT_RESOLVED_BINARY]}" + if [[ -n $key ]] && command_takes_only_flags "$key"; then + load_commands + show_prefix_help "$DIRECT_RESOLVED_ROUTE ${remaining[0]}" && return 0 + fi + fi + if remaining_has_help_flag "${remaining[@]}"; then if ! load_command_by_binary "$DIRECT_RESOLVED_BINARY"; then echo "Binary is missing or not executable: $DIRECT_RESOLVED_BINARY" >&2 diff --git a/bin/omarchy-agent b/bin/omarchy-agent index e35bcf04..55bd11ce 100755 --- a/bin/omarchy-agent +++ b/bin/omarchy-agent @@ -136,6 +136,18 @@ pi) ;; esac +# Claude and Codex start as the active subscription account, the same as the +# shell functions in default/bash/fns/agent-accounts. The home travels in the +# command itself, since a terminal launched through uwsm doesn't inherit this +# process's environment. +if [[ $agent == "claude" && -z ${CLAUDE_CONFIG_DIR:-} ]]; then + home=$(omarchy-agent-account-home claude) + [[ -n $home ]] && command=(env "CLAUDE_CONFIG_DIR=$home" "${command[@]}") +elif [[ $agent == "codex" && -z ${CODEX_HOME:-} ]]; then + home=$(omarchy-agent-account-home codex) + [[ -n $home ]] && command=(env "CODEX_HOME=$home" "${command[@]}") +fi + if [[ $inline == "true" ]]; then exec "${command[@]}" else diff --git a/bin/omarchy-agent-account-add b/bin/omarchy-agent-account-add new file mode 100755 index 00000000..0bd1b5f7 --- /dev/null +++ b/bin/omarchy-agent-account-add @@ -0,0 +1,240 @@ +#!/bin/bash + +# omarchy:summary=Sign in to a Claude, Codex, or Grok subscription +# omarchy:args=[--check] [--events] [--reauth ] [claude|codex|grok] [label] +# omarchy:examples=omarchy agent account add claude | omarchy agent account add Work | omarchy agent account add codex Side + +set -uo pipefail + +# --check prints, per provider, whether adding one now would be its first +# sign-in, an additional account, or not possible yet. --events is for the +# agents panel: progress arrives as `@@omarchy ` lines mixed +# into the CLI's own output (which carries codes the panel shows), nothing +# waits on a terminal, and the result also lands as a notification in case +# the panel has closed. +# --reauth signs an account that's already here in again, in its +# own home, for when its sign-in has lapsed. +check=false +events=false +reauth="" +while [[ ${1:-} == --* ]]; do + case "$1" in + --check) check=true ;; + --events) events=true ;; + --reauth) + reauth="${2:-}" + shift + ;; + *) break ;; + esac + shift +done + +# Whether a provider's own home (~/.claude, ~/.codex, ~/.grok) is signed in. +primary_signed_in() { + if [[ $1 == "grok" ]]; then + [[ -s $HOME/.grok/auth.json ]] + else + omarchy-agent-account-state list "$1" | jq -e '.[0].accounts | any(.primary and .signedIn)' >/dev/null + fi +} + +if [[ $check == "true" ]]; then + for p in claude codex grok; do + if ! primary_signed_in "$p"; then + echo "$p first" + elif [[ $p == "grok" ]]; then + echo "$p unsupported" + else + echo "$p additional" + fi + done + exit 0 +fi + +# The provider can be left out when it's your default agent: +# `omarchy agent account add Work`. +if [[ ${1:-} != "claude" && ${1:-} != "codex" && ${1:-} != "grok" ]]; then + default=$(omarchy-default-agent) + [[ $default == "claude" || $default == "codex" || $default == "grok" ]] && set -- "$default" "$@" +fi + +case "${1:-}" in +claude) provider="claude"; name="Claude"; site="claude.ai"; package="claude" ;; +codex) provider="codex"; name="Codex"; site="chatgpt.com"; package="codex" ;; +grok) provider="grok"; name="Grok"; site="x.ai"; package="npm:@xai-official/grok" ;; +*) + echo "Usage: omarchy-agent-account-add [--check] [--events] [--reauth ] [claude|codex|grok] [label]" >&2 + exit 1 + ;; +esac +label="${2:-}" + +say() { + if [[ $events == "true" ]]; then + printf '@@omarchy status %s\n' "$*" + else + echo "$*" + fi +} + +finish() { + if [[ $events == "true" ]]; then + printf '@@omarchy done %s\n' "$1" + omarchy-notification-send -g 󰚩 "$1" + else + echo + echo "$1" + fi +} + +fail() { + if [[ $events == "true" ]]; then + printf '@@omarchy error %s\n' "$1" + omarchy-notification-send -g 󰚩 -u critical "Couldn't add the $name account" "$1" + else + echo "$1" >&2 + fi + exit 1 +} + +# Cancelling from the panel stops the CLI's login too, and leaves no +# half-made home behind. +home="" +browser="" +cleanup() { + [[ -n $browser ]] && rm -f "$browser" + [[ -n $home && -d $home && $home == */.pending/* ]] && rm -rf "$home" +} +trap cleanup EXIT +trap 'pkill -TERM -P $$; exit 130' TERM INT + +# Bash holds a trap until a foreground command finishes, and a login waits on +# the browser for as long as it takes. Running it in the background behind a +# `wait` lets Cancel stop it at once. It keeps this script's stdin, where the +# panel sends a pasted code. +interruptible() { + "$@" <&0 & + wait $! +} + +if omarchy-cmd-missing "$provider"; then + say "Installing the $name CLI…" + if [[ $events == "true" ]]; then + mise use -g "$package" >/dev/null 2>&1 || fail "Installing the $name CLI didn't finish." + else + mise use -g "$package" || exit 1 + fi +fi + +# Your main browser is almost certainly signed in to the account you already +# have, and a login would quietly reuse it. The CLIs open their login page +# through $BROWSER, so this points it at a private window instead. +private_browser() { + browser=$(mktemp "${XDG_RUNTIME_DIR:-/tmp}/omarchy-account-browser.XXXXXX") + printf '#!/bin/bash\nexec omarchy-launch-browser --private "$1"\n' >"$browser" + chmod +x "$browser" +} + +case "$provider" in +claude) login_command=(claude auth login) ;; +codex) login_command=(codex login) ;; +grok) login_command=(grok login) ;; +esac + +# The CLI's own home, whatever this shell has pointed it at. +login() { + interruptible env -u CLAUDE_CONFIG_DIR -u CODEX_HOME -u GROK_HOME "${login_command[@]}" +} + +# Signing an account that's already here in again, where it lives: the +# primary home through the normal browser, an added one through a private +# window. Its home is never recreated or removed. +if [[ -n $reauth ]]; then + if [[ $provider == "grok" || $reauth == ":primary" ]]; then + account_home="$HOME/.$provider" + else + account_home=$(omarchy-agent-account-state homes "$provider" | awk -F'\t' -v id="$reauth" '$1 == id { print $2 }') + fi + [[ -n $account_home ]] || fail "No $name account named $reauth." + + if [[ $account_home == "$HOME/.$provider" ]]; then + say "Sign in to $name again in the browser window that opens." + login || fail "The $name sign-in didn't finish." + else + say "Sign in again as this account in the private window that opens." + private_browser + if [[ $provider == "claude" ]]; then + interruptible env BROWSER="$browser" CLAUDE_CONFIG_DIR="$account_home" claude auth login || fail "The $name sign-in didn't finish." + else + interruptible env BROWSER="$browser" CODEX_HOME="$account_home" codex login || fail "The $name sign-in didn't finish." + fi + fi + + omarchy-agent-usage-update --force "$provider" >/dev/null 2>&1 & + finish "Signed in to $name again." + exit 0 +fi + +# The first subscription signs in where the CLI always looks, through your +# normal browser: nothing is signed in there yet that it could be mistaken for. +if ! primary_signed_in "$provider"; then + say "Sign in to $name in the browser window that opens." + if login && primary_signed_in "$provider"; then + # The first agent signed in on a machine with none chosen becomes the + # default, so `omarchy agent` and the panel's starter prompts have one to + # run. A default someone already picked is never touched. + if [[ -z $(omarchy-default-agent) ]]; then + mkdir -p "$HOME/.config/omarchy/defaults" + printf '%s\n' "$provider" >"$HOME/.config/omarchy/defaults/agent" + fi + omarchy-agent-usage-update --force "$provider" >/dev/null 2>&1 & + finish "Signed in to $name." + else + fail "The $name sign-in didn't finish." + fi +elif [[ $provider == "grok" ]]; then + fail "Grok is already signed in, and a second Grok account isn't supported yet." +else + if [[ -z $label && $events == "false" && -t 0 ]]; then + label=$(gum input --placeholder "Work" --prompt "Name this $name account: ") || exit 130 + fi + + if [[ $events == "true" ]]; then + say "Sign in as the account you're adding in the private window that opens." + else + cat <&1); then + home="" + added=$(jq -r '.label + (if .email != "" then " (" + .email + ")" else "" end)' <<<"$account") + omarchy-agent-usage-update --force "$provider" >/dev/null 2>&1 & + if [[ $events == "true" ]]; then + finish "Added $added." + else + finish "Added $added." + echo "Switch to it with: omarchy agent account use $provider $(jq -r .id <<<"$account")" + fi + else + home="" + fail "$account" + fi +fi diff --git a/bin/omarchy-agent-account-home b/bin/omarchy-agent-account-home new file mode 100755 index 00000000..5b1e387a --- /dev/null +++ b/bin/omarchy-agent-account-home @@ -0,0 +1,18 @@ +#!/bin/bash + +# omarchy:summary=Print the config home of the active Claude or Codex account +# omarchy:args= +# omarchy:hidden=true + +# Prints nothing while the primary account (~/.claude, ~/.codex) is active, so +# a launch only sets CLAUDE_CONFIG_DIR / CODEX_HOME when there's somewhere +# else to point it. Without a registry, no Python starts at all: this runs on +# every `claude` and `codex` typed at a prompt. +registry="${XDG_STATE_HOME:-$HOME/.local/state}/omarchy/agents/accounts/${1:-}.json" + +if [[ ${1:-} != "claude" && ${1:-} != "codex" ]]; then + echo "Usage: omarchy-agent-account-home " >&2 + exit 1 +elif [[ -f $registry ]]; then + exec omarchy-agent-account-state home "$1" +fi diff --git a/bin/omarchy-agent-account-list b/bin/omarchy-agent-account-list new file mode 100755 index 00000000..41b1203a --- /dev/null +++ b/bin/omarchy-agent-account-list @@ -0,0 +1,58 @@ +#!/bin/bash + +# omarchy:summary=List Claude and Codex subscription accounts and their limits +# omarchy:args=[claude|codex] [--json] +# omarchy:examples=omarchy agent account list | omarchy agent account list claude --json + +set -uo pipefail + +provider="" +json=false +for arg in "$@"; do + case "$arg" in + claude | codex) provider="$arg" ;; + --json) json=true ;; + *) + echo "Usage: omarchy-agent-account-list [claude|codex] [--json]" >&2 + exit 1 + ;; + esac +done + +usage_dir="${XDG_STATE_HOME:-$HOME/.local/state}/omarchy/agents/usage" +accounts=$(omarchy-agent-account-state list $provider) || exit 1 + +# Join each account with the limits its provider's usage record last saw. With +# a single account the collectors keep its limits at the record's top level. +accounts=$(jq -c --arg dir "$usage_dir" ' + map(. as $p | .accounts |= map(. + {limits: []}))' <<<"$accounts") +for p in claude codex; do + record="$usage_dir/$p.json" + [[ -f $record ]] || continue + accounts=$(jq -c --arg p "$p" --slurpfile record "$record" ' + map(if .provider == $p then + .accounts |= map(. as $a | . + {limits: ( + if ($record[0].accounts // []) | length > 0 then + [$record[0].accounts[] | select(.id == $a.id) | .limits] | first // [] + elif $a.primary then + $record[0].limits // [] + else + [] + end)}) + else . end)' <<<"$accounts") +done + +if [[ $json == "true" ]]; then + echo "$accounts" +else + # A provider nobody has signed in to has nothing to list. + for (( i = 0; i < $(jq length <<<"$accounts"); i++ )); do + jq -e ".[$i].accounts | any(.signedIn)" <<<"$accounts" >/dev/null || continue + jq -r ".[$i] | \"\(.name) — \(.switch) switching at \(.threshold)%\"" <<<"$accounts" + jq -r ".[$i].accounts[] | + \"\(if .active then \"*\" else \" \" end)\t\(.id)\t\(.label)\t\(.plan)\t\(.email)\t\" + + ([.limits[]? | select(.percent != null) | \"\(.label): \((.percent * 100) | round)%\"] | join(\", \"))" <<<"$accounts" | + column -t -s $'\t' | sed 's/^/ /' + echo + done +fi diff --git a/bin/omarchy-agent-account-mode b/bin/omarchy-agent-account-mode new file mode 100755 index 00000000..a10a7fd8 --- /dev/null +++ b/bin/omarchy-agent-account-mode @@ -0,0 +1,30 @@ +#!/bin/bash + +# omarchy:summary=Switch Claude or Codex accounts automatically near a limit, or only notify +# omarchy:args=[claude|codex] [manual|auto] [threshold] +# omarchy:examples=omarchy agent account mode auto | omarchy agent account mode auto 90 | omarchy agent account mode codex manual + +set -uo pipefail + +# The provider can be left out when it's your default agent: +# `omarchy agent account use work`. +if [[ ${1:-} != "claude" && ${1:-} != "codex" ]]; then + default=$(omarchy-default-agent) + [[ $default == "claude" || $default == "codex" ]] && set -- "$default" "$@" +fi + +if [[ ${1:-} != "claude" && ${1:-} != "codex" ]]; then + echo "Usage: omarchy-agent-account-mode [claude|codex] [manual|auto] [threshold]" >&2 + exit 1 +fi + +if [[ -n ${2:-} ]]; then + omarchy-agent-account-state mode "$@" || exit 1 +fi + +omarchy-agent-account-state registry "$1" | jq -r ' + if .switch == "auto" then + "\(.name) switches new sessions to the account with the most headroom at \(.threshold)%." + else + "\(.name) notifies at \(.threshold)% and switches when you say so." + end' diff --git a/bin/omarchy-agent-account-remove b/bin/omarchy-agent-account-remove new file mode 100755 index 00000000..7abc9d76 --- /dev/null +++ b/bin/omarchy-agent-account-remove @@ -0,0 +1,30 @@ +#!/bin/bash + +# omarchy:summary=Forget an added Claude or Codex account +# omarchy:args=[claude|codex] +# omarchy:examples=omarchy agent account remove work | omarchy agent account remove codex side + +set -uo pipefail + +# The provider can be left out when it's your default agent: +# `omarchy agent account use work`. +if [[ ${1:-} != "claude" && ${1:-} != "codex" ]]; then + default=$(omarchy-default-agent) + [[ $default == "claude" || $default == "codex" ]] && set -- "$default" "$@" +fi + +if [[ ${1:-} != "claude" && ${1:-} != "codex" ]] || [[ -z ${2:-} ]]; then + echo "Usage: omarchy-agent-account-remove [claude|codex] " >&2 + exit 1 +fi + +# Removing only deletes the account's home here. It never logs out, which +# would revoke the sign-in for every other copy of it too. +if [[ -t 0 ]] && ! gum confirm "Forget the $1 account $2? You'd have to sign in again to add it back."; then + exit 130 +fi + +account=$(omarchy-agent-account-state remove "$1" "$2") || exit 1 +echo "Forgot $(jq -r .label <<<"$account")." +# The panel reads which accounts exist from the usage record. +omarchy-agent-usage-update --limits-only "$1" >/dev/null 2>&1 & diff --git a/bin/omarchy-agent-account-rename b/bin/omarchy-agent-account-rename new file mode 100755 index 00000000..1d5f1bdc --- /dev/null +++ b/bin/omarchy-agent-account-rename @@ -0,0 +1,24 @@ +#!/bin/bash + +# omarchy:summary=Rename a Claude or Codex subscription account +# omarchy:args=[claude|codex] +# omarchy:examples=omarchy agent account rename primary Hey | omarchy agent account rename codex side Side project + +set -uo pipefail + +# The provider can be left out when it's your default agent: +# `omarchy agent account rename primary Hey`. +if [[ ${1:-} != "claude" && ${1:-} != "codex" ]]; then + default=$(omarchy-default-agent) + [[ $default == "claude" || $default == "codex" ]] && set -- "$default" "$@" +fi + +if [[ ${1:-} != "claude" && ${1:-} != "codex" ]] || (( $# < 3 )); then + echo "Usage: omarchy-agent-account-rename [claude|codex] " >&2 + exit 1 +fi + +account=$(omarchy-agent-account-state rename "$@") || exit 1 +echo "Renamed to $(jq -r .label <<<"$account"). Switch to it with: omarchy agent account use $(jq -r .id <<<"$account")" +# The panel reads account names from the usage record. +omarchy-agent-usage-update --limits-only "$1" >/dev/null 2>&1 & diff --git a/bin/omarchy-agent-account-state b/bin/omarchy-agent-account-state new file mode 100755 index 00000000..2a82a0d8 --- /dev/null +++ b/bin/omarchy-agent-account-state @@ -0,0 +1,758 @@ +#!/usr/bin/python3 +# omarchy:summary=Read and change the registry of Claude and Codex subscription accounts +# omarchy:args= [provider] [args...] +# omarchy:hidden=true +"""The account registry behind `omarchy agent account`. + +An account is a Claude Code or Codex login Omarchy keeps a home for. The +login you already have stays where the CLI put it (~/.claude, ~/.codex) and +is the primary account; added accounts get a home of their own under +~/.local/state/omarchy/agents/accounts/// that holds only that +account's credentials and links everything else back to the primary home, so +conversation history, settings and skills stay one set across accounts. + +Each provider's registry is one JSON file beside those homes. It records which +account is active — the one new sessions start as — how switching happens, and +display-safe identity. Tokens never enter it and never leave this process. +""" + +from __future__ import annotations + +import base64 +import datetime as dt +import fcntl +import json +import os +import re +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Any + +PROVIDERS = { + "claude": { + "name": "Claude", + "home_env": "CLAUDE_CONFIG_DIR", + # Everything a session reads or writes that isn't tied to the login. A + # missing directory is created in the primary home first; a missing file + # is linked anyway, so whichever account writes it first writes it there. + "shared_dirs": ["projects", "skills", "agents", "commands", "plugins", "hooks", "themes", "output-styles", "file-history", "todos", "plans"], + "shared_files": ["settings.json", "CLAUDE.md", "history.jsonl", "keybindings.json"], + # Keys copied from the primary's ~/.claude.json into a new account's, so a + # second login keeps MCP servers, folder trust and onboarding choices. + "carried_keys": ["mcpServers", "projects", "theme", "hasCompletedOnboarding", "lastOnboardingVersion", "editorMode"], + }, + "codex": { + "name": "Codex", + "home_env": "CODEX_HOME", + "shared_dirs": ["sessions", "archived_sessions", "prompts", "skills", "rules", "plugins"], + "shared_files": ["config.toml", "AGENTS.md", "history.jsonl", "hooks.json"], + "carried_keys": [], + }, +} + +DEFAULT_THRESHOLD = 95 +# Ids an account can't take because something else already answers to them: +# `use next` cycles, a provider name is read as the provider when it comes +# first, and the menu's Add Account row is `.add`. +RESERVED_IDS = {"next", "add", "claude", "codex", "primary"} +SWITCH_MODES = ("manual", "auto") + + +class AccountError(Exception): + pass + + +# ------------------------------------------------------------------- paths + + +def state_root() -> Path: + return Path(os.environ.get("XDG_STATE_HOME") or Path.home() / ".local" / "state") / "omarchy" / "agents" + + +def accounts_root() -> Path: + return state_root() / "accounts" + + +def registry_path(provider: str) -> Path: + return accounts_root() / f"{provider}.json" + + +def usage_record_path(provider: str) -> Path: + return state_root() / "usage" / f"{provider}.json" + + +# The primary home is where the CLI keeps a login when nothing tells it +# otherwise. It is deliberately not read from CLAUDE_CONFIG_DIR / CODEX_HOME: +# those are how an added account is selected, so honoring them here would make +# whichever account launched this command look like the primary. +def primary_home(provider: str) -> Path: + return Path.home() / (".claude" if provider == "claude" else ".codex") + + +def account_home(provider: str, account: dict[str, Any]) -> Path: + home = str(account.get("home") or "") + return Path(home) if home else primary_home(provider) + + +# Claude keeps its account file beside the config dir for the default home +# (~/.claude.json) and inside it for any other. +def claude_account_file(home: Path) -> Path: + if home == primary_home("claude"): + return Path.home() / ".claude.json" + return home / ".claude.json" + + +# --------------------------------------------------------------- identity + + +def read_json(path: Path) -> Any: + try: + return json.loads(path.read_text(encoding="utf-8")) + except Exception: + return None + + +def plan_label(tier: str, subscription: str) -> str: + match = re.search(r"max_(\d+x)", tier or "", re.IGNORECASE) + if match: + return "Max " + match.group(1) + if subscription: + return subscription[0].upper() + subscription[1:] + return "" + + +def jwt_claims(token: str) -> dict[str, Any]: + try: + payload = token.split(".")[1] + payload += "=" * (-len(payload) % 4) + claims = json.loads(base64.urlsafe_b64decode(payload.encode("ascii"))) + return claims if isinstance(claims, dict) else {} + except Exception: + return {} + + +# Who a home is signed in as, from files the CLI wrote at login. Only +# display-safe fields come out; an empty accountId means nobody is signed in. +def identity(provider: str, home: Path) -> dict[str, str]: + if provider == "claude": + account = (read_json(claude_account_file(home)) or {}).get("oauthAccount") or {} + login = (read_json(home / ".credentials.json") or {}).get("claudeAiOauth") or {} + return { + "accountId": str(account.get("accountUuid") or ""), + "email": str(account.get("emailAddress") or ""), + "org": str(account.get("organizationName") or ""), + "plan": plan_label(str(login.get("rateLimitTier") or ""), str(login.get("subscriptionType") or "")), + } + + auth = read_json(home / "auth.json") or {} + tokens = auth.get("tokens") or {} + claims = jwt_claims(str(tokens.get("id_token") or "")) + openai = claims.get("https://api.openai.com/auth") or {} + plan = str(openai.get("chatgpt_plan_type") or "") + return { + "accountId": str(tokens.get("account_id") or openai.get("chatgpt_account_id") or ""), + "email": str(claims.get("email") or ""), + "org": "", + "plan": plan[0].upper() + plan[1:] if plan else "", + } + + +# ---------------------------------------------------------------- registry + + +def slug(label: str) -> str: + return re.sub(r"[^a-z0-9]+", "-", label.lower()).strip("-") + + +def empty_registry(provider: str) -> dict[str, Any]: + return {"active": "main", "switch": "manual", "threshold": DEFAULT_THRESHOLD, "alert": "", "accounts": []} + + +def primary_entry(provider: str) -> dict[str, Any]: + entry = {"id": "main", "label": "Main", "home": "", "primary": True} + entry.update(identity(provider, primary_home(provider))) + return entry + + +def normalized(provider: str, data: Any) -> dict[str, Any]: + registry = empty_registry(provider) + if isinstance(data, dict): + registry.update({key: data[key] for key in registry if key in data}) + accounts = [a for a in registry["accounts"] if isinstance(a, dict) and a.get("id")] + if not any(a.get("primary") for a in accounts): + accounts.insert(0, primary_entry(provider)) + registry["accounts"] = accounts + if registry["switch"] not in SWITCH_MODES: + registry["switch"] = "manual" + try: + registry["threshold"] = max(50, min(100, int(registry["threshold"]))) + except Exception: + registry["threshold"] = DEFAULT_THRESHOLD + if not find(registry, registry["active"]): + registry["active"] = accounts[0]["id"] + return registry + + +def load(provider: str) -> dict[str, Any]: + return normalized(provider, read_json(registry_path(provider))) + + +def save(provider: str, registry: dict[str, Any]) -> None: + path = registry_path(provider) + path.parent.mkdir(parents=True, exist_ok=True) + os.chmod(path.parent, 0o700) + fd, tmp = tempfile.mkstemp(prefix=path.name + ".", dir=str(path.parent)) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(registry, handle, indent=2) + handle.write("\n") + os.chmod(tmp, 0o600) + os.replace(tmp, path) + except Exception: + Path(tmp).unlink(missing_ok=True) + raise + + +# Every change is read-modify-write under one lock per provider, so the panel, +# the autoswitch run after each usage update, and a command typed at a prompt +# can't interleave and drop each other's edits. +class Locked: + def __init__(self, provider: str): + self.provider = provider + + def __enter__(self) -> dict[str, Any]: + accounts_root().mkdir(parents=True, exist_ok=True) + self.handle = open(accounts_root() / f".{self.provider}.lock", "w") + fcntl.flock(self.handle, fcntl.LOCK_EX) + self.registry = load(self.provider) + return self.registry + + def __exit__(self, kind, value, traceback) -> None: + try: + if kind is None: + save(self.provider, self.registry) + finally: + fcntl.flock(self.handle, fcntl.LOCK_UN) + self.handle.close() + + +def find(registry: dict[str, Any], account_id: str) -> dict[str, Any] | None: + for account in registry["accounts"]: + if account.get("id") == account_id: + return account + # `primary` names the login that was there first, whatever it's called now. + if account_id == "primary": + for account in registry["accounts"]: + if account.get("primary"): + return account + return None + + +def active(registry: dict[str, Any]) -> dict[str, Any]: + return find(registry, registry["active"]) or registry["accounts"][0] + + +# ------------------------------------------------------------------- homes + + +# Lay the links that make an added account share everything but its login +# with the primary home. Safe to repeat: an existing link or real file is left +# as it is, so nothing an account wrote for itself is ever replaced. +def link_shared(provider: str, home: Path) -> None: + spec = PROVIDERS[provider] + primary = primary_home(provider) + home.mkdir(parents=True, exist_ok=True) + os.chmod(home, 0o700) + for name in spec["shared_dirs"]: + (primary / name).mkdir(parents=True, exist_ok=True) + for name in spec["shared_dirs"] + spec["shared_files"]: + target = home / name + if target.exists() or target.is_symlink(): + continue + target.symlink_to(primary / name) + + +def carry_settings(provider: str, home: Path) -> None: + keys = PROVIDERS[provider]["carried_keys"] + if not keys: + return + source = read_json(claude_account_file(primary_home(provider))) or {} + path = claude_account_file(home) + data = read_json(path) or {} + for key in keys: + if key in source and key not in data: + data[key] = source[key] + fd, tmp = tempfile.mkstemp(prefix=path.name + ".", dir=str(path.parent)) + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(data, handle, indent=2) + os.chmod(tmp, 0o600) + os.replace(tmp, path) + + +# A login in progress gets a scratch home beside the real ones, so an +# abandoned or duplicate login never shows up as an account. +def begin(provider: str) -> Path: + pending = accounts_root() / provider / ".pending" + pending.mkdir(parents=True, exist_ok=True) + os.chmod(pending.parent, 0o700) + home = Path(tempfile.mkdtemp(prefix="login-", dir=str(pending))) + link_shared(provider, home) + return home + + +def register(provider: str, label: str, pending: Path) -> dict[str, Any]: + pending = pending.resolve() + if pending.parent != (accounts_root() / provider / ".pending").resolve(): + raise AccountError(f"{pending} is not a pending {PROVIDERS[provider]['name']} login") + + found = identity(provider, pending) + if not found["accountId"]: + shutil.rmtree(pending, ignore_errors=True) + raise AccountError("The login didn't finish, so no account was added.") + + with Locked(provider) as registry: + # Each home is asked who it is now, not trusted from when it was added: a + # `claude auth login` run by hand in ~/.claude changes the primary. + for account in registry["accounts"]: + account.update(identity(provider, account_home(provider, account))) + for account in registry["accounts"]: + if account.get("accountId") == found["accountId"]: + shutil.rmtree(pending, ignore_errors=True) + raise AccountError( + f"That's {account['label']} ({found['email'] or 'already added'}). " + "Try again, and sign in as the other account in the private window." + ) + + label = label.strip() or found["email"].split("@")[0] or "Account" + account_id = unique_id(provider, registry, label) + + # Everything that can fail happens while the login is still pending, where + # the add command cleans it up, or is undone: a home that isn't in the + # registry would keep a sign-in nothing can manage. + carry_settings(provider, pending) + home = accounts_root() / provider / account_id + pending.rename(home) + entry = {"id": account_id, "label": label, "home": str(home), "primary": False} + entry.update(found) + registry["accounts"].append(entry) + try: + save(provider, registry) + except Exception: + registry["accounts"].pop() + home.rename(pending) + raise + return entry + + +# The id an account answers to, from its label. It can't clash with another +# account, a routing keyword, or a home left under that name by an account +# that has since been renamed. +def unique_id(provider: str, registry: dict[str, Any], label: str, keep: dict[str, Any] | None = None) -> str: + base = slug(label) or "account" + account_id, n = base, 2 + while True: + taken = find(registry, account_id) + if account_id not in RESERVED_IDS and (taken is None or taken is keep): + if keep is not None or not (accounts_root() / provider / account_id).exists(): + return account_id + account_id, n = f"{base}-{n}", n + 1 + + +# A new label, and the id that follows from it. The home stays where it is: +# sessions already running in it hold that path. +def rename(provider: str, target: str, label: str) -> dict[str, Any]: + label = label.strip() + if not label: + raise AccountError("Give the account a name") + with Locked(provider) as registry: + account = find(registry, target) + if not account: + raise AccountError(f"No {PROVIDERS[provider]['name']} account named {target}") + new_id = unique_id(provider, registry, label, keep=account) + if not account.get("primary") and not account.get("home"): + account["home"] = str(account_home(provider, account)) + if registry["active"] == account["id"]: + registry["active"] = new_id + registry["alert"] = "" + account["id"] = new_id + account["label"] = label + return dict(account) + + +# Save who each home is signed in as now. A login run in a home, through the +# panel or the CLI, changes it, and the usage records take their names and +# emails from here. +def refresh(provider: str) -> None: + with Locked(provider) as registry: + for account in registry["accounts"]: + account.update(identity(provider, account_home(provider, account))) + + +def use(provider: str, target: str) -> tuple[dict[str, Any], dict[str, Any]]: + with Locked(provider) as registry: + previous = active(registry) + if target == "next": + ids = [a["id"] for a in registry["accounts"]] + target = ids[(ids.index(previous["id"]) + 1) % len(ids)] + account = find(registry, target) + if not account: + raise AccountError(f"No {PROVIDERS[provider]['name']} account named {target}") + registry["active"] = account["id"] + registry["alert"] = "" + if not account.get("primary"): + link_shared(provider, account_home(provider, account)) + return previous, account + + +# Whether a running session was started in this home. Removing it would pull +# the login out from under that session, which must never be touched. +def home_in_use(provider: str, home: Path) -> bool: + wanted = f"{PROVIDERS[provider]['home_env']}=".encode() + resolved = home.resolve() + for environ in Path("/proc").glob("[0-9]*/environ"): + try: + entries = environ.read_bytes().split(b"\0") + except OSError: + continue + for entry in entries: + if entry.startswith(wanted) and Path(entry[len(wanted):].decode(errors="replace")).resolve() == resolved: + return True + return False + + +def remove(provider: str, target: str) -> dict[str, Any]: + with Locked(provider) as registry: + account = find(registry, target) + if not account: + raise AccountError(f"No {PROVIDERS[provider]['name']} account named {target}") + if account.get("primary"): + raise AccountError(f"{account['label']} is the login in {primary_home(provider)}; it can't be removed here.") + home = account_home(provider, account) + if home_in_use(provider, home): + raise AccountError(f"A running {PROVIDERS[provider]['name']} session uses {account['label']}. Quit it first, then remove the account.") + # Only the account's own directory goes. Its links point into the primary + # home, and rmtree removes a link without following it. + if home.resolve().parent == (accounts_root() / provider).resolve(): + shutil.rmtree(home, ignore_errors=True) + registry["accounts"] = [a for a in registry["accounts"] if a["id"] != account["id"]] + if registry["active"] == account["id"]: + registry["active"] = registry["accounts"][0]["id"] + return account + + +def set_mode(provider: str, mode: str, threshold: str | None) -> dict[str, Any]: + if mode not in SWITCH_MODES: + raise AccountError("Switch mode is manual or auto") + with Locked(provider) as registry: + registry["switch"] = mode + if threshold is not None: + try: + registry["threshold"] = max(50, min(100, int(threshold))) + except ValueError: + raise AccountError("The threshold is a percentage between 50 and 100") + registry["alert"] = "" + return dict(registry) + + +# ------------------------------------------------------------- autoswitch + + +def window_open(limit: dict[str, Any], now: dt.datetime) -> bool: + raw = str(limit.get("resetsAt") or "") + if not raw: + return True + try: + resets_at = dt.datetime.fromisoformat(raw.replace("Z", "+00:00")) + except Exception: + return True + if resets_at.tzinfo is None: + resets_at = resets_at.replace(tzinfo=dt.timezone.utc) + return resets_at > now + + +# How spent an account is: its fullest window that hasn't reset yet, in +# percent. A window whose reset has passed counts as empty even when the +# numbers are stale, since that's exactly what a reset means. None when +# nothing is known at all. +def peak(limits: Any, now: dt.datetime) -> float | None: + if not isinstance(limits, list): + return None + known = [l for l in limits if isinstance(l, dict) and isinstance(l.get("percent"), (int, float))] + if not known: + return None + open_windows = [float(l["percent"]) * 100 for l in known if window_open(l, now)] + return max(open_windows) if open_windows else 0.0 + + +def next_reset(limits: Any, now: dt.datetime) -> dt.datetime | None: + times = [] + for limit in limits if isinstance(limits, list) else []: + try: + when = dt.datetime.fromisoformat(str(limit.get("resetsAt") or "").replace("Z", "+00:00")) + except Exception: + continue + if when.tzinfo is None: + when = when.replace(tzinfo=dt.timezone.utc) + if when > now: + times.append(when) + return min(times) if times else None + + +# When an account over the threshold drops back under it: once every window +# holding it there has reset, which is the latest of those resets. A 10% +# session resetting in an hour doesn't help an account whose weekly window is +# full for six more days. None when a blocking window never says. +def available_at(limits: Any, threshold: float, now: dt.datetime) -> dt.datetime | None: + latest = None + for limit in limits if isinstance(limits, list) else []: + if not isinstance(limit, dict) or not isinstance(limit.get("percent"), (int, float)): + continue + if float(limit["percent"]) * 100 < threshold or not window_open(limit, now): + continue + try: + when = dt.datetime.fromisoformat(str(limit.get("resetsAt") or "").replace("Z", "+00:00")) + except Exception: + return None + if when.tzinfo is None: + when = when.replace(tzinfo=dt.timezone.utc) + latest = when if latest is None or when > latest else latest + return latest + + +def duration(delta: dt.timedelta) -> str: + minutes = max(1, int(delta.total_seconds() // 60)) + hours, minutes = divmod(minutes, 60) + days, hours = divmod(hours, 24) + if days: + return f"{days}d {hours}h" + if hours: + return f"{hours}h {minutes}m" + return f"{minutes}m" + + +# The whole switching policy, with no I/O: given the registry and the per +# account limits from the usage record, what should happen. Returns None when +# nothing should, or {"action": "switch"|"notify"|"exhausted", ...}. +# +# Only the active account crossing the threshold starts anything, so auto +# mode never flaps back to an account that just reset. `alert` remembers the +# last thing said, so a shell restart or the next refresh doesn't repeat it. +# +# Numbers kept from an earlier check may be out of date, but a parked account's +# sign-in lapses within hours and can't be checked again until a session runs +# in it. So a stale account is still a candidate: among accounts with room, one +# checked just now wins, but room comes first. An account within 15 points of +# the threshold (where the panel starts checking more often) is near its limit +# however fresh its numbers are. +def decide(registry: dict[str, Any], usage: dict[str, Any], now: dt.datetime, stale: frozenset[str] = frozenset()) -> dict[str, Any] | None: + threshold = float(registry["threshold"]) + current = active(registry) + current_peak = peak(usage.get(current["id"]), now) + if current_peak is None or current_peak < threshold: + return {"action": "clear"} if registry.get("alert") else None + + candidates = [] + for account in registry["accounts"]: + if account["id"] == current["id"]: + continue + level = peak(usage.get(account["id"]), now) + if level is None or level >= threshold: + continue + reset = next_reset(usage.get(account["id"]), now) or dt.datetime.max.replace(tzinfo=dt.timezone.utc) + candidates.append((level >= threshold - 15, account["id"] in stale, level, reset, account)) + candidates.sort(key=lambda c: c[:4]) + + if candidates: + _, _, level, _, target = candidates[0] + action = "switch" if registry["switch"] == "auto" else "notify" + alert = f"{action}:{current['id']}:{target['id']}" + if registry.get("alert") == alert: + return None + return {"action": action, "alert": alert, "from": current, "to": target, "fromPeak": current_peak, "toPeak": level} + + alert = f"exhausted:{current['id']}" + if registry.get("alert") == alert or len(registry["accounts"]) < 2: + return None + # Only when every account is known to be over. One that's signed out or + # couldn't be checked might have room, so there's nothing honest to say. + if any(peak(usage.get(a["id"]), now) is None for a in registry["accounts"]): + return None + soonest = None + for account in registry["accounts"]: + reset = available_at(usage.get(account["id"]), threshold, now) + if reset and (soonest is None or reset < soonest[0]): + soonest = (reset, account) + return {"action": "exhausted", "alert": alert, "from": current, "fromPeak": current_peak, "soonest": soonest} + + +# Limits per account id. An account nobody is signed in to is left out, so it +# is never a candidate: its cached headroom is no use to a session that can't +# start. One whose token merely expired stays in, since a session refreshes it. +def usage_by_account(provider: str) -> tuple[dict[str, Any], frozenset[str]]: + record = read_json(usage_record_path(provider)) or {} + accounts = [ + a for a in record.get("accounts") or [] + if isinstance(a, dict) and a.get("id") and a.get("usageStatusText") != "Waiting for auth" + ] + return ( + {str(a["id"]): a.get("limits") for a in accounts}, + frozenset(str(a["id"]) for a in accounts if a.get("stale") is True), + ) + + +def notify(title: str, body: str, action: list[str] | None = None) -> None: + command = ["omarchy-notification-send", "--app-name", "Omarchy Agents", "-g", "󰚩", title, body] + if action: + command += ["--exec"] + action + subprocess.run(command, check=False) + + +# Returns True when the active account changed. +def autoswitch(provider: str, now: dt.datetime | None = None) -> bool: + now = now or dt.datetime.now(dt.timezone.utc) + name = PROVIDERS[provider]["name"] + with Locked(provider) as registry: + if len(registry["accounts"]) < 2: + return False + usage, stale = usage_by_account(provider) + decision = decide(registry, usage, now, stale) + if decision is None: + return False + if decision["action"] == "clear": + registry["alert"] = "" + return False + + registry["alert"] = decision["alert"] + source = decision["from"] + threshold = registry["threshold"] + if decision["action"] == "exhausted": + soonest = decision["soonest"] + when = f" {soonest[1]['label']} resets in {duration(soonest[0] - now)}." if soonest else "" + notify(f"All {name} accounts are over {threshold}%", f"Staying on {source['label']}.{when}") + return False + + target = decision["to"] + headline = f"{source['label']} is at {round(decision['fromPeak'])}% of a {name} limit" + if decision["action"] == "switch": + registry["active"] = target["id"] + if not target.get("primary"): + link_shared(provider, account_home(provider, target)) + notify(headline, f"New {name} sessions now use {target['label']} ({round(decision['toPeak'])}%). Running sessions stay on {source['label']}.") + return True + notify( + headline, + f"{target['label']} is at {round(decision['toPeak'])}%. Click to switch new sessions to it.", + ["omarchy-agent-account-use", provider, target["id"]], + ) + return False + + +# --------------------------------------------------------------------- CLI + + +def provider_arg(args: list[str]) -> str: + if not args or args[0] not in PROVIDERS: + raise AccountError("Name a provider: claude or codex") + return args[0] + + +def summary(provider: str, registry: dict[str, Any]) -> dict[str, Any]: + # Identity is re-read from each home rather than trusted from the registry: + # a `claude auth login` run by hand in a home changes who it is. + for account in registry["accounts"]: + account.update(identity(provider, account_home(provider, account))) + return { + "provider": provider, + "name": PROVIDERS[provider]["name"], + "active": registry["active"], + "switch": registry["switch"], + "threshold": registry["threshold"], + "accounts": [ + {key: account.get(key, "") for key in ("id", "label", "email", "org", "plan", "primary")} + | { + "home": str(account_home(provider, account)), + "active": account["id"] == registry["active"], + "signedIn": bool(account.get("accountId")), + } + for account in registry["accounts"] + ], + } + + +def main(argv: list[str]) -> int: + if not argv: + print(__doc__.strip().splitlines()[0], file=sys.stderr) + return 1 + command, args = argv[0], argv[1:] + try: + if command == "registry": + provider = provider_arg(args) + print(json.dumps(summary(provider, load(provider)))) + elif command == "list": + providers = [provider_arg(args)] if args else list(PROVIDERS) + print(json.dumps([summary(p, load(p)) for p in providers])) + elif command == "home": + # Empty for the primary, so launching it needs no environment at all. + provider = provider_arg(args) + registry = read_json(registry_path(provider)) + if registry: + account = active(normalized(provider, registry)) + if not account.get("primary"): + print(account_home(provider, account)) + elif command == "homes": + provider = provider_arg(args) + registry = load(provider) + for account in registry["accounts"]: + print("\t".join([account["id"], str(account_home(provider, account)), "1" if account["id"] == registry["active"] else "0"])) + elif command == "refresh": + for provider in ([provider_arg(args)] if args else list(PROVIDERS)): + if registry_path(provider).exists(): + refresh(provider) + elif command == "begin": + print(begin(provider_arg(args))) + elif command == "register": + provider = provider_arg(args) + if len(args) < 3: + raise AccountError("register