mirror of
https://github.com/arthur-pbty/arthur-os.git
synced 2026-08-03 12:47:50 +02:00
239 lines
8.6 KiB
Python
Executable File
239 lines
8.6 KiB
Python
Executable File
#!/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())
|