analyze_opencode_cache.py

#!/usr/bin/env python3
"""Analyze OpenCode token-cache usage for one session.

OpenCode stores cache counters on assistant messages as:

    tokens.cache.read
    tokens.cache.write

The database does not currently contain an explicit cache-invalidation event.
This script therefore reports "inferred invalidations": a transition from one
or more cache reads to a later token-bearing assistant message with no cache
read. Treat those transitions as a useful heuristic, not ground truth.
"""

from __future__ import annotations

import argparse
import curses
import json
import sqlite3
import sys
from pathlib import Path
from typing import Any


DEFAULT_DB = Path("~/.local/share/opencode/opencode.db").expanduser()


def integer(value: Any) -> int:
    """Convert nullable or numeric JSON values to a non-negative integer."""
    try:
        return max(0, int(value or 0))
    except (TypeError, ValueError):
        return 0


def number(value: Any) -> float:
    try:
        return max(0.0, float(value or 0))
    except (TypeError, ValueError):
        return 0.0


def percent(numerator: int, denominator: int) -> float:
    return (numerator / denominator * 100) if denominator else 0.0


def model_name(value: Any) -> str | None:
    if isinstance(value, dict):
        provider = value.get("providerID")
        model = value.get("id") or value.get("modelID")
        if provider and model:
            return f"{provider}/{model}"
        return model or provider
    if isinstance(value, str):
        try:
            return model_name(json.loads(value))
        except json.JSONDecodeError:
            return value
    return None


def message_tokens(data: dict[str, Any]) -> dict[str, int]:
    tokens = data.get("tokens") or {}
    cache = tokens.get("cache") or {}
    return {
        "input": integer(tokens.get("input")),
        "output": integer(tokens.get("output")),
        "reasoning": integer(tokens.get("reasoning")),
        "total": integer(tokens.get("total")),
        "cache_read": integer(cache.get("read", tokens.get("cache_read"))),
        "cache_write": integer(cache.get("write", tokens.get("cache_write"))),
    }


def load_session(
    connection: sqlite3.Connection, session_id: str
) -> tuple[dict[str, Any], list[dict[str, Any]], int]:
    session_row = connection.execute(
        """
        SELECT id, title, directory, version, model, time_created, time_updated,
               cost, tokens_input, tokens_output, tokens_reasoning,
               tokens_cache_read, tokens_cache_write
        FROM session
        WHERE id = ?
        """,
        (session_id,),
    ).fetchone()
    if session_row is None:
        raise ValueError(f"session not found: {session_id}")

    session = dict(session_row)
    messages: list[dict[str, Any]] = []
    invalid_json_count = 0

    rows = connection.execute(
        """
        SELECT id, time_created, data
        FROM message
        WHERE session_id = ?
        ORDER BY time_created, id
        """,
        (session_id,),
    )
    for row in rows:
        try:
            data = json.loads(row["data"])
        except (TypeError, json.JSONDecodeError):
            invalid_json_count += 1
            continue

        if data.get("role") != "assistant":
            continue

        tokens = message_tokens(data)
        cost = number(data.get("cost"))
        if not any(tokens.values()) and not cost:
            continue

        messages.append(
            {
                "id": row["id"],
                "time_created": row["time_created"],
                "model_id": data.get("modelID"),
                "provider_id": data.get("providerID"),
                "finish": data.get("finish"),
                "cost": cost,
                **tokens,
            }
        )

    return session, messages, invalid_json_count


def load_sessions(connection: sqlite3.Connection) -> list[dict[str, Any]]:
    rows = connection.execute(
        """
        SELECT id, title, directory, time_created, time_updated
        FROM session
        ORDER BY time_updated DESC, id DESC
        """
    )
    return [dict(row) for row in rows]


def display_title(session: dict[str, Any]) -> str:
    title = " ".join(str(session.get("title") or "").split())
    return title or "Untitled session"


def clipped(text: str, width: int) -> str:
    if width <= 0:
        return ""
    if len(text) <= width:
        return text
    if width <= 3:
        return text[:width]
    return text[: width - 3] + "..."


