cleveragents-core/tools/token_usage_audit.py

"""Aggregate historical token usage across the two surfaces that
consume LLM tokens in this project:

1. **Claude Code interactive sessions** — JSONL transcripts at
   ``~/.claude/projects/-home-drew-repos-cleveragents-core/*.jsonl``.
   Each line carries an event; assistant-message events embed a
   ``message.usage`` object with input / cache-read /
   cache-creation / output token counts.

2. **OpenCode worker sessions** (the deterministic dispatcher
   pipeline — reviewer, implementer, merge, conflict) — JSON
   archives at ``<repo>/.dispatcher-logs/sessions/*.json``. The
   ``per_turn`` array holds per-turn token counts (input, output,
   reasoning).

Reads everything, sums by day and by agent/source, and emits a
structured JSON summary. Designed to be run BEFORE installing a new
LLM-affecting tool (e.g. graphify) to establish a baseline, then
again later for a before/after comparison.

Usage:
    python3 tools/token_usage_audit.py [--out PATH]

Without ``--out`` the summary prints to stdout. With ``--out PATH``
it writes the summary to that path and prints a one-line digest.

The summary is intentionally serialisable: dates as ISO strings,
totals as integers, agents/sessions as sortable keys. A later
re-run produces a comparable shape so a diff is straightforward.
"""

from __future__ import annotations

import argparse
import json
import sys
from collections import defaultdict
from collections.abc import Iterable
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

REPO_ROOT = Path(__file__).resolve().parent.parent
CLAUDE_PROJECT_DIR = (
    Path.home() / ".claude" / "projects" / "-home-drew-repos-cleveragents-core"
)
OPENCODE_ARCHIVE_DIR = REPO_ROOT / ".dispatcher-logs" / "sessions"


