token_count.py

#!/usr/bin/env python3
"""Sum Claude Code API token usage for a project's transcripts."""
import argparse
import glob
import json
import os
import sys
import tempfile

FIELDS = (
    "input_tokens",
    "output_tokens",
    "cache_creation_input_tokens",
    "cache_read_input_tokens",
)

# gCO2e per million tokens; Jegham et al. 2025 (AWS inference energy measurements)
CARBON_FACTORS_G_PER_MTOK = {
    "fable": {"input": 156, "output": 3304},
    "opus": {"input": 78, "output": 1652},
    "sonnet": {"input": 39, "output": 826},
    "haiku": {"input": 20, "output": 413},
}


def transcript_dir(project_path):
    slug = os.path.abspath(project_path).replace("/", "-")
    return os.path.expanduser(f"~/.claude/projects/{slug}")


def count_tokens(root):
    totals = {f: 0 for f in FIELDS}
    by_model = {}
    seen_ids = set()
    files = sorted(
        glob.glob(os.path.join(root, "*.jsonl"))
        + glob.glob(os.path.join(root, "*", "subagents", "*.jsonl"))
    )
    messages = 0
    for path in files:
        with open(path) as fh:
            for line in fh:
                try:
                    record = json.loads(line)
                except json.JSONDecodeError:
                    continue
                message = record.get("message") or {}
                usage = message.get("usage")
                if message.get("role") != "assistant" or not usage:
                    continue
                msg_id = message.get("id") or record.get("requestId")
                if msg_id in seen_ids:
                    continue
                seen_ids.add(msg_id)
                messages += 1
                model_totals = by_model.setdefault(message.get("model") or "unknown", {f: 0 for f in FIELDS})
                for field in FIELDS:
                    count = usage.get(field, 0)
                    totals[field] += count
                    model_totals[field] += count
    return totals, by_model, len(files), messages


def estimate_carbon(by_model):
    estimates = {}
    for model, totals in by_model.items():
        factors = next((f for key, f in CARBON_FACTORS_G_PER_MTOK.items() if key in model.lower()), None)
        if not factors:
            continue
        # cache reads skip most of the compute a fresh input token costs
        input_like = (
            totals["input_tokens"]
            + totals["cache_creation_input_tokens"]
            + 0.1 * totals["cache_read_input_tokens"]
        )
        estimates[model] = (input_like / 1e6) * factors["input"] + (totals["output_tokens"] / 1e6) * factors["output"]
    return estimates


def selftest():
    with tempfile.TemporaryDirectory() as root:
        session = os.path.join(root, "session.jsonl")
        usage = {
            "input_tokens": 10,
            "output_tokens": 5,
            "cache_creation_input_tokens": 2,
            "cache_read_input_tokens": 1,
        }
        lines = [
            {"message": {"role": "assistant", "id": "msg_1", "model": "claude-sonnet-5", "usage": usage}},
            # same message id repeated (multi-block response) must not double-count
            {"message": {"role": "assistant", "id": "msg_1", "model": "claude-sonnet-5", "usage": usage}},
            {"message": {"role": "assistant", "id": "msg_2", "model": "claude-sonnet-5", "usage": usage}},
            {"message": {"role": "user", "content": "hi"}},
        ]
        with open(session, "w") as fh:
            for line in lines:
                fh.write(json.dumps(line) + "\n")

        agent_dir = os.path.join(root, "sess2", "subagents")
        os.makedirs(agent_dir)
        with open(os.path.join(agent_dir, "agent-1.jsonl"), "w") as fh:
            fh.write(json.dumps({"message": {"role": "assistant", "id": "msg_3", "model": "model-b", "usage": usage}}) + "\n")

        totals, by_model, nfiles, messages = count_tokens(root)
        assert nfiles == 2, nfiles
        assert messages == 3, messages
        assert totals["input_tokens"] == 30, totals
        assert totals["output_tokens"] == 15, totals
        assert by_model["claude-sonnet-5"]["input_tokens"] == 20, by_model
        assert by_model["model-b"]["input_tokens"] == 10, by_model

        carbon = estimate_carbon(by_model)
        assert "model-b" not in carbon, carbon  # no factor match for made-up name
        # 2 deduped msgs on claude-sonnet-5: input=20, output=10, cache_creation=4, cache_read=2
        expected = (20 + 4 + 0.1 * 2) / 1e6 * 39 + 10 / 1e6 * 826
        assert abs(carbon["claude-sonnet-5"] - expected) < 1e-9, carbon
    print("selftest OK")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("project", nargs="?", default=".", help="project directory (default: cwd)")
    parser.add_argument("--selftest", action="store_true", help="run internal self-check and exit")
    args = parser.parse_args()

    if args.selftest:
        selftest()
        return

    root = transcript_dir(args.project)
    if not os.path.isdir(root):
        sys.exit(f"no transcripts found at {root}")

    totals, by_model, nfiles, messages = count_tokens(root)
    grand_total = sum(totals.values())

    print(f"transcripts: {root}")
    print(f"files scanned: {nfiles}, assistant messages: {messages}")
    for field in FIELDS:
        print(f"  {field}: {totals[field]:,}")
    print(f"total tokens: {grand_total:,}")

    print("by model:")
    for model in sorted(by_model):
        model_totals = by_model[model]
        print(f"  {model}:")
        for field in FIELDS:
            print(f"    {field}: {model_totals[field]:,}")
        print(f"    total: {sum(model_totals.values()):,}")

    carbon = estimate_carbon(by_model)
    if carbon:
        print("estimated carbon footprint (Jegham et al. 2025):")
        for model in sorted(carbon):
            print(f"  {model}: {carbon[model]:,.1f} gCO2e")
        print(f"  total: {sum(carbon.values()):,.1f} gCO2e")


if __name__ == "__main__":
    main()