def choose_session(stdscr: Any, sessions: list[dict[str, Any]]) -> str | None:
    """Run a small keyboard-driven session picker and return the chosen ID."""
    if not sessions:
        return None

    try:
        curses.curs_set(0)
    except curses.error:
        pass
    stdscr.keypad(True)

    selected = 0
    offset = 0
    while True:
        height, width = stdscr.getmaxyx()
        stdscr.erase()

        if height < 7 or width < 42:
            message = "Terminal too small. Resize, or press q to quit."
            try:
                stdscr.addstr(0, 0, clipped(message, max(width - 1, 1)))
            except curses.error:
                pass
            stdscr.refresh()
            key = stdscr.getch()
            if key in (ord("q"), ord("Q"), 27):
                return None
            continue

        list_height = height - 5
        if selected < offset:
            offset = selected
        if selected >= offset + list_height:
            offset = selected - list_height + 1

        try:
            stdscr.addstr(0, 2, clipped("Select an OpenCode session", width - 4), curses.A_BOLD)
            stdscr.addstr(1, 2, clipped("Recent sessions, newest first", width - 4), curses.A_DIM)

            for row in range(list_height):
                index = offset + row
                if index >= len(sessions):
                    break
                session = sessions[index]
                label = f"{display_title(session)} [{session['id']}]"
                attr = curses.A_REVERSE if index == selected else curses.A_NORMAL
                stdscr.addstr(row + 3, 2, clipped(label, width - 4).ljust(width - 4), attr)

            footer = "↑/↓ move  PgUp/PgDn page  Enter select  q/Esc cancel"
            stdscr.addstr(height - 1, 2, clipped(footer, width - 4), curses.A_DIM)
        except curses.error:
            pass
        stdscr.refresh()

        key = stdscr.getch()
        if key in (curses.KEY_UP, ord("k")):
            selected = max(0, selected - 1)
        elif key in (curses.KEY_DOWN, ord("j")):
            selected = min(len(sessions) - 1, selected + 1)
        elif key in (curses.KEY_PPAGE,):
            selected = max(0, selected - list_height)
        elif key in (curses.KEY_NPAGE,):
            selected = min(len(sessions) - 1, selected + list_height)
        elif key == curses.KEY_HOME:
            selected = 0
        elif key == curses.KEY_END:
            selected = len(sessions) - 1
        elif key in (curses.KEY_ENTER, 10, 13):
            return sessions[selected]["id"]
        elif key in (ord("q"), ord("Q"), 27):
            return None


def pick_session(connection: sqlite3.Connection) -> str | None:
    sessions = load_sessions(connection)
    return curses.wrapper(lambda stdscr: choose_session(stdscr, sessions))


