pennyrouter 0.3.19 → 0.3.21

package/standalone/usage.py

@@ -0,0 +1,575 @@
+"""Reads token counts out of the transcripts Claude Code and Codex write as they run, plus
+the per-step records OpenCode stores in SQLite.
+
+The JSONL transcripts and SQLite store land ahead of any gateway round trip and need no proxy in
+front of the session. Only usage numbers and model metadata are read; message content is never
+parsed.
+
+The two formats differ in kind: Claude writes per-request counts and repeats a message's usage
+block on every streamed line, so totals are summed over distinct message ids. Codex writes a
+running cumulative total, so its last event wins.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sqlite3
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import pricing
+
+# Set this to False to restore the original parent-transcript-only behavior if a harness
+# changes its child-session format or a comparison needs to exclude subagent work.
+INCLUDE_SUBAGENT_USAGE = True
+
+
+@dataclass
+class Usage:
+    cache_read: int = 0
+    cache_write: int = 0
+    # Split for cost accuracy; cache_write above is their sum, shown as one row in the HUD.
+    write_5m: int = 0
+    write_1h: int = 0
+    fresh_input: int = 0
+    output: int = 0
+    context_used: int = 0
+    context_window: int | None = None
+    model: str = ""
+
+    @property
+    def total(self) -> int:
+        return self.cache_read + self.cache_write + self.fresh_input + self.output
+
+    @property
+    def context_fraction(self) -> float | None:
+        if not self.context_window:
+            return None
+        return min(1.0, self.context_used / self.context_window)
+
+
+@dataclass
+class UsageCall:
+    """One billable model call as recorded by a harness transcript."""
+
+    timestamp: str | None
+    recorded_at: float
+    model: str
+    cache_read: int
+    cache_write: int
+    write_5m: int
+    write_1h: int
+    fresh_input: int
+    output: int
+    reasoning_output: int = 0
+
+    @property
+    def total(self) -> int:
+        # Codex reports reasoning_output_tokens as a subset of output_tokens, so it is useful
+        # detail but must not be added to the billable total a second time.
+        return self.cache_read + self.cache_write + self.fresh_input + self.output
+
+
+@dataclass
+class _Progress:
+    offset: int = 0
+    usage: Usage = field(default_factory=Usage)
+    counted: set = field(default_factory=set)
+
+
+class UsageReader:
+    """Follows one transcript, parsing only what was appended since the last read."""
+
+    def __init__(self, harness: str, path: str | None = None, on_call=None):
+        self.harness = harness
+        self.path = path
+        self.on_call = on_call
+        self._progress: dict[str, _Progress] = {}
+
+    def _emit(self, call: UsageCall) -> None:
+        if self.on_call:
+            self.on_call(call)
+
+    def read(self, path: str | None = None) -> Usage | None:
+        target = path or self.path
+        if self.harness == "opencode" and target and "#" in target:
+            state = self._progress.setdefault(target, _Progress())
+            return self._read_opencode(target, state)
+        if not target or not os.path.exists(target):
+            return None
+        state = self._progress.setdefault(target, _Progress())
+
+        size = os.path.getsize(target)
+        # A shorter file is a different conversation reusing the path; start over.
+        if size < state.offset:
+            self._progress[target] = state = _Progress()
+
+        with open(target, "rb") as handle:
+            handle.seek(state.offset)
+            chunk = handle.read()
+
+        if chunk:
+            # Stop at the last newline so a partially written line is re-read once complete.
+            cut = chunk.rfind(b"\n")
+            if cut >= 0:
+                state.offset += cut + 1
+                apply = self._apply_claude if self.harness == "claude-code" else self._apply_codex
+                for raw in chunk[: cut + 1].splitlines():
+                    if raw.strip():
+                        apply(raw, state)
+        return state.usage
+
+    def _read_opencode(self, target: str, state: _Progress) -> Usage | None:
+        """Read OpenCode's per-step usage records from its SQLite store."""
+        database, session_id = target.split("#", 1)
+        try:
+            db = sqlite3.connect(f"file:{database}?mode=ro", uri=True, timeout=0.2)
+            row = db.execute("select model from session where id = ?", (session_id,)).fetchone()
+            if not row:
+                db.close()
+                return None
+            try:
+                model = str((json.loads(row[0] or "{}") or {}).get("id") or "")
+            except ValueError:
+                model = ""
+            state.usage.model = model
+            for part_id, raw in db.execute(
+                "select id, data from part where session_id = ? order by time_created, id",
+                (session_id,),
+            ):
+                if part_id in state.counted:
+                    continue
+                try:
+                    obj = json.loads(raw)
+                except (TypeError, ValueError):
+                    state.counted.add(part_id)
+                    continue
+                if obj.get("type") != "step-finish" or not isinstance(obj.get("tokens"), dict):
+                    continue
+                tokens = obj["tokens"]
+                cache = tokens.get("cache") or {}
+                read = int(cache.get("read") or 0)
+                write = int(cache.get("write") or 0)
+                fresh = int(tokens.get("input") or 0)
+                output = int(tokens.get("output") or 0)
+                state.usage.cache_read += read
+                state.usage.cache_write += write
+                state.usage.write_5m += write
+                state.usage.fresh_input += fresh
+                state.usage.output += output
+                state.usage.context_used = fresh + read + write
+                state.counted.add(part_id)
+                self._emit(UsageCall(
+                    timestamp=None, recorded_at=time.time(), model=model,
+                    cache_read=read, cache_write=write, write_5m=write, write_1h=0,
+                    fresh_input=fresh, output=output,
+                    reasoning_output=int(tokens.get("reasoning") or 0),
+                ))
+            db.close()
+            return state.usage
+        except (OSError, sqlite3.Error, ValueError):
+            return state.usage if state.usage.model else None
+
+    def read_many(self, paths: list[str]) -> Usage | None:
+        """Read a root transcript and any related child transcripts as one usage stream."""
+        if not paths:
+            return None
+        usages = [self.read(path) for path in paths]
+        present = [usage for usage in usages if usage is not None]
+        if not present:
+            return None
+        combined = Usage(model=present[-1].model,
+                         context_window=present[-1].context_window)
+        for usage in present:
+            combined.cache_read += usage.cache_read
+            combined.cache_write += usage.cache_write
+            combined.write_5m += usage.write_5m
+            combined.write_1h += usage.write_1h
+            combined.fresh_input += usage.fresh_input
+            combined.output += usage.output
+            combined.context_used = max(combined.context_used, usage.context_used)
+            if usage.model:
+                combined.model = usage.model
+            if usage.context_window:
+                combined.context_window = usage.context_window
+        return combined
+
+    def _apply_claude(self, raw: bytes, state: _Progress) -> None:
+        try:
+            obj = json.loads(raw)
+        except (ValueError, UnicodeDecodeError):
+            return
+        message = obj.get("message")
+        if not isinstance(message, dict):
+            return
+        usage = message.get("usage")
+        if not isinstance(usage, dict):
+            return
+
+        mid = message.get("id")
+        if mid:
+            if mid in state.counted:
+                return
+            state.counted.add(mid)
+
+        read = int(usage.get("cache_read_input_tokens") or 0)
+        write = int(usage.get("cache_creation_input_tokens") or 0)
+        creation = usage.get("cache_creation") or {}
+        write_5m = int(creation.get("ephemeral_5m_input_tokens") or 0)
+        write_1h = int(creation.get("ephemeral_1h_input_tokens") or 0)
+        # Older responses omit the split; treat an unsplit write as the (cheaper) 5-minute tier
+        # rather than silently dropping it, since 0/0 there would undercount the total.
+        if not write_5m and not write_1h and write:
+            write_5m = write
+        fresh = int(usage.get("input_tokens") or 0)
+
+        state.usage.cache_read += read
+        state.usage.cache_write += write
+        state.usage.write_5m += write_5m
+        state.usage.write_1h += write_1h
+        state.usage.fresh_input += fresh
+        state.usage.output += int(usage.get("output_tokens") or 0)
+        state.usage.context_used = read + write + fresh
+        if message.get("model"):
+            state.usage.model = message["model"]
+            # Anthropic records no window in the transcript, so it is inferred from the model.
+            if state.usage.context_window is None:
+                state.usage.context_window = pricing.context_window(message["model"])
+        # The long-context variants are a runtime setting rather than a distinct model id, so an
+        # observed context larger than the assumed window is the only signal one is in use.
+        if state.usage.context_window and state.usage.context_used > state.usage.context_window:
+            state.usage.context_window = pricing.next_window_above(state.usage.context_used)
+
+        self._emit(UsageCall(
+            timestamp=obj.get("timestamp"), recorded_at=time.time(),
+            model=str(message.get("model") or state.usage.model),
+            cache_read=read, cache_write=write, write_5m=write_5m,
+            write_1h=write_1h, fresh_input=fresh,
+            output=int(usage.get("output_tokens") or 0),
+            reasoning_output=0,
+        ))
+
+    def _apply_codex(self, raw: bytes, state: _Progress) -> None:
+        try:
+            obj = json.loads(raw)
+        except (ValueError, UnicodeDecodeError):
+            return
+        payload = obj.get("payload")
+        if not isinstance(payload, dict):
+            return
+        # Codex records its line kind on the envelope; only usage carries it in the payload.
+        if obj.get("type") == "turn_context" and payload.get("model"):
+            state.usage.model = payload["model"]
+            return
+        if payload.get("type") != "token_count":
+            return
+        info = payload.get("info")
+        if not isinstance(info, dict):
+            return
+        totals = info.get("total_token_usage")
+        if not isinstance(totals, dict):
+            return
+
+        # Codex occasionally emits the same cumulative snapshot again with an empty
+        # last_token_usage block. It is status, not another model call.
+        fingerprint = tuple(int(totals.get(key) or 0) for key in (
+            "input_tokens", "cached_input_tokens", "cache_write_input_tokens",
+            "output_tokens", "reasoning_output_tokens",
+        ))
+        if fingerprint in state.counted:
+            return
+        state.counted.add(fingerprint)
+
+        cached = int(totals.get("cached_input_tokens") or 0)
+        write = int(totals.get("cache_write_input_tokens") or 0)
+        # Codex reports cached reads and writes inside input_tokens; the HUD keeps
+        # them as separate, non-overlapping buckets.
+        state.usage.cache_read = cached
+        state.usage.cache_write = write
+        # Codex has no 1-hour cache tier, so its entire write is the 5-minute rate.
+        state.usage.write_5m = write
+        state.usage.fresh_input = max(
+            0, int(totals.get("input_tokens") or 0) - cached - write
+        )
+        # reasoning_output_tokens is a subset of output_tokens, not an extra bucket.
+        state.usage.output = int(totals.get("output_tokens") or 0)
+
+        window = info.get("model_context_window")
+        if isinstance(window, int) and window > 0:
+            state.usage.context_window = window
+        last = info.get("last_token_usage")
+        if isinstance(last, dict):
+            used = int(last.get("input_tokens") or 0)
+            # A duplicate event can report zero; it would otherwise blank the gauge.
+            if used > 0:
+                state.usage.context_used = used
+                last_cached = int(last.get("cached_input_tokens") or 0)
+                last_write = int(last.get("cache_write_input_tokens") or 0)
+                self._emit(UsageCall(
+                    timestamp=obj.get("timestamp"), recorded_at=time.time(),
+                    model=state.usage.model,
+                    cache_read=last_cached, cache_write=last_write,
+                    write_5m=last_write, write_1h=0,
+                    # last_token_usage.input_tokens includes both cached reads and
+                    # cache writes, just like total_token_usage.input_tokens.
+                    fresh_input=max(0, used - last_cached - last_write),
+                    output=int(last.get("output_tokens") or 0),
+                    reasoning_output=int(last.get("reasoning_output_tokens") or 0),
+                ))
+
+
+def session_root(harness: str) -> Path | None:
+    home = Path.home()
+    if harness == "claude-code":
+        return home / ".claude/projects"
+    if harness == "codex":
+        return home / ".codex/sessions"
+    if harness == "opencode":
+        return home / ".local/share/opencode/opencode.db"
+    return None
+
+
+def flattened(directory: str) -> str:
+    """Claude Code names a project directory after its cwd with separators turned to dashes."""
+    return str(directory).replace("/", "-")
+
+
+def newest_session(harness: str, project_dir: str | None = None,
+                   since: float | None = None, exclude: set | None = None,
+                   provider: bool | None = None) -> str | None:
+    """The transcript to follow.
+
+    A project directory narrows the search to one working directory. `since` restricts it to
+    sessions created after a moment in time, and `exclude` rules out transcripts another pane has
+    already claimed, which is what separates two panes running the same harness there.
+
+    `provider` (Codex only) selects by how the session was launched: True for one routed through
+    PennyRouter, False for a native one. OpenCode owns provider selection in its own config, so
+    its sessions are selected by directory/time regardless of provider."""
+    root = session_root(harness)
+    if harness == "opencode":
+        if not root or not root.exists():
+            return None
+        excluded = exclude or set()
+        try:
+            db = sqlite3.connect(f"file:{root}?mode=ro", uri=True, timeout=0.2)
+            query = "select id, directory, time_created, model from session"
+            params = ()
+            if project_dir:
+                query += " where directory = ?"
+                params = (str(project_dir),)
+            query += " order by time_created desc"
+            for session_id, directory, created_ms, model_raw in db.execute(query, params):
+                target = f"{root}#{session_id}"
+                if target in excluded or (since is not None and created_ms / 1000 <= since):
+                    continue
+                try:
+                    model = json.loads(model_raw or "{}") or {}
+                except ValueError:
+                    model = {}
+                is_penny = str(model.get("providerID") or "").lower() == "pennyrouter"
+                if provider is not None and is_penny != provider:
+                    continue
+                db.close()
+                return target
+            db.close()
+        except sqlite3.Error:
+            return None
+        return None
+    if not root or not root.exists():
+        return None
+    if project_dir and harness == "claude-code":
+        scoped = root / flattened(project_dir)
+        if scoped.exists():
+            root = scoped
+    newest, newest_mtime = None, -1.0
+    excluded = exclude or set()
+    for path in root.rglob("*.jsonl"):
+        if str(path) in excluded:
+            continue
+        try:
+            mtime = path.stat().st_mtime
+        except OSError:
+            continue
+        if mtime <= newest_mtime:
+            continue
+        if since is not None:
+            try:
+                if path.stat().st_ctime < since:
+                    continue
+            except OSError:
+                continue
+        if harness == "claude-code" and provider is False and _claude_is_penny(path):
+            # Legacy fallback only: compare normally assigns Claude's session id before launch.
+            # If a caller does use discovery, never adopt a transcript Penny explicitly claimed.
+            continue
+        if harness == "codex" and (project_dir or provider is not None):
+            # Codex keeps every session in one tree, so its working directory comes from the
+            # session header rather than the path. That header also names the provider, which
+            # is what separates a penny pane's session from a native one's in the same
+            # directory — neither timing nor the path can tell those apart.
+            meta = codex_session_meta(path)
+            if project_dir and meta.get("cwd") != str(project_dir):
+                continue
+            if provider is not None and _is_penny(meta) != provider:
+                continue
+        newest, newest_mtime = str(path), mtime
+    return newest
+
+
+def _is_penny(meta: dict) -> bool:
+    """Whether a Codex session was launched through PennyRouter."""
+    if str(meta.get("model_provider") or "").lower() == "pennyrouter":
+        return True
+    return _codex_is_penny(meta)
+
+
+def _codex_is_penny(meta: dict) -> bool:
+    """Whether the gateway claimed this Codex session as Penny-routed.
+
+    A subscription session reaches Codex's built-in OpenAI provider, which is what gives it the
+    native tool runtime, so its header names `openai` exactly like a native session's. Codex
+    sends the rollout's own conversation UUID on the WebSocket handshake, and the gateway
+    records it there; matching on that id keeps the two apart without timing or ordering.
+    """
+    session_id = str(meta.get("session_id") or "")
+    if not session_id:
+        return False
+    data = Path.home() / ".local" / "share" / "pennyrouter"
+    try:
+        return (data / f"codex-session-{session_id}").exists()
+    except OSError:
+        return False
+
+
+def _claude_is_penny(path) -> bool:
+    """Whether Penny's exact Claude session-id handoff claims this transcript.
+
+    Transcript text is not an ownership marker: native sessions can mention PennyRouter through
+    MCP instructions, and native proxy sessions can use provider-qualified model names. Penny's
+    title-session files contain the actual Claude session ids it launched.
+    """
+    session_id = Path(path).stem
+    data = Path.home() / ".local" / "share" / "pennyrouter"
+    try:
+        entries = data.glob("title-session-*")
+        for entry in entries:
+            try:
+                if entry.read_text().strip() == session_id:
+                    return True
+            except OSError:
+                continue
+    except OSError:
+        pass
+    return False
+
+
+def transcript_for_session(session_id: str, project_dir: str | None = None) -> str | None:
+    """The transcript file belonging to a Claude Code session id.
+
+    Claude Code names a transcript after its session id and records the same id in every line,
+    so the path is derived directly and the field is only consulted if that file is absent.
+    """
+    root = session_root("claude-code")
+    if not root or not session_id:
+        return None
+    search = root
+    if project_dir:
+        scoped = root / flattened(project_dir)
+        if scoped.exists():
+            search = scoped
+    direct = search / f"{session_id}.jsonl"
+    if direct.exists():
+        return str(direct)
+    for path in sorted(search.rglob("*.jsonl"), key=lambda p: -p.stat().st_mtime):
+        try:
+            with open(path) as handle:
+                first = handle.readline()
+        except OSError:
+            continue
+        try:
+            if json.loads(first).get("sessionId") == session_id:
+                return str(path)
+        except ValueError:
+            continue
+    return None
+
+
+def codex_session_meta(path) -> dict:
+    """The session_meta payload Codex writes as a session's first line.
+
+    It carries both the working directory and the provider the session was launched against,
+    which is what tells a penny pane's rollout apart from a native one's in the same directory.
+    """
+    try:
+        with open(path, "rb") as handle:
+            first = handle.readline()
+        return json.loads(first).get("payload") or {}
+    except (OSError, ValueError):
+        return {}
+
+
+def related_transcripts(harness: str, root_path: str | None) -> list[str]:
+    """Return a root transcript plus its known subagent descendants.
+
+    Claude stores sidechains below the root transcript. Codex stores each thread beside the
+    others and records the parent id in session_meta, so lineage must be resolved explicitly.
+    """
+    if not root_path or not INCLUDE_SUBAGENT_USAGE:
+        return [root_path] if root_path else []
+    root = Path(root_path)
+    if harness == "claude-code":
+        child_dir = root.with_suffix("") / "subagents"
+        children = sorted(str(path) for path in child_dir.rglob("*.jsonl")) if child_dir.exists() else []
+        return [root_path, *children]
+    if harness != "codex":
+        return [root_path]
+
+    root_meta = codex_session_meta(root_path)
+    root_ids = {str(root_meta.get("id") or root_meta.get("session_id") or "")}
+    root_ids.discard("")
+    discovered = [root_path]
+    session_root_path = session_root("codex")
+    if not session_root_path or not session_root_path.exists() or not root_ids:
+        return discovered
+    candidates = []
+    for path in session_root_path.rglob("*.jsonl"):
+        if str(path) == root_path:
+            continue
+        meta = codex_session_meta(path)
+        session_id = str(meta.get("id") or meta.get("session_id") or "")
+        parent_id = str(meta.get("parent_thread_id") or meta.get("forked_from_id") or "")
+        candidates.append((str(path), session_id, parent_id))
+    while True:
+        added = False
+        for path, session_id, parent_id in candidates:
+            if path in discovered or parent_id not in root_ids:
+                continue
+            discovered.append(path)
+            if session_id:
+                root_ids.add(session_id)
+            added = True
+        if not added:
+            break
+    return discovered
+
+
+def codex_cwd(path) -> str | None:
+    """The working directory Codex recorded in a session's first line."""
+    return codex_session_meta(path).get("cwd")
+
+
+def session_id_for_transcript(harness: str, path: str | None) -> str | None:
+    """Return the stable client session id stored by either harness."""
+    if not path:
+        return None
+    if harness == "claude-code":
+        return Path(path).stem
+    if harness == "codex":
+        meta = codex_session_meta(path)
+        return str(meta.get("session_id") or meta.get("id") or Path(path).stem.replace("rollout-", ""))
+    if harness == "opencode" and "#" in path:
+        return path.split("#", 1)[1]
+    return None