ZacharyZhang-NY/omarchycn
2
2
Fork 0
Code Issues 1 Pull Requests 2 Actions Packages Projects Releases 3 Wiki Activity

add model-usage plugin

Browse Source
This commit is contained in:
Ryan Hughes
2026-05-27 16:44:24 -04:00
parent dfc450f2cb
commit 752ecaf6e7
12 changed files with 3011 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
shell/plugins/model-usage/scripts/claude_usage_scanner.py
Executable
+238
View File
@@ -0,0 +1,238 @@
#!/usr/bin/env python3
"""Stream Claude Code project JSONL files and emit compact usage stats.
This replaces the QML-side `rg --json ... | StdioCollector` path, which can
materialize 100MB+ of ripgrep JSON in the Quickshell process. The helper keeps
that work in a short-lived Python process, parses line-by-line, and returns a
single compact JSON object that matches the fields Claude.qml expects.
"""
from __future__ import annotations
import argparse
import datetime as dt
import fcntl
import hashlib
import json
import os
import sys
import time
from pathlib import Path
from typing import Any
def expand_path(value: str) -> Path:
return Path(os.path.expandvars(os.path.expanduser(value))).resolve()
def date_string(value: dt.date) -> str:
return value.strftime("%Y-%m-%d")
def recent_date_strings() -> list[str]:
today = dt.datetime.now().date()
return [date_string(today - dt.timedelta(days=offset)) for offset in range(6, -1, -1)]
def local_date_string() -> str:
return date_string(dt.datetime.now().date())
def local_date_from_timestamp(value: Any) -> str:
if value is None:
return local_date_string()
if isinstance(value, (int, float)):
try:
seconds = float(value) / 1000.0 if float(value) > 10_000_000_000 else float(value)
return date_string(dt.datetime.fromtimestamp(seconds).date())
except Exception:
return local_date_string()
raw = str(value).strip()
if not raw:
return local_date_string()
# Claude JSONL timestamps are usually ISO-8601. Python accepts offsets but
# not a trailing Z until we normalize it to +00:00.
try:
parsed = dt.datetime.fromisoformat(raw.replace("Z", "+00:00"))
if parsed.tzinfo is not None:
parsed = parsed.astimezone()
return date_string(parsed.date())
except Exception:
return local_date_string()
def usage_token(usage: dict[str, Any], snake_key: str, camel_key: str) -> int:
value = usage.get(snake_key, usage.get(camel_key, 0))
try:
return round(float(value or 0))
except Exception:
return 0
def empty_bucket() -> dict[str, int]:
return {
"inputTokens": 0,
"outputTokens": 0,
"cacheReadInputTokens": 0,
"cacheCreationInputTokens": 0,
}
def iter_jsonl_files(projects_path: Path):
if not projects_path.is_dir():
return
yield from projects_path.rglob("*.jsonl")
def scan(projects_path: Path) -> 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()
sessions: set[str] = set()
today_sessions: set[str] = set()
today_tokens: dict[str, int] = {}
usage_by_model: dict[str, dict[str, int]] = {}
prompts = 0
today_prompt_count = 0
today_token_total = 0
malformed_lines = 0
scanned_files = 0
for path in iter_jsonl_files(projects_path) or []:
scanned_files += 1
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. Matches the old rg
# search and keeps files with unrelated lines inexpensive.
if '"usage":' not in line:
continue
try:
entry = json.loads(line)
except Exception:
malformed_lines += 1
continue
message = entry.get("message") if isinstance(entry.get("message"), dict) else {}
if entry.get("type") != "assistant" and message.get("role") != "assistant":
continue
usage = message.get("usage") or entry.get("usage")
if not isinstance(usage, dict):
continue
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)
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)
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:
# Preserve existing QML behavior: recentDays.messageCount
# is actually a token total, despite the legacy name.
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)
recent_days = [recent[day] for day in recent_dates]
return {
"schemaVersion": 1,
"todayPrompts": today_prompt_count,
"todaySessions": len(today_sessions),
"todayTotalTokens": today_token_total,
"todayTokensByModel": today_tokens,
"recentDays": recent_days,
"modelUsage": usage_by_model,
"totalPrompts": prompts,
"totalSessions": len(sessions),
"dailyActivity": recent_days,
"scannedFiles": scanned_files,
"malformedLines": malformed_lines,
}
def cache_paths(projects_path: Path) -> tuple[Path, Path]:
cache_root = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) / "omarchy" / "model-usage"
cache_root.mkdir(parents=True, exist_ok=True)
digest = hashlib.sha1(str(projects_path).encode("utf-8")).hexdigest()[:16]
return cache_root / f"claude-projects-{digest}.json", cache_root / f"claude-projects-{digest}.lock"
def read_fresh_cache(path: Path, max_age_seconds: int) -> str | None:
if max_age_seconds <= 0 or not path.exists():
return None
try:
if time.time() - path.stat().st_mtime <= max_age_seconds:
return path.read_text(encoding="utf-8")
except Exception:
return None
return None
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("projects_path", nargs="?", default="~/.claude/projects")
parser.add_argument("--cache-seconds", type=int, default=20)
parser.add_argument("--force", action="store_true")
args = parser.parse_args()
projects_path = expand_path(args.projects_path)
cache_file, lock_file = cache_paths(projects_path)
if not args.force:
cached = read_fresh_cache(cache_file, args.cache_seconds)
if cached is not None:
print(cached, end="" if cached.endswith("\n") else "\n")
return 0
with lock_file.open("w") as lock:
fcntl.flock(lock, fcntl.LOCK_EX)
if not args.force:
cached = read_fresh_cache(cache_file, args.cache_seconds)
if cached is not None:
print(cached, end="" if cached.endswith("\n") else "\n")
return 0
summary = scan(projects_path)
output = json.dumps(summary, separators=(",", ":"), sort_keys=True) + "\n"
tmp = cache_file.with_suffix(".json.tmp")
tmp.write_text(output, encoding="utf-8")
tmp.replace(cache_file)
print(output, end="")
return 0
if __name__ == "__main__":
raise SystemExit(main())
shell/plugins/model-usage/scripts/codex_usage_scanner.py
+344
View File
@@ -0,0 +1,344 @@
import json
import os
import select
import shutil
import subprocess
import sys
import time
from datetime import datetime, timedelta, timezone
from pathlib import Path
def local_day(value):
if value is None:
return datetime.now().strftime("%Y-%m-%d")
if isinstance(value, (int, float)):
# pi message timestamps are milliseconds; Codex timestamps are usually seconds.
if value > 10_000_000_000:
value = value / 1000
return datetime.fromtimestamp(value).strftime("%Y-%m-%d")
text = str(value)
try:
if text.endswith("Z"):
dt = datetime.fromisoformat(text[:-1] + "+00:00")
else:
dt = datetime.fromisoformat(text)
if dt.tzinfo is not None:
dt = dt.astimezone()
return dt.strftime("%Y-%m-%d")
except Exception:
return datetime.now().strftime("%Y-%m-%d")
def number(value):
try:
return int(value or 0)
except Exception:
return 0
def model_name(raw):
value = str(raw or "codex")
return value if value else "codex"
def runtime_env():
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_command(name):
return shutil.which(name, path=ENV.get("PATH"))
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)]
recent_set = set(recent_dates)
recent = {day: {"date": day, "messageCount": 0} for day in recent_dates}
today_tokens_by_model = {}
model_usage = {}
sessions_by_day = {day: set() for day in recent_dates}
today_sessions = set()
today_prompts = 0
today_total_tokens = 0
total_prompts = 0
total_sessions = set()
seen_pi_messages = set()
usage_status = ""
usage_help = ""
def add_usage(day, session_key, model, input_tokens, output_tokens, cache_read, cache_write):
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)
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
sessions_by_day[day].add(session_key)
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
def scan_pi_sessions():
root = Path.home() / ".pi" / "agent" / "sessions"
if not root.exists():
return
try:
rg = find_command("rg") or "rg"
proc = subprocess.Popen(
[rg, "--json", "-e", '"provider":"openai-codex"', "-e", '"api":"openai-codex"', str(root)],
stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL,
text=True,
errors="replace",
env=ENV,
)
except FileNotFoundError:
return
assert proc.stdout is not None
for raw in proc.stdout:
try:
event = json.loads(raw)
if event.get("type") != "match":
continue
line = event.get("data", {}).get("lines", {}).get("text", "")
path = event.get("data", {}).get("path", {}).get("text", "pi-session")
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
seen_pi_messages.add(message_key)
message = entry.get("message") or {}
if message.get("role") != "assistant":
continue
provider = str(message.get("provider") or "")
api = str(message.get("api") or "")
if provider != "openai-codex" and not api.startswith("openai-codex"):
continue
usage = message.get("usage") or {}
if not usage:
continue
total = number(usage.get("totalTokens"))
input_tokens = number(usage.get("input"))
output_tokens = number(usage.get("output"))
cache_read = number(usage.get("cacheRead"))
cache_write = number(usage.get("cacheWrite"))
if total and not (input_tokens or output_tokens or cache_read or cache_write):
input_tokens = total
if not (input_tokens or output_tokens or cache_read or cache_write):
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)
try:
proc.wait(timeout=1)
except Exception:
proc.kill()
def scan_native_codex_sessions():
codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex"))
roots = [codex_home / "sessions", codex_home / "archived_sessions"]
files = []
cutoff = time.time() - 30 * 24 * 60 * 60
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 {}
usage = info.get("total_token_usage") or {}
input_tokens = number(usage.get("input_tokens"))
output_tokens = number(usage.get("output_tokens")) + number(usage.get("reasoning_output_tokens"))
cache_read = number(usage.get("cached_input_tokens"))
cache_write = 0
if not (input_tokens or output_tokens or cache_read):
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
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)
def fetch_codex_rpc():
result = {
"rateLimitPercent": -1,
"rateLimitLabel": "",
"rateLimitResetAt": "",
"secondaryRateLimitPercent": -1,
"secondaryRateLimitLabel": "",
"secondaryRateLimitResetAt": "",
"tierLabel": "",
}
codex = find_command("codex")
if not codex:
result["usageStatusText"] = "Codex unavailable"
result["authHelpText"] = "codex not found in PATH"
return result
try:
proc = subprocess.Popen(
[codex, "-s", "read-only", "-a", "untrusted", "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": "omarchy-model-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 {}
limits = (limits_msg.get("result") or {}).get("rateLimits") or {}
plan = limits.get("planType") or account.get("planType") or account.get("type") or ""
result["tierLabel"] = str(plan) if plan else ""
def fill(prefix, window):
if not isinstance(window, dict):
return
used = window.get("usedPercent")
if used is not None:
result[prefix + "Percent"] = float(used) / 100.0
mins = number(window.get("windowDurationMins"))
if mins:
if mins == 10080:
result[prefix + "Label"] = "Weekly (7-day)"
elif mins % 60 == 0:
result[prefix + "Label"] = f"{mins // 60}h window"
else:
result[prefix + "Label"] = f"{mins}m window"
reset = window.get("resetsAt")
if reset:
result[prefix + "ResetAt"] = datetime.fromtimestamp(number(reset), timezone.utc).isoformat()
fill("rateLimit", limits.get("primary"))
fill("secondaryRateLimit", limits.get("secondary"))
except Exception as exc:
result["usageStatusText"] = "Codex limits unavailable"
result["authHelpText"] = str(exc)
finally:
try:
proc.terminate()
proc.wait(timeout=1)
except Exception:
try:
proc.kill()
except Exception:
pass
return result
scan_pi_sessions()
scan_native_codex_sessions()
rpc = fetch_codex_rpc()
out = {
"ready": True,
"hasLocalStats": True,
"todayPrompts": today_prompts,
"todaySessions": len(today_sessions),
"todayTotalTokens": today_total_tokens,
"todayTokensByModel": today_tokens_by_model,
"recentDays": [recent[day] for day in recent_dates],
"totalPrompts": total_prompts,
"totalSessions": len(total_sessions),
"modelUsage": model_usage,
"usageStatusText": usage_status,
"authHelpText": usage_help,
}
out.update(rpc)
print(json.dumps(out, separators=(",", ":")))