536 lines
21 KiB
Python
Executable File
536 lines
21 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
"""Print one display-ready Codex usage record as JSON.
|
|
|
|
Local statistics come from the Codex CLI's own session files; the plan and the
|
|
rate limits come from the Codex app-server over JSON-RPC. The usage panel reads
|
|
only the JSON this prints.
|
|
|
|
Adapted from Omarchy (bin/omarchy-agent-usage-codex).
|
|
|
|
Copyright (c) David Heinemeier Hansson
|
|
|
|
Permission is hereby granted, free of charge, to any person obtaining a copy of
|
|
this software and associated documentation files (the "Software"), to deal in
|
|
the Software without restriction, including without limitation the rights to
|
|
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
|
|
the Software, and to permit persons to whom the Software is furnished to do so,
|
|
subject to the following conditions:
|
|
|
|
The above copyright notice and this permission notice shall be included in all
|
|
copies or substantial portions of the Software.
|
|
|
|
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
|
|
FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
|
|
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
|
|
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
|
|
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
|
|
|
── Secrets ───────────────────────────────────────────────────────────────────
|
|
|
|
This never reads the Codex credential store. The app-server is asked for the
|
|
account and its limits over a pipe, and it authenticates itself; no token
|
|
reaches this process, its argv, or the record. The app-server is spawned
|
|
read-only and untrusted so a usage query can never change anything.
|
|
|
|
── Seams ─────────────────────────────────────────────────────────────────────
|
|
|
|
CODEX_HOME Codex's home directory (~/.codex)
|
|
PANAMA_AGENT_CODEX_BIN the codex binary, for tests
|
|
PANAMA_AGENT_USAGE_CACHE the session-scan cache directory
|
|
|
|
The output path is not a seam here: this prints, and panama-agent-usage-update
|
|
owns the atomic write into $XDG_STATE_HOME/panama/agents/usage/.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import fcntl
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import select
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from datetime import datetime, timedelta, timezone
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
AGENT_ID = "codex"
|
|
AGENT_NAME = "Codex"
|
|
AUTH_HELP = "Run `codex login` to authenticate."
|
|
|
|
# A scan this recent is only reused to dedup concurrent collector runs (the
|
|
# fan-out backgrounds one per agent). --limits-only promises only fresh limits,
|
|
# so it may reuse a scan for far longer.
|
|
SCAN_REUSE_SECONDS = 20
|
|
LIMITS_ONLY_REUSE_SECONDS = 900
|
|
|
|
# Sessions older than this are not what anyone is looking at, and walking them
|
|
# on every refresh is the difference between a scan and a stall.
|
|
SESSION_MAX_AGE_SECONDS = 30 * 24 * 60 * 60
|
|
|
|
|
|
def expand_path(value: str) -> Path:
|
|
return Path(os.path.expandvars(os.path.expanduser(value)))
|
|
|
|
|
|
def codex_home() -> Path:
|
|
return expand_path(os.environ.get("CODEX_HOME") or "~/.codex")
|
|
|
|
|
|
def runtime_env() -> dict[str, str]:
|
|
# Codex is commonly a user-level npm or mise install, and a collector run
|
|
# from the shell's environment does not always inherit those directories.
|
|
home = str(Path.home())
|
|
path_parts = [
|
|
os.environ.get("PATH", ""),
|
|
f"{home}/.local/bin",
|
|
f"{home}/.npm-global/bin",
|
|
f"{home}/.local/share/mise/shims",
|
|
]
|
|
env = os.environ.copy()
|
|
env["PATH"] = os.pathsep.join(part for part in path_parts if part)
|
|
return env
|
|
|
|
|
|
ENV = runtime_env()
|
|
|
|
|
|
def find_codex() -> str | None:
|
|
override = os.environ.get("PANAMA_AGENT_CODEX_BIN")
|
|
if override:
|
|
return override if os.access(override, os.X_OK) else None
|
|
return shutil.which("codex", path=ENV.get("PATH"))
|
|
|
|
|
|
def cache_root() -> Path:
|
|
override = os.environ.get("PANAMA_AGENT_USAGE_CACHE")
|
|
root = expand_path(override) if override else (
|
|
Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "panama" / "agent-usage"
|
|
)
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
return root
|
|
|
|
|
|
def number(value: Any) -> int:
|
|
try:
|
|
return int(value or 0)
|
|
except Exception:
|
|
return 0
|
|
|
|
|
|
def model_name(raw: Any) -> str:
|
|
value = str(raw or "codex")
|
|
return value if value else "codex"
|
|
|
|
|
|
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)]
|
|
|
|
|
|
def local_day(value: Any) -> str:
|
|
if value is None:
|
|
return today
|
|
if isinstance(value, (int, float)):
|
|
# Codex timestamps are usually seconds; anything larger is milliseconds.
|
|
if value > 10_000_000_000:
|
|
value = value / 1000
|
|
return datetime.fromtimestamp(value).strftime("%Y-%m-%d")
|
|
text = str(value)
|
|
try:
|
|
parsed = datetime.fromisoformat(text[:-1] + "+00:00" if text.endswith("Z") else text)
|
|
if parsed.tzinfo is not None:
|
|
parsed = parsed.astimezone()
|
|
return parsed.strftime("%Y-%m-%d")
|
|
except Exception:
|
|
return today
|
|
|
|
|
|
class Tally:
|
|
"""Everything the session scan accumulates, in one place."""
|
|
|
|
def __init__(self) -> None:
|
|
self.recent = {day: {"date": day, "messageCount": 0} for day in recent_dates}
|
|
self.today_tokens_by_model: dict[str, int] = {}
|
|
self.model_usage: dict[str, dict[str, int]] = {}
|
|
self.today_sessions: set[str] = set()
|
|
self.total_sessions: set[str] = set()
|
|
self.active_days: set[str] = set()
|
|
self.today_prompts = 0
|
|
self.today_total_tokens = 0
|
|
self.total_prompts = 0
|
|
|
|
def add(self, day: str, session_key: str, model: str,
|
|
input_tokens: int, output_tokens: int, cache_read: int, cache_write: int) -> None:
|
|
total = input_tokens + output_tokens + cache_read + cache_write
|
|
self.total_prompts += 1
|
|
self.total_sessions.add(session_key)
|
|
self.active_days.add(day)
|
|
|
|
bucket = self.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 self.recent:
|
|
# recentDays.messageCount is a token total despite the legacy name.
|
|
self.recent[day]["messageCount"] += total
|
|
|
|
if day == today:
|
|
self.today_prompts += 1
|
|
self.today_sessions.add(session_key)
|
|
self.today_total_tokens += total
|
|
self.today_tokens_by_model[model] = self.today_tokens_by_model.get(model, 0) + total
|
|
|
|
def stats(self) -> dict[str, Any]:
|
|
return {
|
|
"todayPrompts": self.today_prompts,
|
|
"todaySessions": len(self.today_sessions),
|
|
"todayTotalTokens": self.today_total_tokens,
|
|
"todayTokensByModel": self.today_tokens_by_model,
|
|
"recentDays": [self.recent[day] for day in recent_dates],
|
|
"totalPrompts": self.total_prompts,
|
|
"totalSessions": len(self.total_sessions),
|
|
"activeDays": len(self.active_days),
|
|
"activeDates": sorted(self.active_days),
|
|
"modelUsage": self.model_usage,
|
|
}
|
|
|
|
|
|
def scan_native_sessions(tally: Tally) -> None:
|
|
home = codex_home()
|
|
roots = [home / "sessions", home / "archived_sessions"]
|
|
files = []
|
|
cutoff = time.time() - SESSION_MAX_AGE_SECONDS
|
|
for root in roots:
|
|
if not root.exists():
|
|
continue
|
|
for path in root.rglob("*.jsonl"):
|
|
try:
|
|
if path.stat().st_mtime >= cutoff:
|
|
files.append(path)
|
|
except OSError:
|
|
pass
|
|
|
|
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:
|
|
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 only the
|
|
# last turn is counted.
|
|
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 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)
|
|
tally.add(day, str(path), current_model,
|
|
input_tokens, output_tokens, cache_read, cache_write)
|
|
except Exception:
|
|
continue
|
|
|
|
|
|
# ───────────────────────────────────────────────────────────── scan cache ──
|
|
|
|
|
|
def scan_cache_paths() -> tuple[Path, Path]:
|
|
digest = hashlib.sha1(str(codex_home()).encode("utf-8")).hexdigest()[:16]
|
|
root = cache_root()
|
|
return root / f"codex-scan-{digest}.json", root / f"codex-scan-{digest}.lock"
|
|
|
|
|
|
def read_fresh_json(path: Path, max_age_seconds: float) -> Any:
|
|
if max_age_seconds <= 0 or not path.exists():
|
|
return None
|
|
try:
|
|
# A negative age means the mtime is in the future: the clock moved
|
|
# backwards since the write, so the cache cannot be trusted.
|
|
age = time.time() - path.stat().st_mtime
|
|
if 0 <= age <= max_age_seconds:
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
except Exception:
|
|
return None
|
|
return None
|
|
|
|
|
|
def write_json(path: Path, payload: dict[str, Any]) -> None:
|
|
# A temp name unique to this writer, not derived from the target: several
|
|
# collectors can run at once, and a shared temp path means the second
|
|
# replace finds the first one's file already moved away.
|
|
handle_fd, tmp_name = tempfile.mkstemp(dir=path.parent, prefix=path.name + ".", suffix=".tmp")
|
|
tmp = Path(tmp_name)
|
|
try:
|
|
with os.fdopen(handle_fd, "w", encoding="utf-8") as handle:
|
|
handle.write(json.dumps(payload, separators=(",", ":"), sort_keys=True) + "\n")
|
|
tmp.chmod(0o644)
|
|
tmp.replace(path)
|
|
except BaseException:
|
|
tmp.unlink(missing_ok=True)
|
|
raise
|
|
|
|
|
|
# The cache payload is a versioned envelope around the stats dict, so a
|
|
# corrupted or foreign-shaped file is a miss (rescan and rewrite) rather than a
|
|
# crash or a garbage record.
|
|
def read_cached_stats(cache_file: Path, max_age_seconds: float) -> dict[str, Any] | None:
|
|
cached = read_fresh_json(cache_file, max_age_seconds)
|
|
if not isinstance(cached, dict) or cached.get("schemaVersion") != 1:
|
|
return None
|
|
# today* fields only mean "today" on the day they were scanned.
|
|
if cached.get("scanDate") != today:
|
|
return None
|
|
stats = cached.get("stats")
|
|
if not isinstance(stats, dict):
|
|
return None
|
|
if not all(key in stats for key in ("todayPrompts", "todayTotalTokens", "recentDays", "activeDates", "modelUsage")):
|
|
return None
|
|
return stats
|
|
|
|
|
|
def local_stats(max_age: float) -> dict[str, Any]:
|
|
"""Local stats, with the cache as a pure optimization.
|
|
|
|
A cache-layer failure must never take the collector down: it degrades to a
|
|
direct scan. The printed record is the contract; the cache is not.
|
|
"""
|
|
try:
|
|
cache_file, lock_file = scan_cache_paths()
|
|
|
|
cached = read_cached_stats(cache_file, max_age)
|
|
if cached is not None:
|
|
return cached
|
|
|
|
with lock_file.open("w") as lock:
|
|
fcntl.flock(lock, fcntl.LOCK_EX)
|
|
cached = read_cached_stats(cache_file, max_age)
|
|
if cached is not None:
|
|
return cached
|
|
tally = Tally()
|
|
scan_native_sessions(tally)
|
|
stats = tally.stats()
|
|
write_json(cache_file, {"schemaVersion": 1, "scanDate": today, "stats": stats})
|
|
return stats
|
|
except Exception as exc:
|
|
print(f"panama-agent-usage-codex: cache unavailable ({exc}); scanning directly", file=sys.stderr)
|
|
tally = Tally()
|
|
scan_native_sessions(tally)
|
|
return tally.stats()
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────── app-server ──
|
|
|
|
|
|
def rpc_request(proc: subprocess.Popen, request_id: int, method: str,
|
|
params: dict[str, Any] | None = None, timeout: float = 8) -> dict[str, Any]:
|
|
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)
|
|
|
|
|
|
def limit_window(window: Any, prefix: str = "") -> dict[str, Any] | None:
|
|
if not isinstance(window, dict):
|
|
return None
|
|
used = window.get("usedPercent")
|
|
if used is None:
|
|
return None
|
|
mins = number(window.get("windowDurationMins"))
|
|
if mins == 10080:
|
|
label = "Weekly (7-day)"
|
|
elif mins == 300:
|
|
label = "Session (5-hour)"
|
|
elif mins and mins % 60 == 0:
|
|
label = f"{mins // 60}h window"
|
|
elif mins:
|
|
label = f"{mins}m window"
|
|
else:
|
|
label = "Limit"
|
|
if prefix:
|
|
label = f"{prefix} {label}"
|
|
reset = window.get("resetsAt")
|
|
try:
|
|
percent = min(1.0, max(0.0, float(used) / 100.0))
|
|
except Exception:
|
|
return None
|
|
return {
|
|
"label": label,
|
|
"percent": percent,
|
|
"resetsAt": datetime.fromtimestamp(number(reset), timezone.utc).isoformat() if reset else "",
|
|
}
|
|
|
|
|
|
# The account's own windows come back under `primary` and `secondary`.
|
|
# `rateLimitsByLimitId` repeats those under the account's limit id and adds the
|
|
# model-scoped ones beside them -- a window only one model draws from, named by
|
|
# `limitName`. An entry with no name is the account limit again, so only the
|
|
# named ones are worth a row of their own.
|
|
def scoped_limits(limits: dict[str, Any]) -> list[dict[str, Any]]:
|
|
by_id = limits.get("rateLimitsByLimitId")
|
|
if not isinstance(by_id, dict):
|
|
return []
|
|
out: list[dict[str, Any]] = []
|
|
for entry in by_id.values():
|
|
if not isinstance(entry, dict):
|
|
continue
|
|
name = str(entry.get("limitName") or "").strip()
|
|
if name == "":
|
|
continue
|
|
for window in (entry.get("primary"), entry.get("secondary")):
|
|
row = limit_window(window, name)
|
|
if row:
|
|
out.append(row)
|
|
return out
|
|
|
|
|
|
def fetch_rpc() -> dict[str, Any]:
|
|
result: dict[str, Any] = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": AUTH_HELP}
|
|
codex = find_codex()
|
|
if not codex:
|
|
result["usageStatusText"] = "Codex unavailable"
|
|
result["authHelpText"] = "codex was not found on PATH."
|
|
return result
|
|
|
|
try:
|
|
# Read-only, and never asking for approval: a usage query has no
|
|
# business being able to change anything on this machine, and nothing is
|
|
# watching a prompt it might raise. (`-a untrusted` was the flag Omarchy
|
|
# used; codex 0.149 takes on-request or never.)
|
|
proc = subprocess.Popen(
|
|
[codex, "-s", "read-only", "-a", "never", "app-server"],
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.DEVNULL,
|
|
text=True,
|
|
env=ENV,
|
|
)
|
|
except Exception as exc:
|
|
result["usageStatusText"] = "Codex unavailable"
|
|
result["authHelpText"] = str(exc)
|
|
return result
|
|
|
|
try:
|
|
rpc_request(proc, 1, "initialize",
|
|
{"clientInfo": {"name": "panama-agent-usage", "version": "1"}}, timeout=8)
|
|
proc.stdin.write(json.dumps({"method": "initialized", "params": {}}) + "\n")
|
|
proc.stdin.flush()
|
|
account_msg = rpc_request(proc, 2, "account/read", timeout=4)
|
|
limits_msg = rpc_request(proc, 3, "account/rateLimits/read", timeout=4)
|
|
|
|
account = (account_msg.get("result") or {}).get("account") or {}
|
|
payload = limits_msg.get("result") or {}
|
|
limits = payload.get("rateLimits") or {}
|
|
plan = limits.get("planType") or account.get("planType") or account.get("type") or ""
|
|
result["tierLabel"] = str(plan) if plan else ""
|
|
|
|
for window in (limits.get("primary"), limits.get("secondary")):
|
|
entry = limit_window(window)
|
|
if entry:
|
|
result["limits"].append(entry)
|
|
result["limits"].extend(scoped_limits(payload))
|
|
except TimeoutError as exc:
|
|
result["usageStatusText"] = "Codex limits unavailable"
|
|
result["authHelpText"] = (
|
|
f"The Codex app-server did not answer `{exc}` in time. "
|
|
"Start Codex once, or run `codex login`, and the limits come back."
|
|
)
|
|
except Exception as exc:
|
|
result["usageStatusText"] = "Codex limits unavailable"
|
|
result["authHelpText"] = str(exc) or "The Codex app-server could not be reached."
|
|
finally:
|
|
try:
|
|
proc.terminate()
|
|
proc.wait(timeout=1)
|
|
except Exception:
|
|
try:
|
|
proc.kill()
|
|
except Exception:
|
|
pass
|
|
return result
|
|
|
|
|
|
# ───────────────────────────────────────────────────────────────── record ──
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Print the Codex usage record as JSON")
|
|
parser.add_argument("--force", action="store_true",
|
|
help="rescan sessions and re-probe limits, ignoring caches")
|
|
parser.add_argument("--limits-only", action="store_true",
|
|
help="reuse any recent session scan; only the limits probe must be fresh")
|
|
args = parser.parse_args()
|
|
|
|
max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS)
|
|
stats = local_stats(max_age)
|
|
rpc = fetch_rpc()
|
|
|
|
record = {
|
|
"schemaVersion": 1,
|
|
"id": AGENT_ID,
|
|
"name": AGENT_NAME,
|
|
"updatedAt": datetime.now(timezone.utc).isoformat(),
|
|
# Honest about having nothing to say: a Codex that is not signed in and
|
|
# has never run leaves the panel's tab empty rather than showing zeros.
|
|
"ready": number(stats.get("totalPrompts")) > 0 or len(rpc["limits"]) > 0,
|
|
"hasLocalStats": True,
|
|
}
|
|
record.update(stats)
|
|
record.update(rpc)
|
|
print(json.dumps(record, separators=(",", ":"), sort_keys=True))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|