diff --git a/bin/omarchy-agent-usage-claude b/bin/omarchy-agent-usage-claude index 1a800631..d8443094 100755 --- a/bin/omarchy-agent-usage-claude +++ b/bin/omarchy-agent-usage-claude @@ -122,12 +122,209 @@ def empty_bucket() -> dict[str, int]: # ---------------------------------------------------------------- local scan -def scan_projects(projects_path: Path) -> dict[str, Any]: +# Transcripts are append-only, and most of the corpus never changes between +# ticks, so the scan keeps a per-file index: the unique usage records already +# parsed out of each file plus the byte offset they end at. A file whose size +# and mtime match the index is not opened; a file that grew is read from the +# stored offset; anything else (truncated, rewritten, replaced, unknown) is +# read from the start. Without this the whole history was re-parsed every +# refresh, and the cost grew for the life of the machine. + +INDEX_VERSION = 4 + +# How much of each end of what was read the index remembers a hash of, so a +# file replaced or rewritten with more than it had is not mistaken for an +# append. +INDEX_HEAD_BYTES = 4096 + + +# Records hold local days, so an index is only true in the timezone that +# wrote it; after a change every transcript is read again. +def index_zone() -> list[Any]: + return [os.environ.get("TZ"), list(time.tzname), time.timezone, time.altzone] + + +def load_index(index_file: Path) -> dict[str, Any]: + try: + data = json.loads(index_file.read_text(encoding="utf-8")) + if ( + isinstance(data, dict) + and data.get("version") == INDEX_VERSION + and data.get("zone") == index_zone() + and isinstance(data.get("files"), dict) + ): + return data["files"] + except Exception: + pass + return {} + + +# What was already read, by its first and last few kilobytes: a transcript +# rewritten in place rather than appended to changes one or the other. +def read_digest(path: Path, length: int) -> str: + digest = hashlib.sha1() + with path.open("rb") as handle: + digest.update(handle.read(min(length, INDEX_HEAD_BYTES))) + handle.seek(max(0, length - INDEX_HEAD_BYTES)) + digest.update(handle.read(min(length, INDEX_HEAD_BYTES))) + return digest.hexdigest() + + +def parse_usage_records(path: Path, entry: dict[str, Any] | None) -> dict[str, Any]: + """Return the index entry for one transcript, reusing `entry` where the file + only grew. Records are [key, model, day, input, output, cacheRead, + cacheWrite, session]; `day` is None when the line carries no timestamp.""" + stat = path.stat() + size = stat.st_size + mtime = stat.st_mtime_ns + + if entry and entry.get("size") == size and entry.get("mtime") == mtime: + return entry + + # Only the same file, grown past what was read with what was read still in + # place, is an append. A transcript replaced by a larger one is a new file, + # and one rewritten in place has a different start or end. + if ( + entry + and size > number(entry.get("size")) + and entry.get("file") == [stat.st_dev, stat.st_ino] + and entry.get("head") == read_digest(path, number(entry.get("offset"))) + ): + records = list(entry.get("records") or []) + offset = number(entry.get("offset")) + line_number = number(entry.get("lines")) + else: + records = [] + offset = 0 + line_number = 0 + + # A streamed response is several lines sharing one message id, and the + # first line's output_tokens is a placeholder while the last carries the + # real count. Keep each message's line with the highest output count (the + # last one read wins a tie), whole: a response that fell back to another + # model mid-stream has lines that are separate snapshots, and mixing fields + # across them over-counts cache tokens and credits the wrong model. + positions = {record[0]: index for index, record in enumerate(records)} + with path.open("rb") as handle: + handle.seek(offset) + for raw in handle: + if not raw.endswith(b"\n"): + # A line still being written; pick it up whole next time. + break + offset += len(raw) + line_number += 1 + + # Cheap pre-filter before JSON parsing keeps files with unrelated + # lines inexpensive. + if b'"usage":' not in raw: + continue + + try: + parsed = json.loads(raw.decode("utf-8", errors="replace")) + except Exception: + continue + + message = parsed.get("message") if isinstance(parsed.get("message"), dict) else {} + if parsed.get("type") != "assistant" and message.get("role") != "assistant": + continue + + usage = message.get("usage") or parsed.get("usage") + if not isinstance(usage, dict): + continue + + message_id = message.get("id") or parsed.get("messageId") or "" + key = str(message_id) if message_id else f"{path}:{parsed.get('uuid') or parsed.get('requestId') or line_number}" + + input_tokens = usage_token(usage, "input_tokens", "inputTokens") + output_tokens = usage_token(usage, "output_tokens", "outputTokens") + cache_read = usage_token(usage, "cache_read_input_tokens", "cacheReadInputTokens") + cache_write = usage_token(usage, "cache_creation_input_tokens", "cacheCreationInputTokens") + if input_tokens + output_tokens + cache_read + cache_write <= 0: + continue + + held = positions.get(key) + if held is not None and output_tokens < records[held][4]: + continue + + timestamp = parsed.get("timestamp") or message.get("timestamp") + record = [ + key, + str(message.get("model") or parsed.get("model") or "claude"), + local_date_from_timestamp(timestamp) if timestamp is not None else None, + input_tokens, + output_tokens, + cache_read, + cache_write, + str(parsed.get("sessionId") or path), + ] + if held is None: + positions[key] = len(records) + records.append(record) + else: + records[held] = record + + return { + "size": size, + "file": [stat.st_dev, stat.st_ino], + "mtime": mtime, + "head": read_digest(path, offset), + "offset": offset, + "lines": line_number, + "records": records, + } + + +def scan_projects(projects_path: Path, index_file: Path | None = None) -> dict[str, Any]: + previous = load_index(index_file) if index_file else {} + current: dict[str, Any] = {} + + # Directory order, as before: when one API message shows up in several + # transcripts with the same output count, the last file visited wins, and + # the numbers must not move because the scan got faster. + order: list[str] = [] + unreadable_count = 0 + unreadable_examples: list[str] = [] + files = projects_path.rglob("*.jsonl") if projects_path.is_dir() else [] + for path in files: + key = str(path) + try: + current[key] = parse_usage_records(path, previous.get(key)) + except Exception as exc: + unreadable_count += 1 + if len(unreadable_examples) < 3: + unreadable_examples.append(f"{path}: {exc}") + # What was read before is still true of the file. + if key in previous: + current[key] = previous[key] + if key in current: + order.append(key) + + if unreadable_count: + examples = "; ".join(unreadable_examples) + print( + f"Ignoring {unreadable_count} unreadable Claude project files" + f" (examples: {examples})", + file=sys.stderr, + ) + + if index_file and current != previous: + try: + write_json(index_file, {"version": INDEX_VERSION, "zone": index_zone(), "files": current}) + except Exception as exc: + print(f"Could not write Claude scan index {index_file}: {exc}", file=sys.stderr) + + return summarize_records(current, order) + + +def summarize_records(indexed: dict[str, Any], order: list[str]) -> dict[str, Any]: today = local_date_string() recent_dates = recent_date_strings() recent = {day: {"date": day, "messageCount": 0} for day in recent_dates} - seen: set[str] = set() + # One API message can show up in several transcripts (a resumed session + # carries the earlier ones). As within a file, the record with the highest + # output count wins, the last one visited breaking a tie. + seen: dict[str, list[Any]] = {} sessions: set[str] = set() active_days: set[str] = set() today_sessions: set[str] = set() @@ -137,68 +334,38 @@ def scan_projects(projects_path: Path) -> dict[str, Any]: today_prompt_count = 0 today_token_total = 0 - files = projects_path.rglob("*.jsonl") if projects_path.is_dir() else [] - for path in files: - try: - with path.open("r", encoding="utf-8", errors="replace") as handle: - for line_number, line in enumerate(handle, 1): - # Cheap pre-filter before JSON parsing keeps files with unrelated - # lines inexpensive. - if '"usage":' not in line: - continue + for key in order: + for record in indexed[key].get("records") or []: + held = seen.get(record[0]) + if held is not None and record[4] < held[4]: + continue + seen[record[0]] = record - try: - entry = json.loads(line) - except Exception: - continue + for _, model, day, input_tokens, output_tokens, cache_read, cache_write, session_key in seen.values(): + if day is None: + day = today + total = input_tokens + output_tokens + cache_read + cache_write - message = entry.get("message") if isinstance(entry.get("message"), dict) else {} - if entry.get("type") != "assistant" and message.get("role") != "assistant": - continue + sessions.add(session_key) + active_days.add(day) + prompts += 1 - usage = message.get("usage") or entry.get("usage") - if not isinstance(usage, dict): - continue + bucket = usage_by_model.setdefault(model, empty_bucket()) + bucket["inputTokens"] += input_tokens + bucket["outputTokens"] += output_tokens + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write - message_id = message.get("id") or entry.get("messageId") or "" - unique_key = str(message_id) if message_id else f"{path}:{entry.get('uuid') or entry.get('requestId') or line_number}" - if unique_key in seen: - continue - seen.add(unique_key) + if day in recent: + # recentDays.messageCount is actually a token total, despite the + # legacy name shared with synced snapshots. + recent[day]["messageCount"] += total - input_tokens = usage_token(usage, "input_tokens", "inputTokens") - output_tokens = usage_token(usage, "output_tokens", "outputTokens") - cache_read = usage_token(usage, "cache_read_input_tokens", "cacheReadInputTokens") - cache_write = usage_token(usage, "cache_creation_input_tokens", "cacheCreationInputTokens") - total = input_tokens + output_tokens + cache_read + cache_write - if total <= 0: - continue - - model = str(message.get("model") or entry.get("model") or "claude") - day = local_date_from_timestamp(entry.get("timestamp") or message.get("timestamp")) - session_key = str(entry.get("sessionId") or path) - sessions.add(session_key) - active_days.add(day) - prompts += 1 - - bucket = usage_by_model.setdefault(model, empty_bucket()) - bucket["inputTokens"] += input_tokens - bucket["outputTokens"] += output_tokens - bucket["cacheReadInputTokens"] += cache_read - bucket["cacheCreationInputTokens"] += cache_write - - if day in recent: - # recentDays.messageCount is actually a token total, despite the - # legacy name shared with synced snapshots. - recent[day]["messageCount"] += total - - if day == today: - today_prompt_count += 1 - today_sessions.add(session_key) - today_token_total += total - today_tokens[model] = today_tokens.get(model, 0) + total - except Exception as exc: - print(f"Ignoring unreadable Claude project file {path}: {exc}", file=sys.stderr) + if day == today: + today_prompt_count += 1 + today_sessions.add(session_key) + today_token_total += total + today_tokens[model] = today_tokens.get(model, 0) + total return { "todayPrompts": today_prompt_count, @@ -217,18 +384,28 @@ def scan_projects(projects_path: Path) -> dict[str, Any]: } -def scan_cache_paths(projects_path: Path) -> tuple[Path, Path]: +def scan_cache_paths(projects_path: Path) -> tuple[Path, Path, Path]: digest = hashlib.sha1(str(projects_path).encode("utf-8")).hexdigest()[:16] root = cache_root() - return root / f"claude-scan-{digest}.json", root / f"claude-scan-{digest}.lock" + return root / f"claude-scan-{digest}.json", root / f"claude-scan-{digest}.lock", root / f"claude-index-{digest}.json" + + +def read_json(path: Path) -> dict[str, Any] | None: + try: + return json.loads(path.read_text(encoding="utf-8")) + except Exception: + return None def read_fresh_json(path: Path, max_age_seconds: float) -> dict[str, Any] | None: if max_age_seconds <= 0 or not path.exists(): return None try: - if time.time() - path.stat().st_mtime <= max_age_seconds: - return json.loads(path.read_text(encoding="utf-8")) + # A negative age means the mtime is in the future: the clock moved + # backwards since the write, so the cache's freshness cannot be trusted. + age = time.time() - path.stat().st_mtime + if 0 <= age <= max_age_seconds: + return read_json(path) except Exception: return None return None @@ -252,8 +429,8 @@ def write_json(path: Path, payload: dict[str, Any]) -> None: raise -def cached_scan(projects_path: Path, max_age_seconds: float) -> dict[str, Any]: - cache_file, lock_file = scan_cache_paths(projects_path) +def cached_scan(projects_path: Path, max_age_seconds: float, use_index: bool = True) -> dict[str, Any]: + cache_file, lock_file, index_file = scan_cache_paths(projects_path) cached = read_fresh_json(cache_file, max_age_seconds) if cached is not None: @@ -264,7 +441,9 @@ def cached_scan(projects_path: Path, max_age_seconds: float) -> dict[str, Any]: cached = read_fresh_json(cache_file, max_age_seconds) if cached is not None: return cached - summary = scan_projects(projects_path) + if not use_index: + index_file.unlink(missing_ok=True) + summary = scan_projects(projects_path, index_file) write_json(cache_file, summary) return summary @@ -341,11 +520,27 @@ def today_prompts_from_history(claude_dir: Path) -> tuple[int, int]: # provider, model, and token usage on every assistant message. +# `omp --profile=` (and pi's equivalent) relocates the whole agent tree +# under /profiles//, so a subscription driven entirely through a +# profile leaves the default root empty and its usage uncounted. Sessions are +# keyed by file path, so a profile adds sessions rather than double-counting +# the default root. +def pi_session_roots() -> list[Path]: + roots: list[Path] = [] + for base in (Path.home() / ".pi", Path.home() / ".omp"): + roots.append(base / "agent" / "sessions") + profiles = base / "profiles" + try: + # Sorted so the scan order does not depend on directory order. + roots.extend(sorted(child / "agent" / "sessions" for child in profiles.iterdir() if child.is_dir())) + except OSError: + # No profiles directory, or it is unreadable: the default root stands. + pass + return roots + + def scan_pi_usage(max_age_seconds: float) -> dict[str, Any] | None: - roots = [ - Path.home() / ".pi" / "agent" / "sessions", - Path.home() / ".omp" / "agent" / "sessions", - ] + roots = pi_session_roots() cache_file = cache_root() / "claude-pi-sessions.json" cached = read_fresh_json(cache_file, max_age_seconds) if cached is not None: @@ -477,16 +672,36 @@ def scan_opencode_usage(max_age_seconds: float) -> dict[str, Any] | None: return None try: conn.execute("PRAGMA query_only = ON") - for session_id, raw in conn.execute("SELECT session_id, data FROM message"): + tables = {name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table'")} + # opencode v1 keeps messages in `message`, with the role, provider and model + # in each row's data. v2 writes `session_message` instead: the role moves to + # the `type` column and the provider and model nest under data.model. A + # database can hold either or both. + queries = [] + if "message" in tables: + queries.append(("opencode:", "SELECT session_id, data, NULL FROM message")) + if "session_message" in tables: + queries.append(("opencode-v2:", "SELECT session_id, data, type FROM session_message")) + rows = ((prefix, row) for prefix, query in queries for row in conn.execute(query)) + for session_prefix, (session_id, raw, row_type) in rows: # One malformed row must not abort the scan, so every shape assumption # lives inside the try. try: entry = json.loads(raw) + if not isinstance(entry, dict): + continue + if row_type is None: + role = entry.get("role") + provider = entry.get("providerID") + model_id = entry.get("modelID") + else: + role = row_type + model_info = entry.get("model") or {} + provider = model_info.get("providerID") + model_id = model_info.get("id") # Exact match: opencode provider ids are free-form, and a custom # "anthropic-proxy" gateway is not this subscription. - if not isinstance(entry, dict) or entry.get("role") != "assistant": - continue - if str(entry.get("providerID") or "") != "anthropic": + if role != "assistant" or str(provider or "") != "anthropic": continue tokens = entry.get("tokens") or {} cache = tokens.get("cache") or {} @@ -501,10 +716,10 @@ def scan_opencode_usage(max_age_seconds: float) -> dict[str, Any] | None: created = number((entry.get("time") or {}).get("created")) day = dt.datetime.fromtimestamp(created / 1000).strftime("%Y-%m-%d") if created > 0 else today - model = str(entry.get("modelID") or "claude").rstrip("/").split("/")[-1] + model = str(model_id or "claude").rstrip("/").split("/")[-1] except Exception: continue - session_key = "opencode:" + str(session_id) + session_key = session_prefix + str(session_id) sessions.add(session_key) active_days.add(day) prompts += 1 @@ -585,28 +800,75 @@ def merge_stats(base: dict[str, Any], extra: dict[str, Any]) -> dict[str, Any]: # CLI's login. Nothing else leaves the credential store: the token goes # nowhere but the Authorization header of the limits probe, and only the # plan label may travel into the printed record. -def oauth_login(claude_dir: Path) -> tuple[str, int, str]: +def oauth_login(claude_dir: Path) -> tuple[str, int, int, str]: try: data = json.loads((claude_dir / ".credentials.json").read_text(encoding="utf-8")) except Exception: - return "", 0, "" + return "", 0, 0, "" login = data.get("claudeAiOauth") if not isinstance(login, dict): - return "", 0, "" - plan = plan_label(str(login.get("rateLimitTier") or ""), str(login.get("subscriptionType") or "")) - return str(login.get("accessToken") or ""), number(login.get("expiresAt")), plan + return "", 0, 0, "" + # The profile's tier is asked first, but only one that names a multiplier + # outranks the credentials'. The subscription still leads the label, so a + # Team seat keeps reading Team. + tier = profile_tier(claude_dir) + if not re.search(r"max_(\d+x)", tier, re.IGNORECASE): + tier = str(login.get("rateLimitTier") or "") + plan = plan_label(tier, str(login.get("subscriptionType") or "")) + return ( + str(login.get("accessToken") or ""), + number(login.get("expiresAt")), + number(login.get("refreshTokenExpiresAt")), + plan, + ) -def plan_label(tier: str, subscription: str) -> str: - if tier: - match = re.search(r"max_(\d+x)", tier, re.IGNORECASE) - if match: - return "Max " + match.group(1) - if subscription: - return subscription[0].upper() + subscription[1:] +# The credentials carry the tier their token was minted with, and an upgrade +# does not restamp it: a refresh rewrites the file and keeps the old value, so +# an account moved to Max 20x keeps reading Max 5x. The CLI's own profile is +# refreshed against the account, so the tier it states there is asked first. +def profile_tier(claude_dir: Path) -> str: + try: + account = json.loads(profile_path(claude_dir).read_text(encoding="utf-8")).get("oauthAccount") + except Exception: + return "" + if not isinstance(account, dict): + return "" + # A seat with no tier of its own is limited by its organization, which is + # the shape a personal Max subscription takes. + for key in ("userRateLimitTier", "organizationRateLimitTier"): + tier = str(account.get(key) or "") + if tier: + return tier return "" +# The CLI's profile for a config home: a legacy .config.json inside it, else +# ~/.claude.json for the default home and /.claude.json for one it was +# pointed at with CLAUDE_CONFIG_DIR. A copy anywhere else is never refreshed, +# and another home's profile describes another account. +def profile_path(home: Path) -> Path: + legacy = home / ".config.json" + if legacy.exists(): + return legacy + configured = os.environ.get("CLAUDE_CONFIG_DIR") + if home == expand_path("~/.claude") and not (configured and expand_path(configured) == home): + return Path.home() / ".claude.json" + return home / ".claude.json" + + +# The rate-limit tier names the ceiling, not the plan: a Team premium seat +# runs on `default_claude_max_5x` while paying for Team. Lead with the +# subscription and keep the multiplier as its qualifier, so Max still reads +# "Max 5x" and a Team seat reads "Team 5x" rather than passing as Max. +def plan_label(tier: str, subscription: str) -> str: + plan = subscription[0].upper() + subscription[1:] if subscription else "" + match = re.search(r"max_(\d+x)", tier, re.IGNORECASE) if tier else None + if not match: + return plan + return (plan or "Max") + " " + match.group(1) + + def parse_utilization(value: Any) -> float: try: return float(str(value).strip().replace("%", "")) @@ -795,7 +1057,7 @@ def usable_cached_limits(cached: dict[str, Any]) -> list[dict[str, Any]]: return [entry for entry in entries if isinstance(entry, dict) and limit_window_open(entry, now)] -def collect_limits(access_token: str, expires_at_ms: int, force: bool, cache_name: str = "claude-limits.json") -> dict[str, Any]: +def collect_limits(access_token: str, expires_at_ms: int, refresh_expires_at_ms: int, force: bool, cache_name: str = "claude-limits.json") -> dict[str, Any]: # `live` says the numbers came from Anthropic just now (or within the reuse # window) rather than from a cache kept past a failed or impossible probe. result = {"limits": [], "usageStatusText": "", "authHelpText": AUTH_HELP, "live": False} @@ -809,7 +1071,7 @@ def collect_limits(access_token: str, expires_at_ms: int, force: bool, cache_nam # kept as the answer of record when a later probe fails. Every account # keeps its own, so one account's allowance never stands in for another's. probe_cache = cache_root() / cache_name - cached = read_fresh_json(probe_cache, float("inf")) or {} + cached = read_json(probe_cache) or {} fallback = usable_cached_limits(cached) result["fetchedAtMs"] = number(cached.get("fetchedAtMs")) @@ -824,6 +1086,19 @@ def collect_limits(access_token: str, expires_at_ms: int, force: bool, cache_nam return result if expires_at_ms > 0 and expires_at_ms <= time.time() * 1000: result["limits"] = fallback + # A lapsed access token is routine, not a sign-out: it lives about eight + # hours, and the CLI mints a new one from the refresh token the next time + # it runs. Only a lapsed *refresh* token is a sign-in a person must redo, + # and saying so either way sends people to `claude auth login` nightly -- + # which does not restore anything the CLI would not have restored itself. + if refresh_expires_at_ms > time.time() * 1000: + result["usageStatusText"] = "Limits paused" + result["authHelpText"] = ( + "Claude Code's access token has lapsed" + + (" — showing the last known limits." if fallback else ".") + + " Start Claude Code to refresh it; signing in again is not needed." + ) + return result result["usageStatusText"] = "Sign-in expired" result["authHelpText"] = ( "Claude Code's saved sign-in expired" @@ -836,15 +1111,18 @@ def collect_limits(access_token: str, expires_at_ms: int, force: bool, cache_nam # entirely; the interval is there to absorb repeated panel opens, not to # overrule someone who pressed refresh. fetched_at = number(cached.get("fetchedAtMs")) / 1000 - if fallback and not force and time.time() - fetched_at < PROBE_MIN_INTERVAL_SECONDS: + fetched_age = time.time() - fetched_at + if fallback and not force and 0 <= fetched_age < PROBE_MIN_INTERVAL_SECONDS: result["limits"] = fallback result["live"] = True + result["authHelpText"] = "" return result probe = probe_limits(access_token) if probe["ok"]: result["limits"] = probe["limits"] result["live"] = True + result["authHelpText"] = "" result["fetchedAtMs"] = round(time.time() * 1000) write_json(probe_cache, {"fetchedAtMs": result["fetchedAtMs"], "limits": probe["limits"]}) return result @@ -855,7 +1133,8 @@ def collect_limits(access_token: str, expires_at_ms: int, force: bool, cache_nam result["retryAdvised"] = True if fallback: result["limits"] = fallback - result["live"] = time.time() - result["fetchedAtMs"] / 1000 < CURRENT_ENOUGH_SECONDS + fetched_age = time.time() - result["fetchedAtMs"] / 1000 + result["live"] = 0 <= fetched_age < CURRENT_ENOUGH_SECONDS else: result["usageStatusText"] = "Claude limits unavailable" result["authHelpText"] = probe["helpText"] @@ -895,7 +1174,7 @@ def registered_accounts() -> list[dict[str, Any]]: # know this account is the one with the most room. Each window it last saw # comes back at 0%. def replenished_limits(cache_name: str) -> list[dict[str, Any]]: - cached = read_fresh_json(cache_root() / cache_name, float("inf")) or {} + cached = read_json(cache_root() / cache_name) or {} entries = cached.get("limits") if not isinstance(entries, list): return [] @@ -908,9 +1187,8 @@ def replenished_limits(cache_name: str) -> list[dict[str, Any]]: def current_account_id(home: Path) -> str: - path = Path.home() / ".claude.json" if home == expand_path("~/.claude") else home / ".claude.json" try: - account = json.loads(path.read_text(encoding="utf-8")).get("oauthAccount") or {} + account = json.loads(profile_path(home).read_text(encoding="utf-8")).get("oauthAccount") or {} except Exception: return "" return str(account.get("accountUuid") or "") @@ -918,7 +1196,7 @@ def current_account_id(home: Path) -> str: def account_limits(account: dict[str, Any], force: bool) -> dict[str, Any]: home = expand_path(account["home"]) if account.get("home") else expand_path("~/.claude") - access_token, expires_at_ms, plan = oauth_login(home) + access_token, expires_at_ms, refresh_expires_at_ms, plan = oauth_login(home) # Keyed by the subscription, not the label: forget "Work" and add a different # "Work", and the new one must never fall back on the old one's allowance. # Who the home is signed in as now, not when it was added: a `claude auth @@ -932,7 +1210,7 @@ def account_limits(account: dict[str, Any], force: bool) -> dict[str, Any]: cache_name = "claude-limits.json" else: cache_name = f"claude-limits-{key}.json" - limits = collect_limits(access_token, expires_at_ms, force, cache_name) + limits = collect_limits(access_token, expires_at_ms, refresh_expires_at_ms, force, cache_name) if not limits["live"] and not limits["limits"]: limits["limits"] = replenished_limits(cache_name) return { @@ -956,14 +1234,14 @@ def account_limits(account: dict[str, Any], force: bool) -> dict[str, Any]: def main() -> int: parser = argparse.ArgumentParser() - parser.add_argument("--force", action="store_true", help="rescan transcripts and re-probe limits, ignoring caches") + parser.add_argument("--force", action="store_true", help="rescan transcripts from scratch and re-probe limits, ignoring caches") parser.add_argument("--limits-only", action="store_true", help="reuse any recent transcript scan; only the limits probe must be fresh") parser.add_argument("--cache-seconds", type=float, default=20) args = parser.parse_args() claude_dir = config_dir() scan_age = 0 if args.force else (900 if args.limits_only else args.cache_seconds) - stats = cached_scan(claude_dir / "projects", scan_age) + stats = cached_scan(claude_dir / "projects", scan_age, use_index=not args.force) if number(stats.get("totalPrompts")) <= 0: fallback = stats_cache_fallback(claude_dir) @@ -994,8 +1272,8 @@ def main() -> int: limits = {key: current[key] for key in ("limits", "usageStatusText", "authHelpText", "retryAdvised")} stale, fetched_at = current["stale"], current["fetchedAt"] else: - access_token, expires_at_ms, plan = oauth_login(claude_dir) - limits = collect_limits(access_token, expires_at_ms, args.force) + access_token, expires_at_ms, refresh_expires_at_ms, plan = oauth_login(claude_dir) + limits = collect_limits(access_token, expires_at_ms, refresh_expires_at_ms, args.force) stale, fetched_at = not limits["live"], limits["fetchedAtMs"] record = { diff --git a/bin/omarchy-agent-usage-codex b/bin/omarchy-agent-usage-codex index b9e1b3d4..c7635994 100755 --- a/bin/omarchy-agent-usage-codex +++ b/bin/omarchy-agent-usage-codex @@ -22,6 +22,7 @@ import subprocess import sys import tempfile import time +import tomllib from datetime import datetime, timedelta, timezone from pathlib import Path @@ -29,6 +30,11 @@ AGENT_ID = "codex" AGENT_NAME = "Codex" AUTH_HELP = "Run `codex login` to authenticate." +# session_meta.model_provider for the built-in OpenAI backend, which is the +# only one this subscription pays for. Any other id points somewhere else: +# `--oss`'s ollama or lmstudio, or a custom provider in config.toml. +NATIVE_PROVIDER = "openai" + # A scan this recent is only reused to dedup concurrent collector runs (the # update command backgrounds one per agent while the panel refreshes on its # own); every periodic widget refresh lands a real rescan, however low @@ -37,6 +43,11 @@ AUTH_HELP = "Run `codex login` to authenticate." SCAN_REUSE_SECONDS = 20 LIMITS_ONLY_REUSE_SECONDS = 900 +# Per-file totals of native Codex sessions survive between runs, so a rescan +# only reads the session files that changed. Bumping the schema retires every +# record written by an older layout. +FILE_CACHE_SCHEMA = 1 + def local_day(value): if value is None: @@ -91,6 +102,52 @@ def find_command(name): return shutil.which(name, path=ENV.get("PATH")) +def mise_shims_dir(): + data_home = ENV.get("MISE_DATA_DIR") or os.path.join( + ENV.get("XDG_DATA_HOME") or os.path.join(str(Path.home()), ".local", "share"), "mise") + return os.path.join(data_home, "shims") + + +def is_lazy_launcher(path): + # omarchy-mise-install writes a regular file whose first step is `mise use + # -g`, and mise's shims exec `mise x`; running either installs the tool it + # wraps. A shim is a symlink to the mise binary itself, so the symlink + # check comes after the shim checks: a symlink elsewhere is the user's own + # binary. + if os.path.dirname(os.path.abspath(path)) == mise_shims_dir(): + return True + if os.path.islink(path): + mise = find_command("mise") + return bool(mise) and os.path.realpath(path) == os.path.realpath(mise) + try: + with open(path, "rb") as handle: + head = handle.read(4096) + return b"mise use" in head or b"mise x " in head + except OSError: + return False + + +def find_codex_binary(): + for directory in ENV.get("PATH", "").split(os.pathsep): + candidate = os.path.join(directory, "codex") + if os.path.isfile(candidate) and os.access(candidate, os.X_OK) and not is_lazy_launcher(candidate): + return candidate + # `mise which` only prints a binary that is already installed, so it is the + # safe way to resolve a codex managed by mise without running the launcher. + # `mise which` can resolve a configured `latest` over the network; the usage + # probe stays offline. + mise = find_command("mise") + if mise: + try: + out = subprocess.run([mise, "which", "codex"], capture_output=True, text=True, env=dict(ENV, MISE_OFFLINE="1"), timeout=10) + resolved = out.stdout.strip() + if out.returncode == 0 and os.path.isfile(resolved) and not is_lazy_launcher(resolved): + return resolved + except Exception: + pass + return None + + now = datetime.now() today = now.strftime("%Y-%m-%d") recent_dates = [(now - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)] @@ -107,46 +164,109 @@ total_sessions = set() seen_pi_messages = set() -def add_usage(day, session_key, model, input_tokens, output_tokens, cache_read, cache_write): +# One session file's totals, the unit the file cache stores: day -> model -> +# [input, output, cacheRead, cacheWrite, prompts]. Day-keyed and with no +# "today" in it, so a record stays true however long it is kept. +def new_file_record(stat_result): + return {"mtime": stat_result.st_mtime, "size": stat_result.st_size, "days": {}} + + +def record_usage(record, day, model, input_tokens, output_tokens, cache_read, cache_write): + totals = record["days"].setdefault(day, {}).setdefault(model, [0, 0, 0, 0, 0]) + totals[0] += input_tokens + totals[1] += output_tokens + totals[2] += cache_read + totals[3] += cache_write + totals[4] += 1 + + +def merge_file_record(session_key, record): + """Fold one file's totals, freshly read or replayed from cache, into the run.""" global today_prompts, today_total_tokens, total_prompts - total = input_tokens + output_tokens + cache_read + cache_write - total_prompts += 1 - total_sessions.add(session_key) - active_days.add(day) + for day, models in record["days"].items(): + for model, totals in models.items(): + input_tokens, output_tokens, cache_read, cache_write, prompts = totals + total = input_tokens + output_tokens + cache_read + cache_write + total_prompts += prompts + total_sessions.add(session_key) + active_days.add(day) - bucket = model_usage.setdefault(model, { - "inputTokens": 0, - "outputTokens": 0, - "cacheReadInputTokens": 0, - "cacheCreationInputTokens": 0, - }) - bucket["inputTokens"] += input_tokens - bucket["outputTokens"] += output_tokens - bucket["cacheReadInputTokens"] += cache_read - bucket["cacheCreationInputTokens"] += cache_write + bucket = model_usage.setdefault(model, { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + }) + bucket["inputTokens"] += input_tokens + bucket["outputTokens"] += output_tokens + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write - if day in recent: - recent[day]["messageCount"] += total + if day in recent: + recent[day]["messageCount"] += total - if day == today: - today_prompts += 1 - today_sessions.add(session_key) - today_total_tokens += total - today_tokens_by_model[model] = today_tokens_by_model.get(model, 0) + total + if day == today: + today_prompts += prompts + today_sessions.add(session_key) + today_total_tokens += total + today_tokens_by_model[model] = today_tokens_by_model.get(model, 0) + total + + +def add_usage(day, session_key, model, input_tokens, output_tokens, cache_read, cache_write): + record = {"days": {}} + record_usage(record, day, model, input_tokens, output_tokens, cache_read, cache_write) + merge_file_record(session_key, record) + + +def pi_parent_session(path): + try: + with open(path, encoding="utf-8", errors="replace") as session_file: + header = json.loads(session_file.readline()) + except (OSError, json.JSONDecodeError): + return None + + if not isinstance(header, dict): + return None + parent = header.get("parentSession") + if not isinstance(parent, str) or not parent: + return None + return os.path.abspath(os.path.expanduser(parent)) + + +# `omp --profile=` (and pi's equivalent) relocates the whole agent tree +# under /profiles//, so a subscription driven entirely through a +# profile leaves the default root empty and its usage uncounted. Each root is +# scanned separately and sessions are keyed by path, so a profile adds +# sessions rather than double-counting the default one. +def pi_session_roots(): + roots = [] + for base in (Path.home() / ".pi", Path.home() / ".omp"): + roots.append(base / "agent" / "sessions") + profiles = base / "profiles" + try: + # Sorted so the scan order does not depend on directory order. + roots.extend(sorted(child / "agent" / "sessions" for child in profiles.iterdir() if child.is_dir())) + except OSError: + # No profiles directory, or it is unreadable: the default root stands. + pass + return roots def scan_pi_sessions(): - roots = [ - Path.home() / ".pi" / "agent" / "sessions", - Path.home() / ".omp" / "agent" / "sessions", - ] + roots = pi_session_roots() rg = find_command("rg") or "rg" + session_keys = {} + today_session_keys = {} for root in roots: if not root.exists(): continue try: proc = subprocess.Popen( - [rg, "--json", "-e", r'"provider"\s*:\s*"openai-codex"', "-e", r'"api"\s*:\s*"openai-codex', str(root)], + # --no-ignore: these are data files, not a source tree. When $HOME is + # itself a git checkout (a common dotfiles setup), the parent repo's + # ignore rules would otherwise silently exclude every session file and + # the scan would count zero usage. + [rg, "--json", "--no-ignore", "-e", r'"provider"\s*:\s*"openai-codex"', "-e", r'"api"\s*:\s*"openai-codex', str(root)], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, @@ -163,16 +283,22 @@ def scan_pi_sessions(): if event.get("type") != "match": continue line = event.get("data", {}).get("lines", {}).get("text", "") - path = event.get("data", {}).get("path", {}).get("text", "pi-session") + raw_path = event.get("data", {}).get("path", {}).get("text", "pi-session") + path = os.path.abspath(os.path.expanduser(raw_path)) entry = json.loads(line) except Exception: continue if entry.get("type") != "message": continue - message_key = path + ":" + str(entry.get("id") or "") - if message_key in seen_pi_messages: - continue + message_id = str(entry.get("id") or "") + message_timestamp = str(entry.get("timestamp") or "") + # Forks retain both fields; IDs alone can collide across sessions. + if message_id and message_timestamp: + message_key = ("entry", message_id, message_timestamp) + else: + message_key = ("path", path, message_id) + duplicate = message_key in seen_pi_messages seen_pi_messages.add(message_key) message = entry.get("message") or {} if message.get("role") != "assistant": @@ -196,14 +322,37 @@ def scan_pi_sessions(): continue day = local_day(entry.get("timestamp") or message.get("timestamp")) - session_key = path - add_usage(day, session_key, model_name(message.get("model")), input_tokens, output_tokens, cache_read, cache_write) + session_keys.setdefault(path, set()).add(message_key) + if day == today: + today_session_keys.setdefault(path, set()).add(message_key) + if duplicate: + continue + add_usage(day, path, model_name(message.get("model")), input_tokens, output_tokens, cache_read, cache_write) try: proc.wait(timeout=1) except Exception: proc.kill() + # Resolve session ownership after scanning so rg's file order cannot affect it. + parents = {path: pi_parent_session(path) for path in session_keys} + total_sessions.difference_update(session_keys) + today_sessions.difference_update(session_keys) + for path, keys in session_keys.items(): + inherited = set() + visited = {path} + parent = parents.get(path) + while parent and parent not in visited: + visited.add(parent) + inherited.update(session_keys.get(parent, ())) + parent = parents.get(parent) + + unique_keys = keys - inherited + if unique_keys: + total_sessions.add(path) + if unique_keys & today_session_keys.get(path, set()): + today_sessions.add(path) + def scan_opencode_sessions(): # A subscription burned entirely through opencode leaves no native session @@ -283,7 +432,7 @@ def scan_opencode_sessions(): return True -def scan_native_codex_sessions(): +def scan_native_codex_sessions(cached_files, scanned_files): codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) roots = [codex_home / "sessions", codex_home / "archived_sessions"] files = [] @@ -293,48 +442,98 @@ def scan_native_codex_sessions(): continue for path in root.rglob("*.jsonl"): try: - if path.stat().st_mtime >= cutoff: - files.append(path) + stat_result = path.stat() except OSError: - pass + continue + if stat_result.st_mtime >= cutoff: + files.append((path, stat_result)) - for path in files: - current_model = "codex" - try: - with path.open(errors="replace") as handle: - for raw in handle: - try: - entry = json.loads(raw) - except Exception: + for path, stat_result in files: + key = str(path) + record = reusable_record(cached_files, key, stat_result) + if record is None: + record = read_native_codex_session(path, stat_result) + if record["mtime"] is not None: + scanned_files[key] = record + merge_file_record(key, record) + + +def read_native_codex_session(path, stat_result): + record = new_file_record(stat_result) + current_model = "codex" + seen_meta = False + previous_total_usage = None + try: + with path.open(errors="replace") as handle: + for raw in handle: + # Cheap pre-filter before JSON parsing keeps files with unrelated + # lines inexpensive. session_meta still has to get through: it says + # which provider served the rollout. + if '"token_count"' not in raw and '"turn_context"' not in raw and '"session_meta"' not in raw: + continue + try: + entry = json.loads(raw) + except Exception: + continue + if entry.get("type") == "session_meta": + # A forked rollout copies its parent's session_meta after its own, + # so only the first one, as Codex reads it, says who served this. + if seen_meta: continue - if entry.get("type") == "turn_context": - payload = entry.get("payload") or {} - current_model = model_name(payload.get("model") or payload.get("model_slug") or current_model) + seen_meta = True + # Codex CLI happily fronts any OpenAI-compatible backend (`--oss`, + # or a custom model_provider in config.toml). Those turns bill the + # local box or a third party, never this subscription, so drop the + # whole rollout the way the pi and opencode scans drop foreign + # providers. + meta = entry.get("payload") or {} + provider = str(meta.get("model_provider") or "") + # Rollouts predating the field carry no provider at all; count + # those rather than silently lose the history. + if provider and provider != NATIVE_PROVIDER: + break + continue + if entry.get("type") == "turn_context": + payload = entry.get("payload") or {} + current_model = model_name(payload.get("model") or payload.get("model_slug") or current_model) + continue + payload = entry.get("payload") or entry + if entry.get("type") == "response_item" and isinstance(payload, dict): + payload = payload.get("payload") or payload + if not isinstance(payload, dict): + continue + if payload.get("type") != "token_count": + continue + info = payload.get("info") or {} + # total_token_usage is cumulative for the session. Adding every + # snapshot makes usage grow quadratically, so count the last turn. + usage = info.get("last_token_usage") or {} + cache_read = number(usage.get("cached_input_tokens")) + cache_write = number(usage.get("cache_write_input_tokens")) + # Cached tokens are included in input_tokens, and reasoning tokens + # are included in output_tokens. Keep the cache split without + # counting either category twice. + input_tokens = max(0, number(usage.get("input_tokens")) - cache_read - cache_write) + output_tokens = number(usage.get("output_tokens")) + if not (input_tokens or output_tokens or cache_read or cache_write): + continue + # Quota updates can repeat the previous request's last_token_usage. + # Only an unchanged cumulative snapshot proves this is a repeat: + # separate requests may have identical last-token counts, and older + # records may not carry cumulative counters at all. Keep this state + # per rollout and accept changed counters, including counter resets. + total_usage = info.get("total_token_usage") + if isinstance(total_usage, dict) and total_usage: + if total_usage == previous_total_usage: continue - payload = entry.get("payload") or entry - if entry.get("type") == "response_item" and isinstance(payload, dict): - payload = payload.get("payload") or payload - if not isinstance(payload, dict): - continue - if payload.get("type") != "token_count": - continue - info = payload.get("info") or {} - # total_token_usage is cumulative for the session. Adding every - # snapshot makes usage grow quadratically, so count the last turn. - usage = info.get("last_token_usage") or {} - cache_read = number(usage.get("cached_input_tokens")) - cache_write = number(usage.get("cache_write_input_tokens")) - # Cached tokens are included in input_tokens, and reasoning tokens - # are included in output_tokens. Keep the cache split without - # counting either category twice. - input_tokens = max(0, number(usage.get("input_tokens")) - cache_read - cache_write) - output_tokens = number(usage.get("output_tokens")) - if not (input_tokens or output_tokens or cache_read or cache_write): - continue - day = local_day(entry.get("timestamp") or path.stat().st_mtime) - add_usage(day, str(path), current_model, input_tokens, output_tokens, cache_read, cache_write) - except Exception: - continue + previous_total_usage = total_usage + day = local_day(entry.get("timestamp") or stat_result.st_mtime) + record_usage(record, day, current_model, input_tokens, output_tokens, cache_read, cache_write) + except Exception: + # Keep what was read before the failure, but caching a partial record would + # hide the rest until the file changes, so hand back one the cache refuses. + record["mtime"] = None + return record def cache_root(): @@ -343,16 +542,86 @@ def cache_root(): return root -def scan_cache_paths(): +def scan_cache_digest(): codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" # The digest covers every data path the scan reads: the codex session # roots, the opencode DB, and (via Path.home()) the pi/omp session roots. - digest = hashlib.sha1((str(Path.home()) + "\n" + str(codex_home) + "\n" + str(db)).encode("utf-8")).hexdigest()[:16] + return hashlib.sha1((str(Path.home()) + "\n" + str(codex_home) + "\n" + str(db)).encode("utf-8")).hexdigest()[:16] + + +def scan_cache_paths(): + digest = scan_cache_digest() root = cache_root() return root / f"codex-scan-{digest}.json", root / f"codex-scan-{digest}.lock" +# Records hold local days, so they are only true in the timezone that wrote them. +def file_cache_zone(): + return [os.environ.get("TZ"), list(time.tzname), time.timezone, time.altzone] + + +def file_cache_path(): + return cache_root() / f"codex-files-{scan_cache_digest()}.json" + + +def valid_file_record(value): + try: + counts = [totals for models in value["days"].values() for totals in models.values()] + return ( + type(value["mtime"]) in (int, float) and type(value["size"]) is int + and all(len(totals) == 5 and all(type(n) is int for n in totals) for totals in counts) + ) + except Exception: + return False + + +def read_file_cache(): + """Per-file totals of native sessions from earlier scans, keyed by path. + + Session files are append-only, so a file whose mtime and size both match the + record cannot have grown a turn since: its totals can be replayed instead of + read. That is what keeps a refresh off the gigabytes of history it already + counted. Unlike the aggregate scan cache these records carry no "today" and + never expire on their own -- only the file they describe or a change of + timezone can invalidate one, so midnight and the clock moving leave them true. + + pi and omp sessions stay out of it: a forked pi session repeats its parent's + messages, so they are only deduplicated across the whole tree at once. + """ + try: + payload = json.loads(file_cache_path().read_text(encoding="utf-8")) + except Exception: + return {} + if not isinstance(payload, dict) or payload.get("schemaVersion") != FILE_CACHE_SCHEMA: + return {} + if payload.get("zone") != file_cache_zone(): + return {} + files = payload.get("files") + if not isinstance(files, dict): + return {} + return {path: record for path, record in files.items() if valid_file_record(record)} + + +def write_file_cache(files): + # Only the files this scan saw are written back, so deleted sessions and + # sessions that aged out of the 30-day window leave the cache instead of + # growing it forever. + try: + write_json(file_cache_path(), {"schemaVersion": FILE_CACHE_SCHEMA, "zone": file_cache_zone(), "files": files}) + except Exception as exc: + print(f"omarchy-agent-usage-codex: could not write file cache ({exc})", file=sys.stderr) + + +def reusable_record(cached_files, key, stat_result): + record = cached_files.get(key) + if record is None: + return None + if record["mtime"] != stat_result.st_mtime or record["size"] != stat_result.st_size: + return None + return record + + def read_fresh_json(path, max_age_seconds): if max_age_seconds <= 0 or not path.exists(): return None @@ -389,9 +658,10 @@ def write_json(path, payload): # The cache payload is a versioned envelope around the local-stats dict, so a # corrupted or foreign-shaped file is a cache miss (rescan + rewrite) instead # of a crash or a garbage record. +# Version 2 invalidates totals counted before native notification deduplication. def read_cached_stats(cache_file, max_age_seconds): cached = read_fresh_json(cache_file, max_age_seconds) - if not isinstance(cached, dict) or cached.get("schemaVersion") != 1: + if not isinstance(cached, dict) or cached.get("schemaVersion") != 2: return None # today* fields only mean "today" on the day they were scanned. A cache # from another local date (midnight passed, or the clock moved) is a miss, @@ -408,7 +678,7 @@ def read_cached_stats(cache_file, max_age_seconds): def write_cached_stats(cache_file, stats): try: - write_json(cache_file, {"schemaVersion": 1, "scanDate": today, "stats": stats}) + write_json(cache_file, {"schemaVersion": 2, "scanDate": today, "stats": stats}) except Exception as exc: print(f"omarchy-agent-usage-codex: could not write usage cache ({exc})", file=sys.stderr) @@ -432,10 +702,14 @@ def local_stats(): } -def run_local_scans(): +# A forced scan (no reuse at all) reads every session file again. +def run_local_scans(max_age): + cached_files = read_file_cache() if max_age > 0 else {} + scanned_files = {} scan_pi_sessions() - scan_native_codex_sessions() + scan_native_codex_sessions(cached_files, scanned_files) complete = scan_opencode_sessions() + write_file_cache(scanned_files) return local_stats(), complete @@ -450,7 +724,7 @@ def cached_local_stats(max_age): return _cached_local_stats(max_age) except Exception as exc: print(f"omarchy-agent-usage-codex: cache unavailable ({exc}); scanning directly", file=sys.stderr) - stats, _ = run_local_scans() + stats, _ = run_local_scans(max_age) return stats @@ -466,7 +740,7 @@ def _cached_local_stats(max_age): cached = read_cached_stats(cache_file, max_age) if cached is not None: return cached - stats, complete = run_local_scans() + stats, complete = run_local_scans(max_age) # An interrupted scan still serves this run, but caching it would # suppress the missing usage for every reader until the cache expires. if complete: @@ -474,25 +748,55 @@ def _cached_local_stats(max_age): return stats +def rpc_send(proc, payload, operation): + """Write one JSON-RPC frame. Raise immediately if the app-server is gone.""" + try: + proc.stdin.write(json.dumps(payload) + "\n") + proc.stdin.flush() + except OSError as exc: + try: + proc.stdin.close() + except OSError: + pass + raise RuntimeError(f"Codex app-server exited before {operation}") from exc + + def rpc_request(proc, request_id, method, params=None, timeout=8): payload = {"id": request_id, "method": method, "params": params or {}} - proc.stdin.write(json.dumps(payload) + "\n") - proc.stdin.flush() - deadline = time.time() + timeout - while time.time() < deadline: - ready, _, _ = select.select([proc.stdout], [], [], 0.25) - if not ready: - continue - line = proc.stdout.readline() - if not line: - break - try: - message = json.loads(line) - except Exception: - continue - if message.get("id") == request_id: - return message - raise TimeoutError(method) + rpc_send(proc, payload, method) + # Read the raw fd and split lines ourselves. The app-server can emit + # notifications in the same write as a reply; a buffered readline() would + # pull them all into Python's buffer where select() cannot see them, and + # the reply would sit there until the deadline. + fd = proc.stdout.fileno() + pending = getattr(proc, "_rpc_pending", b"") + deadline = time.monotonic() + timeout + try: + while True: + while b"\n" in pending: + line, pending = pending.split(b"\n", 1) + try: + message = json.loads(line) + except Exception: + continue + if isinstance(message, dict) and message.get("id") == request_id: + return message + remaining = deadline - time.monotonic() + if remaining <= 0: + break + ready, _, _ = select.select([fd], [], [], min(0.25, remaining)) + if not ready: + if proc.poll() is not None: + raise RuntimeError(f"Codex app-server exited before {method}") + continue + chunk = os.read(fd, 65536) + if not chunk: + raise RuntimeError(f"Codex app-server exited before {method}") + pending += chunk + finally: + proc._rpc_pending = pending + # Process still up but silent: keep the method name searchable, but name the stall. + raise TimeoutError(f"Codex app-server did not answer {method}") def limit_window(window): @@ -518,6 +822,33 @@ def limit_window(window): } +def codex_rpc_help(proc, stderr_file, exc): + """Explain an app-server RPC failure: a live server's stall, a dead one's own + stderr, or the login hint when it died saying nothing.""" + try: + proc.wait(timeout=1) + except Exception: + pass + + text = str(exc).strip() + if proc.poll() is None: + # Still up: stall or unparseable payload. Not an auth problem, and any + # stderr so far is its logging, not why it stopped. + return text[:300] or AUTH_HELP + + try: + stderr_file.seek(0) + detail = stderr_file.read() + except Exception: + detail = "" + lines = [line.strip() for line in detail.splitlines() if line.strip()] + if lines: + # The fatal error comes last, after whatever the CLI logged on the way up. + return f"codex app-server exited: {' '.join(lines)[-275:]}" + # Exited cleanly with empty stderr — typical "not logged in" shape. + return AUTH_HELP + + def reset_credits(result): granted = (result.get("rateLimitResetCredits") or {}).get("credits") or [] available = [c for c in granted if isinstance(c, dict) and c.get("status") == "available"] @@ -530,33 +861,70 @@ def reset_credits(result): } +# Keyring and auto stores keep credentials outside auth.json, so only the file +# store can be known to be empty without asking Codex. +def has_codex_credentials(codex_home, env): + if env.get("CODEX_ACCESS_TOKEN"): + return True + try: + with (codex_home / "config.toml").open("rb") as handle: + config = tomllib.load(handle) + except FileNotFoundError: + config = {} + except Exception: + # A config this can't read may still name another store; let Codex say. + return True + if str(config.get("cli_auth_credentials_store") or "file") != "file": + return True + return (codex_home / "auth.json").is_file() + + def fetch_codex_rpc(home=None): result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": AUTH_HELP} env = dict(ENV, CODEX_HOME=str(home)) if home else ENV - codex = find_command("codex") + codex = find_codex_binary() if not codex: result["usageStatusText"] = "Codex unavailable" result["authHelpText"] = "codex not found in PATH" return result + # No auth.json under the file store and no CODEX_ACCESS_TOKEN means + # account/read can only fail — and starting the app-server is not free: it + # syncs the plugin list, a git fetch per refresh that also leaves + # .tmp/git-* folders behind. The home is the one being probed, so a + # signed-out secondary account reads as signed out even while ~/.codex + # holds a login. + if not has_codex_credentials(Path(env.get("CODEX_HOME") or (Path.home() / ".codex")), env): + result["usageStatusText"] = "Waiting for auth" + result["authHelpText"] = AUTH_HELP + return result + + # stderr goes to a regular file, not a PIPE: an unread pipe can fill and + # deadlock a chatty app-server. It is only read once the process has exited. + # Without temporary space the probe still runs, just without its error text. + try: + stderr_file = tempfile.TemporaryFile("w+", errors="replace") + except OSError: + stderr_file = None try: proc = subprocess.Popen( [codex, "-s", "read-only", "-a", "on-request", "app-server"], stdin=subprocess.PIPE, stdout=subprocess.PIPE, - stderr=subprocess.DEVNULL, + stderr=stderr_file or subprocess.DEVNULL, text=True, env=env, ) except Exception as exc: + if stderr_file: + stderr_file.close() result["usageStatusText"] = "Codex unavailable" result["authHelpText"] = str(exc) return result try: rpc_request(proc, 1, "initialize", {"clientInfo": {"name": "omarchy-agent-usage", "version": "1"}}, timeout=8) - proc.stdin.write(json.dumps({"method": "initialized", "params": {}}) + "\n") - proc.stdin.flush() + rpc_send(proc, {"method": "initialized", "params": {}}, "initialized") limits_msg = rpc_request(proc, 2, "account/rateLimits/read", timeout=8) # A home nobody is signed in to answers with an error rather than limits. @@ -573,6 +941,7 @@ def fetch_codex_rpc(home=None): result["authHelpText"] = message return result limits = (limits_msg.get("result") or {}).get("rateLimits") or {} + result["authHelpText"] = "" plan = limits.get("planType") or "" # The limits name the plan themselves. account/read is only a fallback @@ -583,7 +952,7 @@ def fetch_codex_rpc(home=None): account_msg = rpc_request(proc, 3, "account/read", timeout=2) account = (account_msg.get("result") or {}).get("account") or {} plan = account.get("planType") or account.get("type") or "" - except TimeoutError: + except (TimeoutError, RuntimeError): pass result["tierLabel"] = str(plan) if plan else "" @@ -599,7 +968,7 @@ def fetch_codex_rpc(home=None): result["resetCredits"] = credits except Exception as exc: result["usageStatusText"] = "Codex limits unavailable" - result["authHelpText"] = str(exc) + result["authHelpText"] = codex_rpc_help(proc, stderr_file, exc) finally: try: proc.terminate() @@ -609,6 +978,8 @@ def fetch_codex_rpc(home=None): proc.kill() except Exception: pass + if stderr_file: + stderr_file.close() return result diff --git a/bin/omarchy-agent-usage-fireworks b/bin/omarchy-agent-usage-fireworks index e2c9b632..fe04bc1d 100755 --- a/bin/omarchy-agent-usage-fireworks +++ b/bin/omarchy-agent-usage-fireworks @@ -29,7 +29,7 @@ from typing import Any AGENT_ID = "fireworks" AGENT_NAME = "Fireworks" -AUTH_HELP = "Set FIREWORKS_API_KEY, run `firectl set-api-key`, or sign in to Fireworks in opencode." +AUTH_HELP = "Set FIREWORKS_API_KEY, run `firectl set-api-key`, or sign in to Fireworks in pi or opencode." API_BASE_URL = "https://api.fireworks.ai" @@ -196,11 +196,57 @@ def read_auth_file(path: Path) -> tuple[str, str]: return api_key, account_id +def pi_auth_path() -> Path: + agent_dir = Path(os.environ.get("PI_CODING_AGENT_DIR") or (Path.home() / ".pi" / "agent")) + return agent_dir / "auth.json" + + def opencode_auth_path() -> Path: data_home = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) return data_home / "opencode" / "auth.json" +PI_ENV_REFERENCE = re.compile(r"\$(?:\{([A-Za-z_][A-Za-z0-9_]*)\}|([A-Za-z_][A-Za-z0-9_]*))") + + +def resolve_pi_key(raw: Any) -> str: + # pi credential keys are literals, $ENV_VAR/${ENV_VAR} references, or + # !command shell lookups. The collector resolves the first two; command + # lookups stay pi-only and are skipped rather than sent to the API verbatim. + key = str(raw or "").strip() + if not key or key.startswith("!"): + return "" + # Escapes survive interpolation as sentinels, then restore. + key = key.replace("$$", "\x00").replace("$!", "\x01") + + unresolved = False + + def substitute(match: re.Match[str]) -> str: + nonlocal unresolved + value = os.environ.get(match.group(1) or match.group(2), "") + if not value: + unresolved = True + return value + + key = PI_ENV_REFERENCE.sub(substitute, key) + if unresolved: + return "" + return key.replace("\x00", "$").replace("\x01", "!") + + +def read_pi_key(path: Path) -> str: + try: + parsed = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError): + return "" + entry = parsed.get("fireworks") if isinstance(parsed, dict) else None + if not isinstance(entry, dict): + return "" + if entry.get("type") not in (None, "api_key"): + return "" + return resolve_pi_key(entry.get("key")) + + def read_opencode_key(path: Path) -> str: try: parsed = json.loads(path.read_text()) @@ -227,11 +273,12 @@ def read_config() -> dict[str, Any]: def credentials(auth_path: Path, config: dict[str, Any]) -> tuple[str, str]: file_key, file_account = read_auth_file(auth_path) - # opencode is the last resort: an explicit key or a firectl login should - # win over whatever another tool happens to be signed in with. + # Agent logins are the fallback ladder below an explicit key or a firectl + # login: pi first as Omarchy's default agent, opencode as the last resort. api_key = ( str(os.environ.get("FIREWORKS_API_KEY", "")).strip() or file_key + or read_pi_key(pi_auth_path()) or read_opencode_key(opencode_auth_path()) ) account_id = ( diff --git a/bin/omarchy-agent-usage-grok b/bin/omarchy-agent-usage-grok index 09eea0e8..f2d4d7fe 100755 --- a/bin/omarchy-agent-usage-grok +++ b/bin/omarchy-agent-usage-grok @@ -30,6 +30,9 @@ CREDITS_URL = "https://cli-chat-proxy.grok.com/v1/billing?format=credits" # Grok's own CLI watches its subscription once a minute; a panel opened and # shut repeatedly reuses an answer for about that long. PROBE_REUSE_SECONDS = 60 +# Grok's docs give a saved sign-in 30 days; a refresh token older than that +# can't renew the access token, so it no longer counts as signed in. +REFRESH_LIFETIME_SECONDS = 30 * 24 * 3600 # A limits-only refresh needs fresh limits, not a fresh history scan. LIMITS_ONLY_SCAN_REUSE_SECONDS = 900 @@ -190,6 +193,22 @@ def credits_limit(data): } +# A window whose reset has passed starts over empty. A weekly one's next +# reset is a whole number of weeks later; a monthly one's isn't known. +def rolled_forward(limits): + now = datetime.now(timezone.utc) + rolled = [] + for limit in limits: + resets = parse_time(limit.get("resetsAt")) + if resets and resets <= now: + limit = dict(limit, percent=0.0, resetsAt="") + if limit.get("label") == "Weekly": + weeks = (now - resets) // timedelta(weeks=1) + 1 + limit["resetsAt"] = (resets + weeks * timedelta(weeks=1)).isoformat() + rolled.append(limit) + return rolled + + def collect_limits(home, key, force): result = {"limits": [], "usageStatusText": "", "authHelpText": AUTH_HELP, "live": False, "fetchedAtMs": 0} entry = login(home) @@ -203,17 +222,25 @@ def collect_limits(home, key, force): fetched_at = float(cached.get("fetchedAtMs") or 0) result["fetchedAtMs"] = fetched_at - # The CLI refreshes its token while it runs; one left to lapse can't be - # used until Grok starts again. expires = parse_time(entry.get("expires_at")) if expires and expires.timestamp() <= time.time(): result["limits"] = cached.get("limits") or [] - result["usageStatusText"] = "Sign-in expired" - result["authHelpText"] = "Grok's saved sign-in expired. Start Grok to refresh it." + # The access token lives six hours and Grok mints a new one from its + # refresh token whenever it starts, so a lapsed one is routine, not a + # sign-out. And while Grok hasn't run, nothing here has spent its + # allowance, so the last numbers still stand, rolled past any reset. + if entry.get("refresh_token") and time.time() - expires.timestamp() < REFRESH_LIFETIME_SECONDS: + result.update(limits=rolled_forward(result["limits"]), live=True) + if not result["limits"]: + result["usageStatusText"] = "Limits paused" + result["authHelpText"] = "Start Grok to see its limits; signing in again is not needed." + else: + result["usageStatusText"] = "Sign-in expired" + result["authHelpText"] = "Grok's saved sign-in expired. Start Grok to refresh it." return result if cached.get("limits") and not force and time.time() - fetched_at / 1000 < PROBE_REUSE_SECONDS: - result.update(limits=cached["limits"], live=True) + result.update(limits=cached["limits"], live=True, authHelpText="") return result request = urllib.request.Request(CREDITS_URL, headers={"Authorization": f"Bearer {token}", "Accept": "application/json"}) @@ -228,7 +255,7 @@ def collect_limits(home, key, force): limit = None if limit: - result.update(limits=[limit], live=True, fetchedAtMs=round(time.time() * 1000)) + result.update(limits=[limit], live=True, fetchedAtMs=round(time.time() * 1000), authHelpText="") write_json(cache, {"fetchedAtMs": result["fetchedAtMs"], "limits": result["limits"]}) else: result["limits"] = cached.get("limits") or [] diff --git a/shell/plugins/agents/Agent.qml b/shell/plugins/agents/Agent.qml index c9ba4c5b..843901fa 100644 --- a/shell/plugins/agents/Agent.qml +++ b/shell/plugins/agents/Agent.qml @@ -13,22 +13,41 @@ Item { property string path: "" property var record: null + // The text the current record was parsed from, so a reload that finds the + // file unchanged doesn't hand every listener an identical new record. + property string parsedText: "" + FileView { + id: agentFile path: root.path watchChanges: true printErrors: false onFileChanged: reload() onLoaded: root.parse(text()) - onLoadFailed: root.record = null + onLoadFailed: root.clear() + } + + // The usage update rewrites records with an atomic mv, which replaces the + // inode. Should inotify fail to rearm on the new one (watch quota, ENOSPC), + // the FileView goes quiet, so Main.qml also reloads every record after each + // update run (#9974). + function reload() { agentFile.reload() } + + function clear() { + root.parsedText = "" + root.record = null } function parse(content) { + var text = String(content || "") + if (root.record && text === root.parsedText) return try { - var parsed = JSON.parse(String(content || "")) + var parsed = JSON.parse(text) root.record = parsed && typeof parsed === "object" ? parsed : null + root.parsedText = root.record ? text : "" } catch (e) { console.warn("agents", "Ignoring bad usage record", root.path, e) - root.record = null + root.clear() } } } diff --git a/shell/plugins/agents/Main.qml b/shell/plugins/agents/Main.qml index 865a1880..7c3bb625 100644 --- a/shell/plugins/agents/Main.qml +++ b/shell/plugins/agents/Main.qml @@ -64,6 +64,13 @@ Item { onObjectRemoved: (index, object) => root.rebuildAgents() } + // Each record's FileView follows its file on its own; this catches the ones + // whose watch was lost when an update replaced the file underneath them. + function reloadRecords() { + for (var i = 0; i < agents.length; i++) + if (agents[i]) agents[i].reload() + } + function rebuildAgents() { var result = [] for (var i = 0; i < agentInstantiator.count; i++) { @@ -150,6 +157,7 @@ Item { running: false onExited: { root.rescanAgents() + root.reloadRecords() if (root.pendingUpdateKind !== "") { var kind = root.pendingUpdateKind root.pendingUpdateKind = "" diff --git a/shell/plugins/agents/Panel.qml b/shell/plugins/agents/Panel.qml index d7e4333d..a209ddfc 100644 --- a/shell/plugins/agents/Panel.qml +++ b/shell/plugins/agents/Panel.qml @@ -1333,7 +1333,8 @@ Panel { visible: !section.multi && root.otherTrouble(section.provider) !== "" width: parent.width textFormat: Text.PlainText - text: section.provider ? String(section.provider.authHelpText || "") : "" + // Shown for the status, so a record with no help to offer says the status rather than nothing. + text: section.provider ? String(section.provider.authHelpText || section.provider.usageStatusText || "") : "" color: root.urgent font.family: root.fontFamily font.pixelSize: Style.font.caption diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index c9daebba..d1081862 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -71,7 +71,7 @@ light surfaces — and the bar glyph stands in when there is none. | Collector | Limits | Local stats | |---|---|---| | `claude` | Anthropic's OAuth usage endpoint (5-hour session + 7-day weekly) | `~/.claude/projects` transcripts, opencode sessions on an Anthropic provider, plus `stats-cache.json` and `history.jsonl` as fallback | -| `codex` | The Codex app-server RPC | native Codex CLI session files (plus pi and opencode sessions) | +| `codex` | The Codex app-server RPC | native Codex CLI session files on the built-in `openai` provider (plus pi and opencode sessions) | | `grok` | The credits endpoint behind Grok's `/usage` view (the billing period's included usage) | Each session's `usage.json` (the ledger `grok usage` prints: tokens by model per finished turn), plus `summary.json` for sessions | | `fireworks` | Estimated prepaid balance: configured funding minus rated account costs | Fireworks billing API, grouped by day and model for the last 30 days | @@ -89,6 +89,8 @@ since every account shares the primary home's history. After each run, active account crosses its threshold, and re-collects the record if the active account changed. +Codex CLI will front any OpenAI-compatible backend — `--oss`, or a custom `model_provider` in `config.toml` aimed at Ollama, LM Studio, or a gateway — and those rollouts sit in the same sessions directory as OpenAI-backed ones. The Codex collector skips a rollout whose first `session_meta.model_provider` names anything but the built-in `openai` provider. Rollouts written before Codex recorded that field carry no provider and still count. + 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 @@ -98,8 +100,9 @@ sign-in; a sign-in left to lapse shows the last credits until Grok runs again. Fireworks reads `FIREWORKS_API_KEY` and `FIREWORKS_ACCOUNT_ID` first, then `~/.fireworks/auth.ini` (which `firectl set-api-key` creates), then the key -opencode stores in `~/.local/share/opencode/auth.json` when Fireworks is -signed in there. +pi stores in `~/.pi/agent/auth.json` when Fireworks is signed in there +(honoring `PI_CODING_AGENT_DIR`, and pi's literal and `$ENV_VAR` key forms), +and finally the key opencode stores in `~/.local/share/opencode/auth.json`. ### Fireworks balance diff --git a/test/shell.d/agent-usage-accounts-test.sh b/test/shell.d/agent-usage-accounts-test.sh index 72af0241..40b20290 100644 --- a/test/shell.d/agent-usage-accounts-test.sh +++ b/test/shell.d/agent-usage-accounts-test.sh @@ -114,6 +114,27 @@ PY fail "a lapsed account whose windows all reset reads as untouched" "$rested" pass "a lapsed account whose windows all reset reads as untouched" +# A cache stamped ahead of a clock that later stepped back is still what the +# account last saw, so its reset windows still read as untouched. +touch -d "@$(( $(date +%s) + 3600 ))" "$XDG_CACHE_HOME/omarchy/agent-usage/claude-limits-old.json" +rested_future=$(COLLECTOR="$ROOT/bin/omarchy-agent-usage-claude" python3 - <<'PY' +import importlib.machinery, importlib.util, io, json, os, sys + +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) +collector.urllib.request.urlopen = lambda request, timeout=None: io.BytesIO(b'{"five_hour": {"utilization": 30.0}}') +collector.scan_pi_usage = lambda age: None +collector.scan_opencode_usage = lambda age: None +sys.argv = ["omarchy-agent-usage-claude", "--force"] +collector.main() +PY +) +[[ $(jq -c '.accounts[2] | {stale, limits: [.limits[] | {label, empty: (.percent == 0)}]}' <<<"$rested_future") == '{"stale":true,"limits":[{"label":"Session (5-hour)","empty":true}]}' ]] || + fail "a future-dated cache of reset windows still reads as untouched" "$rested_future" +pass "a future-dated cache of reset windows still reads as untouched" + # Signing the primary home in to another subscription must not inherit the # last one's numbers when the first probe for the new one fails. printf '{"oauthAccount":{"accountUuid":"u-new"}}\n' >"$HOME/.claude.json" @@ -212,6 +233,7 @@ for line in sys.stdin: print(json.dumps({"id": message["id"], "result": result}), flush=True) PY chmod +x "$test_tmp/bin/codex" +touch "$HOME/.codex/auth.json" "$accounts/codex/side/auth.json" cat >"$accounts/codex.json" <"$dir/.credentials.json" + if [[ $where == "inside" ]]; then + live="$dir/.claude.json" other="$home/.claude.json" relocated="$dir" + elif [[ $where == "legacy" ]]; then + live="$dir/.config.json" other="$home/.claude.json" relocated="" + else + live="$home/.claude.json" other="$dir/.claude.json" relocated="" + fi + [[ -n $profile ]] && printf '%s' "$profile" >"$live" + [[ -n $stray ]] && printf '%s' "$stray" >"$other" + + plan=$(env -u CLAUDE_CONFIG_DIR ${relocated:+CLAUDE_CONFIG_DIR="$relocated"} HOME="$home" \ + COLLECTOR="$ROOT/bin/omarchy-agent-usage-claude" CLAUDE_DIR="$dir" python3 - <<'PY' +import importlib.machinery, importlib.util, os, pathlib + +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) + +print(collector.oauth_login(pathlib.Path(os.environ["CLAUDE_DIR"]))[-1]) +PY + ) + + printf '%s' "$plan" +} + +credentials='{"claudeAiOauth":{"accessToken":"token","expiresAt":1,"subscriptionType":"max","rateLimitTier":"default_claude_max_5x"}}' +upgraded='{"oauthAccount":{"organizationRateLimitTier":"default_claude_max_20x","userRateLimitTier":null}}' + +# The tier in the credentials is the one the token was minted with. An upgraded +# account keeps it there, so the profile the CLI refreshes decides the label. +plan=$(read_plan "$credentials" "$upgraded") +[[ $plan == "Max 20x" ]] || + fail "Claude collector labels the plan from the refreshed profile" "$plan" +pass "Claude collector labels the plan from the refreshed profile" + +# CLAUDE_CONFIG_DIR takes the profile along with the rest of the config. +plan=$(read_plan "$credentials" "$upgraded" inside) +[[ $plan == "Max 20x" ]] || + fail "Claude collector finds the profile inside a relocated config directory" "$plan" +pass "Claude collector finds the profile inside a relocated config directory" + +# A copy of the profile in the place the CLI does not read is never refreshed, +# so it must not outrank the live one, whichever way the config is laid out. +stale='{"oauthAccount":{"organizationRateLimitTier":"default_claude_max_5x","userRateLimitTier":null}}' +for where in beside inside legacy; do + plan=$(read_plan "$credentials" "$upgraded" "$where" "$stale") + [[ $plan == "Max 20x" ]] || + fail "Claude collector reads only the profile the CLI keeps" "$where -> $plan" +done +pass "Claude collector reads only the profile the CLI keeps" + +# A seat with a tier of its own is limited by that tier, not by its org's. +plan=$(read_plan "$credentials" '{"oauthAccount":{"organizationRateLimitTier":"default_claude_max_20x","userRateLimitTier":"default_claude_max_5x"}}') +[[ $plan == "Max 5x" ]] || + fail "Claude collector prefers the seat's own tier over the organization's" "$plan" +pass "Claude collector prefers the seat's own tier over the organization's" + +# Without a profile — or with one that states no tier, or a tier the label +# cannot read — the credentials are all there is, and they still answer. +for profile in '' '{}' '{"oauthAccount":{"organizationRateLimitTier":null}}' '{"oauthAccount":{"organizationRateLimitTier":"default_claude_unknown"}}'; do + plan=$(read_plan "$credentials" "$profile") + [[ $plan == "Max 5x" ]] || + fail "Claude collector falls back to the credentials when the profile names no readable tier" "$profile -> $plan" +done +pass "Claude collector falls back to the credentials when the profile names no readable tier" + +# A Team seat runs on a Max tier; the refreshed profile changes the multiplier +# but never turns the seat into Max. +team='{"claudeAiOauth":{"accessToken":"token","expiresAt":1,"subscriptionType":"team","rateLimitTier":"default_claude_max_5x"}}' +plan=$(read_plan "$team" "$upgraded") +[[ $plan == "Team 20x" ]] || + fail "Claude collector keeps a Team seat's label when the profile names its tier" "$plan" +pass "Claude collector keeps a Team seat's label when the profile names its tier" + +# A secondary account's home has its own profile; ~/.claude.json describes the +# primary account and must never lend it its tier. +second_home=$(mktemp -d "$SANDBOX/home.XXXXXX") +mkdir -p "$second_home/.claude" "$second_home/accounts/work" +printf '%s' "$upgraded" >"$second_home/.claude.json" +printf '%s' "$credentials" >"$second_home/accounts/work/.credentials.json" +plan=$(env -u CLAUDE_CONFIG_DIR HOME="$second_home" COLLECTOR="$ROOT/bin/omarchy-agent-usage-claude" \ + CLAUDE_DIR="$second_home/accounts/work" python3 - <<'PY' +import importlib.machinery, importlib.util, os, pathlib + +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) + +print(collector.oauth_login(pathlib.Path(os.environ["CLAUDE_DIR"]))[-1]) +PY +) +[[ $plan == "Max 5x" ]] || + fail "Claude collector never labels a secondary account from the primary's profile" "$plan" +pass "Claude collector never labels a secondary account from the primary's profile" diff --git a/test/shell.d/agent-usage-claude-scanner-test.sh b/test/shell.d/agent-usage-claude-scanner-test.sh index e1407f3c..533c2ad6 100644 --- a/test/shell.d/agent-usage-claude-scanner-test.sh +++ b/test/shell.d/agent-usage-claude-scanner-test.sh @@ -5,8 +5,10 @@ source "$(dirname "$0")/base-test.sh" require_command jq require_command python3 -TEST_HOME=$(mktemp -d) -trap 'rm -rf "$TEST_HOME"' EXIT +# Every fixture home lives under one scratch directory, cleaned up at exit. +SCRATCH=$(mktemp -d) +trap 'rm -rf "$SCRATCH"' EXIT +TEST_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") projects="$TEST_HOME/.claude/projects/example" mkdir -p "$projects" @@ -29,14 +31,166 @@ pass "Claude collector counts each API message once" fail "Claude collector keeps mutually exclusive token categories" "$result" pass "Claude collector keeps mutually exclusive token categories" + [[ $(jq -r '.id + "/" + .usageStatusText' <<<"$result") == "claude/Waiting for auth" ]] || fail "Claude collector identifies itself and reports missing auth" "$result" pass "Claude collector identifies itself and reports missing auth" +# Transcripts only ever grow, so a refresh reads what was appended since the +# last one and takes the rest from the scan index. +cat >>"$projects/session.jsonl" </dev/null | head -1) +[[ -n $index && $(jq -r '.files | to_entries[0].value.records | length' "$index") == "3" ]] || + fail "Claude collector keeps one record per API message in the scan index" "$(cat "$index" 2>/dev/null)" +pass "Claude collector keeps one record per API message in the scan index" + +# An unchanged file is not opened again: its records come from the index. +chmod 000 "$projects/session.jsonl" +result=$(HOME="$TEST_HOME" XDG_CACHE_HOME="$TEST_HOME/.cache" XDG_DATA_HOME="$TEST_HOME/.local/share" \ + "$ROOT/bin/omarchy-agent-usage-claude" --cache-seconds 0 2>/dev/null) +chmod 644 "$projects/session.jsonl" + +[[ $(jq -r '.todayTotalTokens' <<<"$result") == "59000" ]] || + fail "Claude collector serves unchanged transcripts from the scan index" "$result" +pass "Claude collector serves unchanged transcripts from the scan index" + +# A file that shrank was rewritten, not appended to; it is read from the start. +head -n 1 "$projects/session.jsonl" >"$projects/session.jsonl.new" +mv "$projects/session.jsonl.new" "$projects/session.jsonl" +result=$(HOME="$TEST_HOME" XDG_CACHE_HOME="$TEST_HOME/.cache" XDG_DATA_HOME="$TEST_HOME/.local/share" \ + "$ROOT/bin/omarchy-agent-usage-claude" --cache-seconds 0) + +[[ $(jq -r '.todayTotalTokens' <<<"$result") == "29090" ]] || + fail "Claude collector rescans a transcript that was rewritten" "$result" +pass "Claude collector rescans a transcript that was rewritten" + +# A transcript replaced by a larger one, or rewritten in place with more than +# it had, is a new file rather than an append: nothing of the old one stays, +# and the new one is read from its start. +INDEX_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") +index_projects="$INDEX_HOME/.claude/projects/example" +mkdir -p "$index_projects" +index_line() { + printf '{"timestamp":"%s","type":"assistant","sessionId":"s","message":{"id":"%s","role":"assistant","model":"claude-test","usage":{"input_tokens":%s,"output_tokens":0}}}\n' "$timestamp" "$1" "$2" +} +index_scan() { + HOME="$INDEX_HOME" XDG_CACHE_HOME="$INDEX_HOME/.cache" XDG_DATA_HOME="$INDEX_HOME/.local/share" \ + "$ROOT/bin/omarchy-agent-usage-claude" --cache-seconds 0 "$@" +} + +index_line old-1 1000 >"$index_projects/session.jsonl" +index_scan >/dev/null +{ index_line new-1 20; index_line new-2 30; } >"$index_projects/session.jsonl.new" +mv "$index_projects/session.jsonl.new" "$index_projects/session.jsonl" +[[ $(index_scan | jq -c '[.todayTotalTokens, .totalPrompts]') == '[50,2]' ]] || + fail "Claude collector reads a transcript replaced by a larger one from its start" "$(index_scan)" + +{ index_line other-1 7; index_line other-2 8; index_line other-3 9; } >"$index_projects/session.jsonl" +[[ $(index_scan | jq -c '[.todayTotalTokens, .totalPrompts]') == '[24,3]' ]] || + fail "Claude collector reads a transcript rewritten in place from its start" "$(index_scan)" +pass "Claude collector rereads a transcript that was replaced or rewritten" + +# A rewrite in place that grows the file and changes only its end keeps the +# first kilobytes, so the end of what was read is checked too. +for i in $(seq 1 60); do index_line "long-$i" 100; done >"$index_projects/session.jsonl" +index_scan >/dev/null +{ for i in $(seq 1 59); do index_line "long-$i" 100; done; index_line long-60 5000; index_line long-61 1; } >"$index_projects/session.jsonl" +[[ $(index_scan | jq -r '.todayTotalTokens') == "10901" ]] || + fail "Claude collector rereads a transcript whose end was rewritten in place" "$(index_scan)" +pass "Claude collector rereads a transcript whose end was rewritten in place" + +# A larger replacement that matches both ends of what was read but differs in +# between is still a different file, so it is read from its start. +for i in $(seq 1 200); do index_line "mid-$i" 100; done >"$index_projects/session.jsonl" +index_scan >/dev/null +{ for i in $(seq 1 200); do if (( i == 100 )); then index_line "mid-$i" 900; else index_line "mid-$i" 100; fi; done; index_line mid-201 1; } >"$index_projects/session.jsonl.new" +mv "$index_projects/session.jsonl.new" "$index_projects/session.jsonl" +[[ $(index_scan | jq -r '.todayTotalTokens') == "20801" ]] || + fail "Claude collector rereads a replacement whose ends match" "$(index_scan)" +pass "Claude collector rereads a replacement whose ends match" + +# The index holds local days, so a new timezone reads every transcript again +# rather than keep the days another timezone gave them. +zone_timestamp="$(date -u +%Y-%m-%d)T01:00:00Z" +printf '{"timestamp":"%s","type":"assistant","sessionId":"s","message":{"id":"zone-1","role":"assistant","model":"claude-test","usage":{"input_tokens":5,"output_tokens":0}}}\n' "$zone_timestamp" >"$index_projects/session.jsonl" +TZ=UTC index_scan >/dev/null +zone_dates=$(TZ=America/Los_Angeles index_scan | jq -r '.activeDates | join(",")') +[[ $zone_dates == "$(TZ=America/Los_Angeles date -d "$zone_timestamp" +%Y-%m-%d)" ]] || + fail "Claude collector recomputes indexed days after a timezone change" "$zone_dates" +pass "Claude collector recomputes indexed days after a timezone change" + +# A streamed response is several lines sharing one message id. The first +# line's output_tokens is a placeholder and the last line has the real count, +# so the message is counted from the line with the highest output. +STREAM_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") +stream_projects="$STREAM_HOME/.claude/projects/example" +mkdir -p "$stream_projects" +cat >"$stream_projects/session.jsonl" <>"$stream_projects/session.jsonl" <>"$projects/session.jsonl" <"$PI_HOME/.pi/agent/sessions/project/pi.jsonl" <"$PI_HOME/.omp/agent/sessions/project/omp.jsonl" <` moves the whole agent tree under profiles//, so a +# subscription spent entirely through a profile leaves the default root empty. +cat >"$PI_HOME/.omp/profiles/work/agent/sessions/project/omp-profile.jsonl" <"$TEST_HOME/bin/codex" <<'EOF' #!/bin/bash @@ -44,11 +58,13 @@ EOF chmod +x "$TEST_HOME/bin/codex" timestamp="$(date +%Y-%m-%d)T12:00:00Z" +collision_timestamp="$(date +%Y-%m-%d)T12:00:01Z" session="$TEST_HOME/.codex/sessions/$(date +%Y/%m/%d)/rollout.jsonl" cat >"$session" <"$PI_HOME/bin/rg" <"$PI_HOME/.pi/agent/sessions/project/pi.jsonl" <"$PI_HOME/.pi/agent/sessions/project/pi-fork.jsonl" <"$PI_HOME/.pi/agent/sessions/project/pi-id-collision.jsonl" <"$PI_HOME/.omp/agent/sessions/project/omp.jsonl" <` moves the whole agent tree under profiles//, so a +# subscription spent entirely through a profile leaves the default root empty. +cat >"$PI_HOME/.omp/profiles/codex/agent/sessions/project/omp-profile.jsonl" <"$GIT_HOME/.pi/agent/sessions/project/pi.jsonl" <"$GIT_HOME/.gitignore" + +result=$(HOME="$GIT_HOME" CODEX_HOME="$GIT_HOME/.codex" XDG_DATA_HOME="$GIT_HOME/.local/share" \ + PATH="$GIT_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-codex") + +[[ $(jq -r '.todayTotalTokens' <<<"$result") == "8" ]] || + fail "Codex collector counts pi sessions when HOME is a git checkout" "$result" +pass "Codex collector counts pi sessions when HOME is a git checkout" # A subscription burned entirely through opencode has no native session files; # usage must come from opencode's message database, filtered to OpenAI. -OPENCODE_HOME=$(mktemp -d) -trap 'rm -rf "$TEST_HOME" "$PI_HOME" "$OPENCODE_HOME"' EXIT -mkdir -p "$OPENCODE_HOME/bin" +OPENCODE_HOME=$(signed_in_home) cp "$TEST_HOME/bin/codex" "$OPENCODE_HOME/bin/codex" python3 - "$OPENCODE_HOME/.local/share/opencode/opencode.db" <<'PY' @@ -151,9 +217,7 @@ pass "Codex collector ignores prefix-colliding providers, user messages, and mal # A warm cache makes --limits-only cheap: local stats come from the last scan # instead of another walk over the opencode database, and --force bypasses it. -CACHE_HOME=$(mktemp -d) -trap 'rm -rf "$TEST_HOME" "$PI_HOME" "$OPENCODE_HOME" "$CACHE_HOME" "$FRESH_HOME"' EXIT -mkdir -p "$CACHE_HOME/bin" +CACHE_HOME=$(signed_in_home) cp "$TEST_HOME/bin/codex" "$CACHE_HOME/bin/codex" python3 - "$CACHE_HOME/.local/share/opencode/opencode.db" <<'PY' @@ -195,10 +259,20 @@ cache_file=$(ls "$CACHE_HOME/.cache/omarchy/agent-usage/"/codex-scan-*.json 2>/d fail "Codex collector leaves a cache file behind" "$result" [[ $(stat -c %a "$cache_file") == "644" ]] || fail "Codex collector keeps cache files readable" "$result" -[[ $(jq -r '.schemaVersion' "$cache_file") == "1" && $(jq -r '.stats.todayTotalTokens' "$cache_file") == "5" ]] || +[[ $(jq -r '.schemaVersion' "$cache_file") == "2" && $(jq -r '.stats.todayTotalTokens' "$cache_file") == "5" ]] || fail "Codex collector writes a versioned cache envelope" "$result" pass "Codex collector writes a local-stats cache on first scan" +# A still-fresh cache from before native notification deduplication must not +# restore inflated counts, even when only quota limits were requested. +jq '.schemaVersion = 1 | .stats.todayTotalTokens = 999' "$cache_file" >"$CACHE_HOME/old-cache.json" +mv "$CACHE_HOME/old-cache.json" "$cache_file" +result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$CACHE_HOME/.local/share" \ + PATH="$CACHE_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-codex" --limits-only) +[[ $(jq -r '.todayTotalTokens' <<<"$result") == "5" && $(jq -r '.schemaVersion' "$cache_file") == "2" ]] || + fail "Codex collector invalidates pre-deduplication cached totals" "$result" +pass "Codex collector invalidates pre-deduplication cached totals" + # A corrupt-but-parseable cache (wrong shape) is a cache miss: rescan and # rewrite instead of emitting a garbage record. printf '[]' >"$cache_file" @@ -207,7 +281,7 @@ result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CAC [[ $(jq -r '.todayTotalTokens' <<<"$result") == "5" ]] || fail "Codex collector recovers from a corrupt cache file" "$result" -[[ $(jq -r '.schemaVersion' "$cache_file") == "1" ]] || +[[ $(jq -r '.schemaVersion' "$cache_file") == "2" ]] || fail "Codex collector rewrites the cache after a corrupt read" "$result" pass "Codex collector recovers from a corrupt cache file" @@ -415,9 +489,7 @@ result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CAC pass "Codex collector treats a future-dated cache as a miss" # First --limits-only on a machine with no cache falls back to a full scan. -FRESH_HOME=$(mktemp -d) -trap 'rm -rf "$TEST_HOME" "$PI_HOME" "$OPENCODE_HOME" "$CACHE_HOME" "$FRESH_HOME"' EXIT -mkdir -p "$FRESH_HOME/bin" +FRESH_HOME=$(signed_in_home) cp "$TEST_HOME/bin/codex" "$FRESH_HOME/bin/codex" python3 - "$FRESH_HOME/.local/share/opencode/opencode.db" <<'PY' @@ -455,9 +527,7 @@ pass "Codex collector --limits-only falls back to a full scan without a cache" # A malformed opencode row must not abort the scan: json_valid() guards the # parse, so the good rows are still counted. Real opencode data also stores # compact JSON, so one row is serialized compactly here on purpose. -MALFORMED_HOME=$(mktemp -d) -trap 'rm -rf "$TEST_HOME" "$PI_HOME" "$OPENCODE_HOME" "$CACHE_HOME" "$FRESH_HOME" "$MALFORMED_HOME"' EXIT -mkdir -p "$MALFORMED_HOME/bin" +MALFORMED_HOME=$(signed_in_home) cp "$TEST_HOME/bin/codex" "$MALFORMED_HOME/bin/codex" python3 - "$MALFORMED_HOME/.local/share/opencode/opencode.db" <<'PY' @@ -509,9 +579,7 @@ result=$(HOME="$MALFORMED_HOME" CODEX_HOME="$MALFORMED_HOME/.codex" XDG_CACHE_HO pass "Codex collector counts good opencode rows past malformed ones" # An unwritable cache must not kill the collector: the record is the contract. -UNWRITABLE_HOME=$(mktemp -d) -trap 'rm -rf "$TEST_HOME" "$PI_HOME" "$OPENCODE_HOME" "$CACHE_HOME" "$FRESH_HOME" "$MALFORMED_HOME" "$UNWRITABLE_HOME"' EXIT -mkdir -p "$UNWRITABLE_HOME/bin" +UNWRITABLE_HOME=$(signed_in_home) cp "$TEST_HOME/bin/codex" "$UNWRITABLE_HOME/bin/codex" python3 - "$UNWRITABLE_HOME/.local/share/opencode/opencode.db" <<'PY' @@ -551,9 +619,7 @@ pass "Codex collector still prints a complete record when the cache is unwritabl # A scan cut short by a database error (schema migration, transient lock, # corruption) must not be cached as the whole story, or the missing usage # would be suppressed for every reader until the cache expires. -INTERRUPTED_HOME=$(mktemp -d) -trap 'rm -rf "$TEST_HOME" "$PI_HOME" "$OPENCODE_HOME" "$CACHE_HOME" "$FRESH_HOME" "$MALFORMED_HOME" "$UNWRITABLE_HOME" "$INTERRUPTED_HOME"' EXIT -mkdir -p "$INTERRUPTED_HOME/bin" +INTERRUPTED_HOME=$(signed_in_home) cp "$TEST_HOME/bin/codex" "$INTERRUPTED_HOME/bin/codex" # A database without the message table makes the scan fail mid-flight. @@ -640,3 +706,537 @@ result=$(HOME="$TEST_HOME" CODEX_HOME="$TEST_HOME/.codex" XDG_DATA_HOME="$TEST_H [[ $(jq -r '.usageStatusText' <<<"$result") == "Waiting for auth" ]] || fail "Codex collector reports a missing sign-in as one" "$result" pass "Codex collector reports a missing sign-in as one" + +# The app-server batches notifications with replies in one write. A reply +# that shares a write with a notification must not be stranded in a read +# buffer, and bytes left over from one request must carry into the next. +BATCHED_HOME=$(signed_in_home) +cat >"$BATCHED_HOME/bin/codex" <<'EOF' +#!/bin/bash + +while read -r request; do + id=$(jq -r '.id // empty' <<<"$request") + method=$(jq -r '.method // empty' <<<"$request") + + case "$method" in + initialize) + # One write: this reply and a trailing notification. + printf '%s\n%s\n' \ + "$(jq -cn --argjson id "$id" '{id: $id, result: {}}')" \ + '{"method":"remoteControl/status/changed","params":{"status":"disabled"}}' + ;; + account/rateLimits/read) + # One write: a notification ahead of this reply. + printf '%s\n%s\n' \ + '{"method":"account/updated","params":{"authMode":"chatgpt","planType":"plus"}}' \ + "$(jq -cn --argjson id "$id" '{id: $id, result: {rateLimits: {planType: "plus", primary: {usedPercent: 5, windowDurationMins: 300, resetsAt: 1790659194}}}}')" + ;; + esac +done +EOF +chmod +x "$BATCHED_HOME/bin/codex" + +result=$(HOME="$BATCHED_HOME" CODEX_HOME="$BATCHED_HOME/.codex" XDG_CACHE_HOME="$BATCHED_HOME/.cache" XDG_DATA_HOME="$BATCHED_HOME/.local/share" \ + PATH="$BATCHED_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-codex") + +[[ $(jq -c '{tierLabel, usageStatusText, limits: [.limits[] | {label, percent}]}' <<<"$result") == '{"tierLabel":"plus","usageStatusText":"","limits":[{"label":"5h window","percent":0.05}]}' ]] || + fail "Codex collector reads replies batched with notifications" "$result" +pass "Codex collector reads replies batched with notifications" + +# Without temporary space the probe still runs; it only loses Codex's error +# text, so stderr goes nowhere instead of failing the probe. +NO_TEMP_PYTHON="$BATCHED_HOME/python" +mkdir -p "$NO_TEMP_PYTHON" +cat >"$NO_TEMP_PYTHON/sitecustomize.py" <<'EOF' +import tempfile + +def no_space(*args, **kwargs): + raise OSError(28, "No space left on device") + +tempfile.TemporaryFile = no_space +EOF + +result=$(HOME="$BATCHED_HOME" CODEX_HOME="$BATCHED_HOME/.codex" XDG_CACHE_HOME="$BATCHED_HOME/.cache" XDG_DATA_HOME="$BATCHED_HOME/.local/share" \ + PYTHONPATH="$NO_TEMP_PYTHON" PATH="$BATCHED_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-codex" --force) + +[[ $(jq -c '[.limits[] | .percent]' <<<"$result") == '[0.05]' ]] || + fail "Codex collector probes without temporary space" "$result" +pass "Codex collector probes without temporary space" + +# The lazy launcher at ~/.local/bin/codex runs `mise use -g` when executed, so +# a read-only usage probe on a machine without Codex must never spawn it. A +# private tools dir keeps a real codex or mise on the host out of the probe: +# PATH holds only the interpreter and file tools the collector may exec. +LAUNCH_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") +SAFE_PATH="$LAUNCH_HOME/tools" +mkdir -p "$SAFE_PATH" +for tool in python3 rg; do + if command -v "$tool" >/dev/null; then + ln -s "$(command -v "$tool")" "$SAFE_PATH/$tool" + fi +done +mkdir -p "$LAUNCH_HOME/bin" "$LAUNCH_HOME/.local/bin" "$LAUNCH_HOME/.codex" +touch "$LAUNCH_HOME/.codex/auth.json" +cat >"$LAUNCH_HOME/.local/bin/codex" <<'LAUNCHER' +#!/bin/bash +export MISE_MINIMUM_RELEASE_AGE=0 +mise use -g --quiet "npm:@openai/codex" || exit 1 +exec mise x "npm:@openai/codex" -- codex "$@" +LAUNCHER +chmod +x "$LAUNCH_HOME/.local/bin/codex" +cat >"$LAUNCH_HOME/bin/mise" <<'STUB' +#!/bin/bash +printf '%s\n' "$*" >>"$MISE_CALLS_FILE" +exit 1 +STUB +chmod +x "$LAUNCH_HOME/bin/mise" + +result=$(HOME="$LAUNCH_HOME" CODEX_HOME="$LAUNCH_HOME/.codex" MISE_CALLS_FILE="$LAUNCH_HOME/mise-calls" XDG_DATA_HOME="$LAUNCH_HOME/.local/share" \ + PATH="$LAUNCH_HOME/bin:$SAFE_PATH" "$ROOT/bin/omarchy-agent-usage-codex" --limits-only) + +[[ $(jq -r '.usageStatusText' <<<"$result") == "Codex unavailable" ]] || + fail "Codex collector reports Codex unavailable when only the launcher exists" "$result" +# The launcher would log `use -g ...` through the mise stub if it ever ran; +# `which codex` is the only permitted call. +[[ ! -s $LAUNCH_HOME/mise-calls || $(cat "$LAUNCH_HOME/mise-calls") == "which codex" ]] || + fail "Codex collector must not execute the lazy codex launcher" "$(cat "$LAUNCH_HOME/mise-calls" 2>/dev/null)" +pass "Codex collector resolves through mise which instead of running the launcher" + +# When mise reports an installed binary, that binary is probed, not the +# launcher that shadows it on PATH. +cp "$TEST_HOME/bin/codex" "$LAUNCH_HOME/real-codex" +cat >"$LAUNCH_HOME/bin/mise" <>"$LAUNCH_HOME/mise-calls" +echo "$LAUNCH_HOME/real-codex" +STUB + +rm -f "$LAUNCH_HOME/mise-calls" +result=$(HOME="$LAUNCH_HOME" CODEX_HOME="$LAUNCH_HOME/.codex" CODEX_ARGS_FILE="$LAUNCH_HOME/codex-args" XDG_DATA_HOME="$LAUNCH_HOME/.local/share" \ + PATH="$LAUNCH_HOME/bin:$SAFE_PATH" "$ROOT/bin/omarchy-agent-usage-codex" --limits-only) + +[[ -f $LAUNCH_HOME/codex-args && $(cat "$LAUNCH_HOME/mise-calls") == "which codex" ]] || + fail "Codex collector probes the binary mise which reports" "$result" +pass "Codex collector probes the mise-resolved binary" + +# A symlink at the launcher path is the user's own binary, so it is probed +# directly without asking mise at all. +LINK_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") +mkdir -p "$LINK_HOME/bin" "$LINK_HOME/.local/bin" "$LINK_HOME/.codex" +touch "$LINK_HOME/.codex/auth.json" +ln -s "$TEST_HOME/bin/codex" "$LINK_HOME/.local/bin/codex" +cat >"$LINK_HOME/bin/mise" <>"$LINK_HOME/mise-calls" +echo "$LINK_HOME/real-codex" +STUB +chmod +x "$LINK_HOME/bin/mise" + +result=$(HOME="$LINK_HOME" CODEX_HOME="$LINK_HOME/.codex" CODEX_ARGS_FILE="$LINK_HOME/codex-args" XDG_DATA_HOME="$LINK_HOME/.local/share" \ + PATH="$LINK_HOME/bin:$SAFE_PATH" "$ROOT/bin/omarchy-agent-usage-codex" --limits-only) + +[[ -f $LINK_HOME/codex-args && ! -s $LINK_HOME/mise-calls ]] || + fail "Codex collector probes a symlinked codex without invoking mise" "$result" +pass "Codex collector probes a user-owned symlink at the launcher path" + +# A mise shim is a symlink to the mise binary itself, so it resolves to mise, +# not to an installed codex — running it would exec `mise x` and install the +# tool. It must be treated as lazy despite being a symlink. +SHIM_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") +mkdir -p "$SHIM_HOME/bin" "$SHIM_HOME/.local/share/mise/shims" +cat >"$SHIM_HOME/bin/mise" <>"$SHIM_HOME/mise-calls" +exit 1 +STUB +chmod +x "$SHIM_HOME/bin/mise" +ln -s "$SHIM_HOME/bin/mise" "$SHIM_HOME/.local/share/mise/shims/codex" + +result=$(HOME="$SHIM_HOME" CODEX_HOME="$SHIM_HOME/.codex" MISE_CALLS_FILE="$SHIM_HOME/mise-calls" XDG_DATA_HOME="$SHIM_HOME/.local/share" \ + PATH="$SHIM_HOME/bin:$SAFE_PATH" "$ROOT/bin/omarchy-agent-usage-codex" --limits-only) + +[[ $(jq -r '.usageStatusText' <<<"$result") == "Codex unavailable" ]] || + fail "Codex collector reports Codex unavailable when only a mise shim exists" "$result" +[[ ! -s $SHIM_HOME/mise-calls || $(cat "$SHIM_HOME/mise-calls") == "which codex" ]] || + fail "Codex collector must not execute a mise shim" "$(cat "$SHIM_HOME/mise-calls" 2>/dev/null)" +pass "Codex collector treats a shim symlink to mise as lazy" + +# Without Codex credentials, account/read can only fail — and starting the +# app-server is not free: it syncs the plugin list, a git fetch per refresh +# that leaves ~/.codex/.tmp/git-* folders behind. The collector must not +# spawn codex at all. +NOAUTH_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") +mkdir -p "$NOAUTH_HOME/bin" +cp "$TEST_HOME/bin/codex" "$NOAUTH_HOME/bin/codex" + +result=$(HOME="$NOAUTH_HOME" CODEX_HOME="$NOAUTH_HOME/.codex" CODEX_ARGS_FILE="$NOAUTH_HOME/codex-args" XDG_DATA_HOME="$NOAUTH_HOME/.local/share" \ + PATH="$NOAUTH_HOME/bin:$PATH" env -u OPENAI_API_KEY -u CODEX_API_KEY -u CODEX_ACCESS_TOKEN "$ROOT/bin/omarchy-agent-usage-codex" --limits-only) + +[[ ! -e $NOAUTH_HOME/codex-args ]] || + fail "Codex collector does not spawn app-server without credentials" "$result" +[[ $(jq -r '.usageStatusText' <<<"$result") == "Waiting for auth" ]] || + fail "Codex collector waits for auth without credentials" "$result" +pass "Codex collector does not spawn app-server without credentials" + +# A non-file credentials store keeps credentials outside auth.json, so the +# probe must still run. +KEYRING_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") +mkdir -p "$KEYRING_HOME/bin" "$KEYRING_HOME/.codex" +cp "$TEST_HOME/bin/codex" "$KEYRING_HOME/bin/codex" +printf 'cli_auth_credentials_store = "keyring"\n' >"$KEYRING_HOME/.codex/config.toml" + +result=$(HOME="$KEYRING_HOME" CODEX_HOME="$KEYRING_HOME/.codex" CODEX_ARGS_FILE="$KEYRING_HOME/codex-args" XDG_DATA_HOME="$KEYRING_HOME/.local/share" \ + PATH="$KEYRING_HOME/bin:$PATH" env -u OPENAI_API_KEY -u CODEX_API_KEY -u CODEX_ACCESS_TOKEN "$ROOT/bin/omarchy-agent-usage-codex" --limits-only) + +[[ -e $KEYRING_HOME/codex-args ]] || + fail "Codex collector probes the app-server when the keyring store is configured" "$result" +pass "Codex collector probes the app-server when the keyring store is configured" + +# The setting is TOML, so a single-quoted value names the keyring just as well, +# while an explicit file store with no auth.json is still signed out. +for store in "'keyring'" '"file"'; do + rm -f "$KEYRING_HOME/codex-args" + printf 'model = "gpt-5"\ncli_auth_credentials_store = %s\n' "$store" >"$KEYRING_HOME/.codex/config.toml" + HOME="$KEYRING_HOME" CODEX_HOME="$KEYRING_HOME/.codex" CODEX_ARGS_FILE="$KEYRING_HOME/codex-args" XDG_DATA_HOME="$KEYRING_HOME/.local/share" \ + PATH="$KEYRING_HOME/bin:$PATH" env -u OPENAI_API_KEY -u CODEX_API_KEY -u CODEX_ACCESS_TOKEN "$ROOT/bin/omarchy-agent-usage-codex" --limits-only >/dev/null + if [[ $store == "'keyring'" ]]; then + [[ -e $KEYRING_HOME/codex-args ]] || fail "Codex collector reads a single-quoted keyring store as TOML" + else + [[ ! -e $KEYRING_HOME/codex-args ]] || fail "Codex collector does not probe an empty file store" + fi +done +pass "Codex collector reads the credentials store as TOML" + +# A CODEX_ACCESS_TOKEN is credentials even without auth.json. +TOKEN_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") +mkdir -p "$TOKEN_HOME/bin" "$TOKEN_HOME/.codex" +cp "$TEST_HOME/bin/codex" "$TOKEN_HOME/bin/codex" + +result=$(HOME="$TOKEN_HOME" CODEX_HOME="$TOKEN_HOME/.codex" CODEX_ARGS_FILE="$TOKEN_HOME/codex-args" XDG_DATA_HOME="$TOKEN_HOME/.local/share" \ + PATH="$TOKEN_HOME/bin:$PATH" env -u OPENAI_API_KEY -u CODEX_API_KEY CODEX_ACCESS_TOKEN=x "$ROOT/bin/omarchy-agent-usage-codex" --limits-only) + +[[ -e $TOKEN_HOME/codex-args ]] || + fail "Codex collector probes the app-server when CODEX_ACCESS_TOKEN is set" "$result" +pass "Codex collector probes the app-server when CODEX_ACCESS_TOKEN is set" + +# An API key alone does not log the app-server in. +APIKEY_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") +mkdir -p "$APIKEY_HOME/bin" "$APIKEY_HOME/.codex" +cp "$TEST_HOME/bin/codex" "$APIKEY_HOME/bin/codex" + +result=$(HOME="$APIKEY_HOME" CODEX_HOME="$APIKEY_HOME/.codex" CODEX_ARGS_FILE="$APIKEY_HOME/codex-args" XDG_DATA_HOME="$APIKEY_HOME/.local/share" \ + PATH="$APIKEY_HOME/bin:$PATH" env -u CODEX_API_KEY -u CODEX_ACCESS_TOKEN OPENAI_API_KEY=x "$ROOT/bin/omarchy-agent-usage-codex" --limits-only) + +[[ ! -e $APIKEY_HOME/codex-args ]] || + fail "Codex collector does not spawn app-server for an API key alone" "$result" +pass "Codex collector does not spawn app-server for an API key alone" + +# A codex that exits before speaking the protocol (rejected flag, crash, etc.) +# must not leave the panel showing the bare RPC method name "initialize". +EXIT_HOME=$(signed_in_home) +cat >"$EXIT_HOME/bin/codex" <<'EOF' +#!/bin/bash +echo "error: invalid value 'untrusted' for '--ask-for-approval '" >&2 +echo " [possible values: on-request, never]" >&2 +exit 2 +EOF +chmod +x "$EXIT_HOME/bin/codex" + +result=$(HOME="$EXIT_HOME" CODEX_HOME="$EXIT_HOME/.codex" XDG_CACHE_HOME="$EXIT_HOME/.cache" XDG_DATA_HOME="$EXIT_HOME/.local/share" \ + PATH="$EXIT_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-codex" --limits-only) + +[[ $(jq -r '.usageStatusText' <<<"$result") == "Codex limits unavailable" ]] || + fail "Codex collector reports limits unavailable when app-server rejects argv" "$result" +help=$(jq -r '.authHelpText' <<<"$result") +[[ $help == *"invalid value 'untrusted'"* ]] || + fail "Codex collector surfaces the CLI's own error" "$result" +[[ $help != "initialize" ]] || + fail "Codex collector must not leak the raw RPC method name" "$result" +pass "Codex collector surfaces a rejected app-server call instead of the RPC method name" + +# EOF on stdout during initialize (process died) is reported as an exit, not a bare method. +DEAD_HOME=$(signed_in_home) +cat >"$DEAD_HOME/bin/codex" <<'EOF' +#!/bin/bash +exec 1>&- +exec sleep 30 +EOF +chmod +x "$DEAD_HOME/bin/codex" + +result=$(HOME="$DEAD_HOME" CODEX_HOME="$DEAD_HOME/.codex" XDG_CACHE_HOME="$DEAD_HOME/.cache" XDG_DATA_HOME="$DEAD_HOME/.local/share" \ + PATH="$DEAD_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-codex" --limits-only) + +help=$(jq -r '.authHelpText' <<<"$result") +[[ $help == "Codex app-server exited before initialize" ]] || + fail "Codex collector identifies an app-server that exits during startup" "$result" +pass "Codex collector identifies an app-server that exits during startup" + +# Failure while sending the initialized notification after initialize answered. +HALF_HOME=$(signed_in_home) +cat >"$HALF_HOME/bin/codex" <<'EOF' +#!/bin/bash +read -r request +exec 0<&- +jq -cn --argjson id "$(jq -r '.id' <<<"$request")" '{id: $id, result: {}}' +exec sleep 30 +EOF +chmod +x "$HALF_HOME/bin/codex" + +result=$(HOME="$HALF_HOME" CODEX_HOME="$HALF_HOME/.codex" XDG_CACHE_HOME="$HALF_HOME/.cache" XDG_DATA_HOME="$HALF_HOME/.local/share" \ + PATH="$HALF_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-codex" --limits-only) + +help=$(jq -r '.authHelpText' <<<"$result") +[[ $help == "Codex app-server exited before initialized" ]] || + fail "Codex collector translates failure to send the initialized notification" "$result" +pass "Codex collector translates failure to send the initialized notification" + +# Silent clean exit with no stderr: fall back to the login hint. +SILENT_HOME=$(signed_in_home) +cat >"$SILENT_HOME/bin/codex" <<'EOF' +#!/bin/bash +exit 0 +EOF +chmod +x "$SILENT_HOME/bin/codex" + +result=$(HOME="$SILENT_HOME" CODEX_HOME="$SILENT_HOME/.codex" XDG_CACHE_HOME="$SILENT_HOME/.cache" XDG_DATA_HOME="$SILENT_HOME/.local/share" \ + PATH="$SILENT_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-codex" --limits-only) + +help=$(jq -r '.authHelpText' <<<"$result") +[[ $help == "Run \`codex login\` to authenticate." ]] || + fail "Codex collector falls back to the login hint when the CLI is silent" "$result" +pass "Codex collector falls back to the login hint when the app-server says nothing" + +# Live app-server that stalls on account/rateLimits/read: keep a clear stall message, not login. +STALL_HOME=$(signed_in_home) +cat >"$STALL_HOME/bin/codex" <<'EOF' +#!/bin/bash +while read -r request; do + id=$(jq -r '.id // empty' <<<"$request") + method=$(jq -r '.method // empty' <<<"$request") + case "$method" in + initialize) jq -cn --argjson id "$id" '{id: $id, result: {}}' ;; + account/rateLimits/read) : ;; + esac +done +EOF +chmod +x "$STALL_HOME/bin/codex" + +result=$(HOME="$STALL_HOME" CODEX_HOME="$STALL_HOME/.codex" XDG_CACHE_HOME="$STALL_HOME/.cache" XDG_DATA_HOME="$STALL_HOME/.local/share" \ + PATH="$STALL_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-codex" --limits-only) + +help=$(jq -r '.authHelpText' <<<"$result") +[[ $help != "Run \`codex login\` to authenticate." ]] || + fail "Codex collector must not blame auth when the app-server is merely stalled" "$result" +[[ $help == "Codex app-server did not answer account/rateLimits/read" ]] || + fail "Codex collector names the stalled RPC method clearly" "$result" +[[ $help != "account/rateLimits/read" && $help != "initialize" ]] || + fail "Codex collector must not leak a bare method name" "$result" +pass "Codex collector names a stalled RPC instead of leaking the method name" + +# A stalled app-server that has logged to stderr is still running: its logging +# is not why it stopped, and it has not exited. +NOISY_HOME=$(signed_in_home) +sed 's/^while read/echo "WARN codex_core: startup notice" >\&2\nwhile read/' "$STALL_HOME/bin/codex" >"$NOISY_HOME/bin/codex" +chmod +x "$NOISY_HOME/bin/codex" + +result=$(HOME="$NOISY_HOME" CODEX_HOME="$NOISY_HOME/.codex" XDG_CACHE_HOME="$NOISY_HOME/.cache" XDG_DATA_HOME="$NOISY_HOME/.local/share" \ + PATH="$NOISY_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-codex" --limits-only) + +[[ $(jq -r '.authHelpText' <<<"$result") == "Codex app-server did not answer account/rateLimits/read" ]] || + fail "Codex collector reports a stall, not an exit, when a live app-server has logged" "$result" +pass "Codex collector reports a stall, not an exit, when a live app-server has logged" + +# A CLI that logs plenty before failing must still show the failure, not the logging. +CHATTY_HOME=$(signed_in_home) +cat >"$CHATTY_HOME/bin/codex" <<'EOF' +#!/bin/bash +for i in {1..20}; do echo "WARN codex_core::config: ignoring unknown key number $i" >&2; done +echo "error: failed to start app-server" >&2 +exit 1 +EOF +chmod +x "$CHATTY_HOME/bin/codex" + +result=$(HOME="$CHATTY_HOME" CODEX_HOME="$CHATTY_HOME/.codex" XDG_CACHE_HOME="$CHATTY_HOME/.cache" XDG_DATA_HOME="$CHATTY_HOME/.local/share" \ + PATH="$CHATTY_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-codex" --limits-only) + +[[ $(jq -r '.authHelpText' <<<"$result") == "codex app-server exited: "*"error: failed to start app-server" ]] || + fail "Codex collector keeps the CLI's final error past its startup logging" "$result" +pass "Codex collector keeps the CLI's final error past its startup logging" + +# Codex CLI can front any OpenAI-compatible backend (`--oss`, or a custom +# model_provider in config.toml). Those rollouts land in the same sessions +# directory but spend a local box or a third party, never this subscription. +PROVIDER_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") +mkdir -p "$PROVIDER_HOME/bin" "$PROVIDER_HOME/.codex/sessions/$(date +%Y/%m/%d)" +cp "$TEST_HOME/bin/codex" "$PROVIDER_HOME/bin/codex" + +provider_session() { + local name=$1 meta=$2 model=$3 input=$4 output=$5 + local file="$PROVIDER_HOME/.codex/sessions/$(date +%Y/%m/%d)/rollout-$name.jsonl" + [[ -n $meta ]] && echo "{\"timestamp\":\"$timestamp\",\"type\":\"session_meta\",\"payload\":{\"model_provider\":\"$meta\"}}" >"$file" + cat >>"$file" <"$session" <"$incremental_session" <"$incremental_pi" </dev/null | head -n 1) +[[ -n $file_cache && -s $file_cache ]] || + fail "Codex collector writes a per-file cache on first scan" +[[ $(jq --arg path "$incremental_session" -r '.files[$path].days | length' "$file_cache") == "1" ]] || + fail "Codex collector records the native session file it read" "$(cat "$file_cache")" +# A forked pi session repeats its parent's messages, so pi sessions are only +# deduplicated across the whole tree at once, never replayed per file. +[[ $(jq --arg path "$incremental_pi" -r '.files | has($path)' "$file_cache") == "false" ]] || + fail "Codex collector keeps pi sessions out of the per-file cache" "$(cat "$file_cache")" +pass "Codex collector records per-file totals as it scans" + +# Unreadable, but unchanged in mtime and size: a run that still reports the +# same totals can only have replayed them. +chmod 000 "$incremental_session" +expire_scan_cache +result=$(run_incremental) +[[ $(jq -r '.todayTotalTokens' <<<"$result") == "20" ]] || + fail "Codex collector rereads session files it has already counted" "$result" +pass "Codex collector replays unchanged session files instead of rereading them" + +# --force means the history itself is re-read, per-file records included. +result=$(run_incremental --force) +[[ $(jq -r '.todayTotalTokens' <<<"$result") == "5" ]] || + fail "Codex collector --force reuses per-file records" "$result" +pass "Codex collector --force rereads the history" + +chmod 644 "$incremental_session" +# The --force run above could not read the session, so cache it again first, +# or the append below is read from scratch whether or not the cache noticed it. +expire_scan_cache +run_incremental >/dev/null +cat >>"$incremental_session" <"$zone_session" </dev/null +expire_scan_cache +TZ=America/Los_Angeles run_incremental >/dev/null +file_cache=$(ls "$INCREMENTAL_HOME/.cache/omarchy/agent-usage/"codex-files-*.json | head -n 1) +[[ $(jq --arg path "$zone_session" -r '.files[$path].days | keys[0]' "$file_cache") == "$(TZ=America/Los_Angeles date -d "$zone_timestamp" +%Y-%m-%d)" ]] || + fail "Codex collector keeps the old timezone's days after a timezone change" "$(cat "$file_cache")" +pass "Codex collector re-reads per-file records after a timezone change" + +# A line that stops the read keeps the usage read before it, as it always has. +broken_session="$INCREMENTAL_HOME/.codex/sessions/broken.jsonl" +cat >"$broken_session" <= 0, `found ${signature} in ${file}`) + const end = source.indexOf('\n }', start) + assert(end > start, `found the end of ${signature} in ${file}`) + return source.slice(start, end + '\n }'.length) +} + +const main = fs.readFileSync(root + '/shell/plugins/agents/Main.qml', 'utf8') +const agent = fs.readFileSync(root + '/shell/plugins/agents/Agent.qml', 'utf8') + +const updateStart = main.indexOf(' id: updateProcess') +const exited = main.slice(main.indexOf(' onExited: {', updateStart), main.indexOf('\n }', updateStart)) +assert(exited.includes('root.reloadRecords()'), 'the update process reloads every record when it exits') + +const panel = { agents: [] } +vm.createContext(panel) +vm.runInContext(extract(main, 'function reloadRecords() {', 'Main.qml'), panel) +const reloaded = [] +panel.agents = [{ reload: () => reloaded.push('claude') }, null, { reload: () => reloaded.push('codex') }] +panel.reloadRecords() +assertDeepEqual(reloaded, ['claude', 'codex'], 'reloadRecords reloads each agent record') + +assert(agent.includes('function reload() { agentFile.reload() }'), 'an agent reloads through its FileView') +const record = { record: null, parsedText: '', path: 'codex.json', console: { warn() {} } } +record.root = record +vm.createContext(record) +vm.runInContext(extract(agent, 'function clear() {', 'Agent.qml') + '\n' + extract(agent, 'function parse(content) {', 'Agent.qml'), record) + +record.parse('{"id":"codex","todayPrompts":1}') +const first = record.record +assertEqual(first.todayPrompts, 1, 'a record is parsed from its file') +record.parse('{"id":"codex","todayPrompts":1}') +assert(record.record === first, 'reloading an unchanged file keeps the same record') +record.parse('{"id":"codex","todayPrompts":2}') +assertEqual(record.record.todayPrompts, 2, 'a changed file replaces the record') +record.parse('not json') +assertEqual(record.record, null, 'a bad file clears the record') +record.parse('{"id":"codex","todayPrompts":2}') +assertEqual(record.record.todayPrompts, 2, 'the record comes back once the file is good again') +JS diff --git a/test/shell.d/agent-usage-fireworks-scanner-test.sh b/test/shell.d/agent-usage-fireworks-scanner-test.sh index a7d3ccf8..6440ce11 100755 --- a/test/shell.d/agent-usage-fireworks-scanner-test.sh +++ b/test/shell.d/agent-usage-fireworks-scanner-test.sh @@ -27,13 +27,13 @@ EOF # Without credentials the collector must still print a full, hidden-by-default # record: the update runner writes whatever valid JSON appears on stdout. no_key=$(HOME="$TEST_HOME" XDG_CONFIG_HOME="$TEST_HOME/.config" XDG_DATA_HOME="$TEST_HOME/.local/share" \ - FIREWORKS_API_KEY="" FIREWORKS_AUTH_PATH="$TEST_HOME/missing.ini" "$ROOT/bin/omarchy-agent-usage-fireworks") + PI_CODING_AGENT_DIR="$TEST_HOME/.pi/agent" FIREWORKS_API_KEY="" FIREWORKS_AUTH_PATH="$TEST_HOME/missing.ini" "$ROOT/bin/omarchy-agent-usage-fireworks") [[ $(jq -r '.id + ":" + (.ready | tostring) + ":" + (.hasPromptStats | tostring)' <<<"$no_key") == "fireworks:false:false" ]] || fail "Fireworks collector prints a valid record without credentials" "$no_key" pass "Fireworks collector prints a valid record without credentials" -result=$(python3 - "$ROOT/bin/omarchy-agent-usage-fireworks" "$auth_file" "$TEST_HOME/.config" "$TEST_HOME/.local/share" <<'PY' +result=$(python3 - "$ROOT/bin/omarchy-agent-usage-fireworks" "$auth_file" "$TEST_HOME/.config" "$TEST_HOME/.local/share" "$TEST_HOME/.pi/agent" <<'PY' import importlib.machinery import importlib.util import json @@ -95,7 +95,11 @@ summary["apiKey"] = api_key summary["accountId"] = account_id summary["money"] = float(scanner.money_value({"units": "12", "nanos": 430000000})) -# The opencode key only wins when no explicit key or firectl login exists. +# The opencode key only wins when no explicit key, firectl login, or pi +# login exists. Pin the pi dir first so a developer's real ~/.pi/agent +# credentials cannot leak into the fixture run. +pi_dir = Path(sys.argv[5]) +os.environ["PI_CODING_AGENT_DIR"] = str(pi_dir) data_home = Path(os.environ["XDG_DATA_HOME"]) opencode_auth = data_home / "opencode" / "auth.json" opencode_auth.parent.mkdir(parents=True, exist_ok=True) @@ -105,6 +109,22 @@ opencode_key, _ = scanner.credentials(Path("/nonexistent/auth.ini"), {}) firectl_key, _ = scanner.credentials(auth_path, {}) summary["opencodeFallback"] = opencode_key == "fw_opencode" and firectl_key == "fw_test" +# A pi login is the next rung below firectl, above opencode. +pi_dir.mkdir(parents=True, exist_ok=True) +(pi_dir / "auth.json").write_text(json.dumps({"fireworks": {"type": "api_key", "key": "fw_pi"}})) +pi_key, _ = scanner.credentials(Path("/nonexistent/auth.ini"), {}) +firectl_still_wins, _ = scanner.credentials(auth_path, {}) +summary["piFallback"] = pi_key == "fw_pi" and firectl_still_wins == "fw_test" + +# pi key syntax: $ENV references resolve, missing ones stay unresolved, and +# !command lookups are skipped and fall through to opencode. +os.environ["FW_TEST_ENV_KEY"] = "fw_from_env" +summary["piEnvKey"] = scanner.resolve_pi_key("$FW_TEST_ENV_KEY") == "fw_from_env" +summary["piMissingEnvKey"] = scanner.resolve_pi_key("$FW_DEFINITELY_MISSING") == "" +(pi_dir / "auth.json").write_text(json.dumps({"fireworks": {"type": "api_key", "key": "!secret-tool lookup fw"}})) +command_key, _ = scanner.credentials(Path("/nonexistent/auth.ini"), {}) +summary["piCommandKeySkipped"] = command_key == "fw_opencode" + class WorkingClient: def __init__(self, api_key, base_url): pass @@ -242,3 +262,19 @@ pass "Fireworks collector dates buckets by local day east of Greenwich" [[ $(jq -r '.opencodeFallback' <<<"$result") == "true" ]] || fail "Fireworks collector falls back to the opencode key last" "$result" pass "Fireworks collector falls back to the opencode key last" + +[[ $(jq -r '.piFallback' <<<"$result") == "true" ]] || + fail "Fireworks collector falls back to the pi login before opencode" "$result" +pass "Fireworks collector falls back to the pi login before opencode" + +[[ $(jq -r '.piEnvKey' <<<"$result") == "true" ]] || + fail "Fireworks collector resolves env-var pi keys" "$result" +pass "Fireworks collector resolves env-var pi keys" + +[[ $(jq -r '.piMissingEnvKey' <<<"$result") == "true" ]] || + fail "Fireworks collector treats missing env references as unresolved" "$result" +pass "Fireworks collector treats missing env references as unresolved" + +[[ $(jq -r '.piCommandKeySkipped' <<<"$result") == "true" ]] || + fail "Fireworks collector skips command-lookup pi keys and falls through" "$result" +pass "Fireworks collector skips command-lookup pi keys and falls through" diff --git a/test/shell.d/agent-usage-grok-test.sh b/test/shell.d/agent-usage-grok-test.sh index 6a31d7e9..eeaf985e 100755 --- a/test/shell.d/agent-usage-grok-test.sh +++ b/test/shell.d/agent-usage-grok-test.sh @@ -74,6 +74,10 @@ record=$(collect) [[ $(jq -c '{ready, tierLabel, stale: .limitsStale, label: .limits[0].label, percent: .limits[0].percent}' <<<"$record") == '{"ready":true,"tierLabel":"X Premium+","stale":false,"label":"Weekly","percent":0.42}' ]] || fail "Grok's plan and credits come from its own home" "$record" [[ -n $(jq -r '.limits[0].resetsAt' <<<"$record") ]] || fail "the credits window says when the period ends" "$record" +[[ $(jq -r '.authHelpText' <<<"$record") == "" ]] || fail "a successful Grok probe drops the login hint" "$record" +record=$(collect) +[[ $(jq -c '{percent: .limits[0].percent, help: .authHelpText}' <<<"$record") == '{"percent":0.42,"help":""}' ]] || + fail "reused Grok limits drop the login hint" "$record" pass "Grok's plan and credits come from its own home" signed_in "$HOME/.grok" token-main u-main "$past" "X Premium+" @@ -82,6 +86,44 @@ record=$(collect) fail "a lapsed sign-in keeps the last credits and says so" "$record" pass "a lapsed sign-in keeps the last credits and says so" +# A lapsed access token with a refresh token is routine: Grok renews it when it +# starts, and while it hasn't run nothing here has spent its allowance, so the +# last numbers stand as current, with no sign-in asked for. +jq '.[].refresh_token = "r"' "$HOME/.grok/auth.json" >"$test_tmp/auth.json" +mv "$test_tmp/auth.json" "$HOME/.grok/auth.json" +record=$(collect) +[[ $(jq -c '{usageStatusText, first: .limits[0].percent, stale: .limitsStale}' <<<"$record") == '{"usageStatusText":"","first":0.42,"stale":false}' ]] || + fail "a lapsed access token with a refresh token keeps showing the last numbers" "$record" +pass "a lapsed access token with a refresh token keeps showing the last numbers" + +# A week that reset while Grok sat idle starts over at 0%, a week later. +cache=$(ls "$XDG_CACHE_HOME"/omarchy/agent-usage/grok-limits-*.json | head -1) +reset_past=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) - dt.timedelta(hours=1)).isoformat())') +jq --arg at "$reset_past" '.limits[0].resetsAt = $at' "$cache" >"$test_tmp/cache.json" +mv "$test_tmp/cache.json" "$cache" +record=$(collect) +next_reset=$(jq -r '.limits[0].resetsAt' <<<"$record") +[[ $(jq -r '.limits[0].percent' <<<"$record") == 0.0 ]] && + python3 -c 'import datetime as dt, sys; n = dt.datetime.fromisoformat(sys.argv[1]); d = n - dt.datetime.now(dt.timezone.utc); sys.exit(0 if dt.timedelta(days=6) < d < dt.timedelta(days=7) else 1)' "$next_reset" || + fail "a week that reset while Grok was idle starts over a week later" "$record" +pass "a week that reset while Grok was idle starts over a week later" + +# Without cached numbers there's nothing to show, so it says how to get them. +rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/grok-limits-*.json +record=$(collect) +[[ $(jq -c '{usageStatusText, limits}' <<<"$record") == '{"usageStatusText":"Limits paused","limits":[]}' ]] || + fail "a lapsed token with nothing cached says to start Grok" "$record" +pass "a lapsed token with nothing cached says to start Grok" + +# A refresh token past Grok's 30-day sign-in can't renew anything: signed out. +long_ago=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) - dt.timedelta(days=31)).isoformat())') +jq --arg at "$long_ago" '.[].expires_at = $at' "$HOME/.grok/auth.json" >"$test_tmp/auth.json" +mv "$test_tmp/auth.json" "$HOME/.grok/auth.json" +record=$(collect) +[[ $(jq -r '.usageStatusText' <<<"$record") == "Sign-in expired" ]] || + fail "a sign-in lapsed past Grok's 30 days asks for a sign-in" "$record" +pass "a sign-in lapsed past Grok's 30 days asks for a sign-in" + # A period with nothing used yet comes without a percentage. signed_in "$HOME/.grok" token-fresh u-main "$future" "X Premium+" rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/grok-limits-*.json diff --git a/test/shell.d/agents-panel-test.sh b/test/shell.d/agents-panel-test.sh index b8e6d8d2..d9e742a2 100755 --- a/test/shell.d/agents-panel-test.sh +++ b/test/shell.d/agents-panel-test.sh @@ -31,6 +31,7 @@ assert(/opacity: stale \? 0\.5 : 1\.0/.test(panelSource) && /"As of " \+ root\.f assert(/onPickingChanged: resetKeys\(\)/.test(panelSource), 'the cursor starts over when the agent list comes or goes') assert(!/t === "a" \|\| t === "A"/.test(panelSource), 'adding an account has no hotkey; the + is the way in') assert(/if \(!accounts\[a\]\.active\) \{/.test(panelSource) && /if \(row\.length > 0\) rows\.push\(row\)/.test(panelSource), 'the active account with nothing to fix is not a keyboard stop') +assert(/text: section\.provider \? String\(section\.provider\.authHelpText \|\| section\.provider\.usageStatusText \|\| ""\) : ""/.test(panelSource), 'the trouble line falls back to the usage status when a record has no auth help') const mainSource = fs.readFileSync(root + '/shell/plugins/agents/Main.qml', 'utf8') assert(/var from = Math\.min\(0\.8, \(threshold - 15\) \/ 100\)/.test(mainSource), 'faster checks start 15 points below the switch threshold') assert(/rows\.push\(\[\{ kind: "provider", index: p \}\]\)/.test(panelSource), "every agent's header is a keyboard stop")