def _iter_jsonl(path: Path) -> Iterable[dict[str, Any]]:
    """Yield every JSON object from a JSONL file, swallowing
    parse errors (malformed lines are rare and shouldn't kill the
    whole audit). The harness writes one event per line."""
    try:
        with path.open("r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    yield json.loads(line)
                except json.JSONDecodeError:
                    continue
    except OSError:
        return


def _date_of(timestamp: str | None) -> str | None:
    """Extract YYYY-MM-DD from an ISO-8601 timestamp; return None
    on any parse failure so the caller can decide."""
    if not timestamp or not isinstance(timestamp, str):
        return None
    try:
        # Trim sub-millisecond precision Python's stdlib parser
        # rejects on some platforms.
        dt = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
        return dt.astimezone(timezone.utc).date().isoformat()
    except (ValueError, TypeError):
        return None


def audit_claude_code() -> dict[str, Any]:
    """Walk every Claude Code JSONL transcript for this project and
    sum per-assistant-message usage by date.

    Returns a dict with:
    - ``sessions``: count of distinct ``sessionId`` values seen
    - ``messages_with_usage``: count of assistant messages whose
      usage block we summed
    - ``by_date``: mapping ``"YYYY-MM-DD"`` → totals dict
    - ``totals``: grand totals across every transcript
    """
    by_date: dict[str, dict[str, int]] = defaultdict(
        lambda: {"input": 0, "cache_read": 0, "cache_create": 0, "output": 0, "messages": 0}
    )
    sessions: set[str] = set()
    messages_with_usage = 0

    if not CLAUDE_PROJECT_DIR.exists():
        return {
            "exists": False,
            "path": str(CLAUDE_PROJECT_DIR),
            "sessions": 0,
            "messages_with_usage": 0,
            "by_date": {},
            "totals": {"input": 0, "cache_read": 0, "cache_create": 0, "output": 0, "messages": 0},
        }

    for jsonl_path in sorted(CLAUDE_PROJECT_DIR.glob("*.jsonl")):
        for event in _iter_jsonl(jsonl_path):
            session_id = event.get("sessionId")
            if isinstance(session_id, str):
                sessions.add(session_id)
            msg = event.get("message")
            if not isinstance(msg, dict):
                continue
            if msg.get("role") != "assistant":
                continue
            usage = msg.get("usage")
            if not isinstance(usage, dict):
                continue
            date = _date_of(event.get("timestamp")) or "unknown"
            bucket = by_date[date]
            bucket["input"] += int(usage.get("input_tokens") or 0)
            bucket["cache_read"] += int(usage.get("cache_read_input_tokens") or 0)
            bucket["cache_create"] += int(usage.get("cache_creation_input_tokens") or 0)
            bucket["output"] += int(usage.get("output_tokens") or 0)
            bucket["messages"] += 1
            messages_with_usage += 1

    totals = {"input": 0, "cache_read": 0, "cache_create": 0, "output": 0, "messages": 0}
    for d in by_date.values():
        for k, v in d.items():
            totals[k] += v

    return {
        "exists": True,
        "path": str(CLAUDE_PROJECT_DIR),
        "sessions": len(sessions),
        "messages_with_usage": messages_with_usage,
        "by_date": dict(sorted(by_date.items())),
        "totals": totals,
    }


def audit_opencode() -> dict[str, Any]:
    """Walk every OpenCode session archive and sum per_turn tokens.

    Returns a dict with:
    - ``archives``: count of archive files inspected
    - ``by_date``: mapping ``"YYYY-MM-DD"`` → totals dict
    - ``by_agent``: mapping ``agent_name`` → totals dict
    - ``totals``: grand totals
    """
    by_date: dict[str, dict[str, int]] = defaultdict(
        lambda: {"input": 0, "output": 0, "reasoning": 0, "turns": 0, "sessions": 0}
    )
    by_agent: dict[str, dict[str, int]] = defaultdict(
        lambda: {"input": 0, "output": 0, "reasoning": 0, "turns": 0, "sessions": 0}
    )
    archives_seen = 0

    if not OPENCODE_ARCHIVE_DIR.exists():
        return {
            "exists": False,
            "path": str(OPENCODE_ARCHIVE_DIR),
            "archives": 0,
            "by_date": {},
            "by_agent": {},
            "totals": {"input": 0, "output": 0, "reasoning": 0, "turns": 0, "sessions": 0},
        }

    for archive_path in sorted(OPENCODE_ARCHIVE_DIR.glob("*.json")):
        try:
            with archive_path.open("r", encoding="utf-8") as f:
                archive = json.load(f)
        except (OSError, json.JSONDecodeError):
            continue
        archives_seen += 1
        agent = str(archive.get("agent") or "unknown")
        date = _date_of(archive.get("started_at")) or "unknown"
        per_turn = archive.get("per_turn") or []
        if not isinstance(per_turn, list):
            continue

        sess_in = sess_out = sess_reason = sess_turns = 0
        for turn in per_turn:
            if not isinstance(turn, dict):
                continue
            sess_in += int(turn.get("input_tokens") or 0)
            sess_out += int(turn.get("output_tokens") or 0)
            sess_reason += int(turn.get("reasoning_tokens") or 0)
            sess_turns += 1

        by_date[date]["input"] += sess_in
        by_date[date]["output"] += sess_out
        by_date[date]["reasoning"] += sess_reason
        by_date[date]["turns"] += sess_turns
        by_date[date]["sessions"] += 1

        by_agent[agent]["input"] += sess_in
        by_agent[agent]["output"] += sess_out
        by_agent[agent]["reasoning"] += sess_reason
        by_agent[agent]["turns"] += sess_turns
        by_agent[agent]["sessions"] += 1

    totals = {"input": 0, "output": 0, "reasoning": 0, "turns": 0, "sessions": 0}
    for d in by_date.values():
        for k, v in d.items():
            totals[k] += v

    return {
        "exists": True,
        "path": str(OPENCODE_ARCHIVE_DIR),
        "archives": archives_seen,
        "by_date": dict(sorted(by_date.items())),
        "by_agent": dict(sorted(by_agent.items())),
        "totals": totals,
    }


def _format_int(n: int) -> str:
    return f"{n:>14,}"


def render_digest(summary: dict[str, Any]) -> str:
    """Produce a human-readable one-shot digest. Numbers are
    right-aligned for easy column comparison between baseline and
    follow-up runs."""
    cc = summary["claude_code"]
    oc = summary["opencode"]
    lines = [
        "Token usage audit",
        f"Generated at: {summary['generated_at']}",
        "",
        "=== Claude Code (interactive Claude Code sessions) ===",
        f"  transcripts dir:        {cc['path']}",
        f"  sessions:               {cc['sessions']}",
        f"  assistant messages:     {cc['messages_with_usage']}",
        f"  total input tokens:     {_format_int(cc['totals']['input'])}",
        f"  total cache_read:       {_format_int(cc['totals']['cache_read'])}",
        f"  total cache_create:     {_format_int(cc['totals']['cache_create'])}",
        f"  total output tokens:    {_format_int(cc['totals']['output'])}",
        "",
        "  by date (assistant messages, input + cache_read + cache_create + output):",
    ]
    for date, d in cc["by_date"].items():
        total = d["input"] + d["cache_read"] + d["cache_create"] + d["output"]
        lines.append(
            f"    {date}  msgs={d['messages']:>4}  total_tokens={total:>14,}  "
            f"(in={d['input']:>9,} c_r={d['cache_read']:>11,} "
            f"c_w={d['cache_create']:>11,} out={d['output']:>8,})"
        )

    lines += [
        "",
        "=== OpenCode (dispatcher worker sessions) ===",
        f"  archives dir:           {oc['path']}",
        f"  sessions:               {oc['totals']['sessions']}",
        f"  assistant turns:        {oc['totals']['turns']}",
        f"  total input tokens:     {_format_int(oc['totals']['input'])}",
        f"  total output tokens:    {_format_int(oc['totals']['output'])}",
        f"  total reasoning tokens: {_format_int(oc['totals']['reasoning'])}",
        "",
        "  by date:",
    ]
    for date, d in oc["by_date"].items():
        total = d["input"] + d["output"] + d["reasoning"]
        lines.append(
            f"    {date}  sess={d['sessions']:>3}  turns={d['turns']:>4}  "
            f"total_tokens={total:>12,}  "
            f"(in={d['input']:>11,} out={d['output']:>9,} reason={d['reasoning']:>9,})"
        )

    lines += ["", "  by agent (top consumers by total tokens):"]
    agent_rows = sorted(
        oc["by_agent"].items(),
        key=lambda kv: -(kv[1]["input"] + kv[1]["output"] + kv[1]["reasoning"]),
    )
    for agent, d in agent_rows:
        total = d["input"] + d["output"] + d["reasoning"]
        lines.append(
            f"    {agent:<32}  sess={d['sessions']:>3}  turns={d['turns']:>4}  "
            f"total_tokens={total:>12,}  "
            f"(in={d['input']:>11,} out={d['output']:>9,} reason={d['reasoning']:>9,})"
        )

    return "\n".join(lines)


def build_summary() -> dict[str, Any]:
    return {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "claude_code": audit_claude_code(),
        "opencode": audit_opencode(),
    }


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--out",
        type=Path,
        help="Write the JSON summary to this path (otherwise prints the digest to stdout).",
    )
    parser.add_argument(
        "--json-only",
        action="store_true",
        help="Print the JSON summary to stdout instead of the human digest.",
    )
    args = parser.parse_args(argv)
    summary = build_summary()
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(json.dumps(summary, indent=2), encoding="utf-8")
        digest = render_digest(summary)
        print(f"wrote {args.out}\n")
        print(digest)
    elif args.json_only:
        print(json.dumps(summary, indent=2))
    else:
        print(render_digest(summary))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())