def analyze(
    session: dict[str, Any], messages: list[dict[str, Any]], invalid_json_count: int
) -> dict[str, Any]:
    totals = {
        key: sum(message[key] for message in messages)
        for key in ("input", "output", "reasoning", "total", "cache_read", "cache_write")
    }
    totals["cost"] = sum(message["cost"] for message in messages)
    read_messages = sum(message["cache_read"] > 0 for message in messages)
    write_messages = sum(message["cache_write"] > 0 for message in messages)

    model_breakdown: dict[str, dict[str, Any]] = {}
    for message in messages:
        provider = message["provider_id"] or "unknown-provider"
        model = message["model_id"] or "unknown-model"
        key = f"{provider}/{model}"
        breakdown = model_breakdown.setdefault(
            key,
            {
                "provider": provider,
                "model": model,
                "messages": 0,
                "cost": 0.0,
                "input": 0,
                "output": 0,
                "cache_read": 0,
                "cache_write": 0,
            },
        )
        breakdown["messages"] += 1
        breakdown["cost"] += message["cost"]
        for token_key in ("input", "output", "cache_read", "cache_write"):
            breakdown[token_key] += message[token_key]

    for breakdown in model_breakdown.values():
        breakdown["cost"] = round(breakdown["cost"], 12)
        breakdown["cost_share_percent"] = round(
            percent(breakdown["cost"], totals["cost"]), 2
        )

    # There is no explicit invalidation record in the current schema. Count
    # one inferred invalidation per contiguous cache-read gap after a hit.
    invalidations: list[dict[str, Any]] = []
    has_seen_cache_read = False
    in_read_gap = False
    previous_read_message_id: str | None = None
    for message in messages:
        if message["cache_read"] > 0:
            has_seen_cache_read = True
            in_read_gap = False
            previous_read_message_id = message["id"]
            continue

        meaningful = message["input"] > 0 or message["cache_write"] > 0
        if has_seen_cache_read and meaningful and not in_read_gap:
            invalidations.append(
                {
                    "message_id": message["id"],
                    "time_created": message["time_created"],
                    "previous_cache_read_message_id": previous_read_message_id,
                    "cache_write_tokens": message["cache_write"],
                    "reason": "cache-read reset",
                }
            )
            in_read_gap = True

    # OpenCode's `input` is the uncached input counter. Cache reads are a
    # separate counter, so this is the fraction of effective prompt tokens
    # served from cache, excluding cache-write accounting.
    effective_prompt_tokens = totals["input"] + totals["cache_read"]
    cache_read_share = percent(totals["cache_read"], effective_prompt_tokens)

    return {
        "session": {
            "id": session["id"],
            "title": session["title"],
            "directory": session["directory"],
            "version": session["version"],
            "model": model_name(session["model"]),
            "time_created": session["time_created"],
            "time_updated": session["time_updated"],
            "cost": session["cost"],
        },
        "message_count": len(messages),
        "invalid_json_count": invalid_json_count,
        "totals_from_messages": totals,
        "cost_source": (
            "sum of assistant message data.cost values; each value is already "
            "provider/model-specific"
        ),
        "cost_by_model": sorted(
            model_breakdown.values(), key=lambda item: item["cost"], reverse=True
        ),
        "stored_session_totals": {
            "input": integer(session["tokens_input"]),
            "output": integer(session["tokens_output"]),
            "reasoning": integer(session["tokens_reasoning"]),
            "cache_read": integer(session["tokens_cache_read"]),
            "cache_write": integer(session["tokens_cache_write"]),
            "cost": number(session["cost"]),
        },
        "derived": {
            "effective_prompt_tokens": effective_prompt_tokens,
            "cache_read_share_percent": round(cache_read_share, 2),
            "uncached_input_share_percent": round(
                percent(totals["input"], effective_prompt_tokens), 2
            ),
            "cache_read_messages": read_messages,
            "cache_write_messages": write_messages,
            "inferred_invalidations": len(invalidations),
            "invalidation_detection": (
                "heuristic: first meaningful no-cache-read message after a cache-read run"
            ),
        },
        "inferred_invalidation_events": invalidations,
        "messages": messages,
    }


def format_number(value: int | float) -> str:
    if isinstance(value, float):
        return f"{value:,.2f}"
    return f"{value:,}"


