engine/core/tokens.py

"""tokens.py — session token rollup.

The per-turn rollup already happens in the loader (sum of message.usage across the
turn's assistant rows). This module sums turns → session totals and ALWAYS exposes
the cacheRead/out ratio (the cost-driver headline: cacheRead dominates and is the
"expensive because big, not wasteful" signal). Pure code, NO model.
"""
from __future__ import annotations

from typing import Any

from engine.contract import Tokens

# Context-WINDOW tiers (the GAUGE denominator) — ascending known Claude windows. The
# JSONL records only a bare model id (e.g. "claude-opus-4-8"), never the 1M-beta "[1m]"
# suffix or a window field, so the tier can't always be read off the name. We therefore
# pick the SMALLEST known tier that fits the session's observed peak occupancy, capped
# at what the model family supports (Haiku tops out at 200k; Opus/Sonnet can opt into
# the 1M beta). Denominator stays honest: never smaller than what the data proves was
# used, never larger than the model can do. NOT hardcoded to 1M.
CONTEXT_TIERS = (200_000, 1_000_000)


def _model_max_window(model: str | None) -> int:
    """The largest context window the model family supports."""
    m = (model or "").lower()
    if "haiku" in m:
        return 200_000   # Haiku has no 1M tier
    return 1_000_000     # Opus / Sonnet (and unknown) can run the 1M beta


def context_limit(model: str | None, peak: int) -> int:
    """Gauge denominator: the smallest known tier that fits `peak`, capped at the model
    family's max. Underfilled sessions read against the standard 200k window (the common
    case); a peak above 200k proves the 1M tier was in use. If peak exceeds every tier
    the model supports, returns the cap (and the caller's overLimit guard flags it)."""
    cap = _model_max_window(model)
    for tier in CONTEXT_TIERS:
        if tier > cap:
            break
        if tier >= peak:
            return tier
    return cap

# A compaction shows up as a sharp DROP in occupancy: a new prompt can only ADD to the
# window, and cache-TTL expiry merely shifts cacheRead↔input without lowering the total,
# so the only things that pull occupancy down are /compact, /clear, or an auto-compact.
# Flag a drop ONLY once the window was substantially full, so small turns never read as
# compactions.
_COMPACT_FLOOR = 100_000  # occupancy must have been at least this full first
_COMPACT_RATIO = 0.6      # ...and then fell below 60% of that


def session_tokens(turns) -> Tokens:
    """Sum per-turn token rollups into a session total."""
    total = Tokens()
    for t in turns:
        total = total.add(t.tokens)
    return total


def cache_read_ratio(tokens: Tokens) -> float:
    """cacheRead / out. Always reported. 0.0 when out == 0 (avoid div-by-zero)."""
    if not tokens.out:
        return 0.0
    return tokens.cacheRead / tokens.out


def rollup(turns) -> dict[str, Any]:
    """Session token summary: totals + the always-present cacheRead/out ratio."""
    total = session_tokens(turns)
    return {
        "tokens": total.to_dict(),
        "cacheReadOverOut": cache_read_ratio(total),
    }


def context_window(turns, model: str | None = None) -> dict[str, Any]:
    """Point-in-time context-WINDOW occupancy — the "fuel gauge", NOT the cumulative
    token sums in rollup(). Returns the peak fill, the model-aware window `limit` (see
    context_limit — NOT hardcoded), a per-turn trajectory, and any inferred compactions
    (sharp occupancy drops). `overLimit` is the genuinely-suspect case: a request whose
    occupancy exceeds the model's window (physically impossible — if it appears, the
    source data or parse is wrong). Pure code, NO model inference at runtime."""
    traj: list[dict[str, int]] = []
    peak = 0
    for t in turns:
        if not t.ctxPeak:  # no main-thread usage seen for this turn
            continue
        traj.append({"i": t.i, "start": t.ctxStart, "peak": t.ctxPeak, "end": t.ctxEnd})
        peak = max(peak, t.ctxPeak)

    limit = context_limit(model, peak)
    over_limit: list[int] = [t.i for t in turns if t.ctxPeak and t.ctxPeak > limit]

    # compaction = a sharp drop, BETWEEN turns (prev end → this start) or WITHIN a
    # turn (this peak → this end), once the window was substantially full.
    compactions: list[dict[str, int]] = []
    prev_end = 0
    for e in traj:
        if prev_end > _COMPACT_FLOOR and e["start"] < prev_end * _COMPACT_RATIO:
            compactions.append({"atTurn": e["i"], "before": prev_end, "after": e["start"]})
        elif e["peak"] > _COMPACT_FLOOR and e["end"] < e["peak"] * _COMPACT_RATIO:
            compactions.append({"atTurn": e["i"], "before": e["peak"], "after": e["end"]})
        prev_end = e["end"]

    return {
        "limit": limit,
        "peak": peak,
        "peakPct": (peak / limit) if limit else 0.0,
        "trajectory": traj,
        "compactions": compactions,
        "overLimit": over_limit,
    }