def print_report(report: dict[str, Any], show_timeline: bool) -> None:
    session = report["session"]
    totals = report["totals_from_messages"]
    stored = report["stored_session_totals"]
    derived = report["derived"]

    print(f"Session: {session['id']}")
    print(f"Title:   {session['title']}")
    print(f"Model:   {session['model'] or 'unknown'}")
    print(f"Messages with token data: {report['message_count']}")
    print()
    print("Token usage")
    print(f"  Uncached input: {format_number(totals['input'])}")
    print(f"  Cache read:     {format_number(totals['cache_read'])}")
    print(f"  Cache write:    {format_number(totals['cache_write'])}")
    print(f"  Output:         {format_number(totals['output'])}")
    print(f"  Reasoning:      {format_number(totals['reasoning'])}")
    print(f"  Effective prompt: {format_number(derived['effective_prompt_tokens'])}")
    print(f"  Cost:           ${totals['cost']:.8f}")
    print()
    print("Cache efficiency")
    print(f"  Input served from cache: {derived['cache_read_share_percent']:.2f}%")
    print(f"  Input not served from cache: {derived['uncached_input_share_percent']:.2f}%")
    print(f"  Messages with cache reads:  {derived['cache_read_messages']}")
    print(f"  Messages with cache writes: {derived['cache_write_messages']}")
    print(f"  Inferred invalidations:     {derived['inferred_invalidations']}")
    print("  Invalidation definition:    cache-read reset after a prior cache hit")
    print()
    print("Stored session totals vs message totals")
    for key in ("input", "output", "reasoning", "cache_read", "cache_write"):
        message_value = totals[key]
        stored_value = stored[key]
        marker = "" if message_value == stored_value else " (DIFF)"
        print(
            f"  {key:11} messages={format_number(message_value):>12} "
            f"session={format_number(stored_value):>12}{marker}"
        )
    cost_marker = "" if abs(totals["cost"] - stored["cost"]) < 1e-9 else " (DIFF)"
    print(
        f"  {'cost':11} messages=${totals['cost']:>11.8f} "
        f"session=${stored['cost']:>11.8f}{cost_marker}"
    )

    print("\nCost by provider/model")
    for breakdown in report["cost_by_model"]:
        print(
            f"  {breakdown['provider']}/{breakdown['model']}: "
            f"${breakdown['cost']:.8f} ({breakdown['cost_share_percent']:.2f}%, "
            f"{breakdown['messages']} messages)"
        )

    if report["invalid_json_count"]:
        print(f"\nWarning: skipped {report['invalid_json_count']} malformed message row(s).")

    if report["inferred_invalidation_events"]:
        print("\nInferred invalidation events")
        for event in report["inferred_invalidation_events"]:
            write = format_number(event["cache_write_tokens"])
            print(f"  {event['message_id']} (cache write: {write})")

    if show_timeline:
        print("\nMessage timeline")
        print("  id                                      input    cache-read  cache-write  output        cost")
        for message in report["messages"]:
            print(
                f"  {message['id']:<40} {message['input']:>8,} "
                f"{message['cache_read']:>11,} {message['cache_write']:>12,} "
                f"{message['output']:>8,} ${message['cost']:.8f}"
            )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Analyze OpenCode cache usage for a session."
    )
    parser.add_argument(
        "session_id",
        nargs="?",
        help="OpenCode session ID, for example ses_... (omit to open the picker)",
    )
    parser.add_argument(
        "--db",
        type=Path,
        default=DEFAULT_DB,
        help=f"SQLite database path (default: {DEFAULT_DB})",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        dest="as_json",
        help="write machine-readable JSON instead of the text report",
    )
    parser.add_argument(
        "--timeline",
        action="store_true",
        help="include one row per assistant message in the text report",
    )
    parser.add_argument(
        "-i",
        "--interactive",
        action="store_true",
        help="open the session picker even when a session ID is supplied",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    db_path = args.db.expanduser()
    if not db_path.is_file():
        print(f"database not found: {db_path}", file=sys.stderr)
        return 2

    try:
        connection = sqlite3.connect(
            f"file:{db_path}?mode=ro", uri=True, timeout=5
        )
        connection.row_factory = sqlite3.Row
        try:
            if args.as_json and (args.interactive or args.session_id is None):
                print("error: --json requires a session ID and cannot use the picker", file=sys.stderr)
                return 2

            session_id = args.session_id
            if args.interactive or session_id is None:
                try:
                    session_id = pick_session(connection)
                except curses.error as error:
                    print(
                        f"error: interactive picker requires a terminal ({error})",
                        file=sys.stderr,
                    )
                    return 2
                if session_id is None:
                    return 0

            session, messages, invalid_json_count = load_session(
                connection, session_id
            )
        finally:
            connection.close()
    except (sqlite3.Error, ValueError) as error:
        print(f"error: {error}", file=sys.stderr)
        return 1

    report = analyze(session, messages, invalid_json_count)
    if args.as_json:
        print(json.dumps(report, indent=2))
    else:
        print_report(report, args.timeline)
    return 0


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