pennyrouter 0.3.19 → 0.3.21

package/standalone/pane.py

@@ -0,0 +1,430 @@
+#!/usr/bin/env python3
+"""Runs one harness with a token readout pinned above it.
+
+Invoked inside a tmux pane by penny-compare. The readout owns the top rows via a scroll region,
+so the harness below scrolls normally and its own rendering is untouched.
+"""
+
+from __future__ import annotations
+
+import argparse
+from dataclasses import asdict
+import json
+import os
+import shutil
+import subprocess
+import sys
+import threading
+import time
+import traceback
+import uuid
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+import hud
+import pricing
+from usage import (UsageReader, newest_session, session_id_for_transcript, session_root,
+                   related_transcripts, transcript_for_session)
+
+# Each variant is the command that already launches that session correctly. The penny variants
+# delegate to `penny`, which owns gateway auth and routing; reproducing that here is what made a
+# pane prompt about a foreign upstream key.
+VARIANTS = {
+    "claude":       ("claude-code", ["claude"],          "NATIVE"),
+    "penny-claude": ("claude-code", ["penny", "claude"], "PENNY"),
+    "codex":        ("codex",       ["codex"],           "NATIVE"),
+    "penny-codex":  ("codex",       ["penny", "codex"],  "PENNY"),
+    "opencode":     ("opencode",    ["opencode"],        "NATIVE"),
+    "penny-opencode": ("opencode", ["penny", "opencode"], "PENNY"),
+}
+
+# Upstream env a NATIVE pane never inherits. The ambient shell may already point every
+# Anthropic client somewhere; a NATIVE pane decides its own upstream, so it starts from a
+# clean slate rather than whatever happens to be exported.
+NATIVE_STRIP_ENV = ("ANTHROPIC_BASE_URL", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_KEY")
+
+# With --native-proxy, the NATIVE pane routes through these instead of going direct.
+NATIVE_PROXY_URL_ENV = "PENNY_COMPARE_NATIVE_BASE_URL"
+NATIVE_PROXY_TOKEN_ENV = "PENNY_COMPARE_NATIVE_AUTH_TOKEN"
+
+
+class Readout:
+    """Paints the fixed block at the top of the pane."""
+
+    def __init__(self, harness, title, delta_path=None, provider=None, label=None,
+                 session_id=None, side=None, variant=None, pricing_profile="standard"):
+        self.harness = harness
+        # The pane's fixed role ("NATIVE"/"PENNY"). It keys the files the two panes use to
+        # coordinate, so it stays constant no matter what the readout is labelled.
+        self.title = title
+        # What the readout actually shows, which the user may rename.
+        self.label = label or title
+        self.delta_path = delta_path
+        # For Codex, whether this pane's session should be the PennyRouter-routed one.
+        self.provider = provider
+        # Claude accepts a caller-supplied session id, so its transcript can be selected exactly
+        # instead of inferred from timing, model names, or Penny's global title handoff files.
+        self.session_id = session_id
+        self.side = side or title.lower()
+        self.variant = variant
+        self.pricing_profile = pricing_profile
+        self.reader = UsageReader(harness, on_call=self._record_call)
+        self.call_cost = 0.0
+        self.call_cost_baseline = 0.0
+        self.counter = hud.CountUp()
+        # Counts recorded before the last reset; subtracted from what is shown.
+        self.baseline = (0, 0, 0, 0)
+        self.write_5m_baseline = 0
+        self.write_1h_baseline = 0
+        self.reset_pending = False
+        self._reset_at = 0.0
+        self.reported_error = False
+        self.path = None
+        # Only sessions created after the pane launched can be this pane's.
+        self.started = time.time()
+        self._preexisting = self._existing_sessions()
+        self.stop = threading.Event()
+        self.last = time.monotonic()
+        self._publish_metadata()
+
+    def _state_path(self, name):
+        state = os.environ.get("PENNY_COMPARE_STATE")
+        return os.path.join(state, name) if state else None
+
+    def _publish_metadata(self):
+        path = self._state_path(f"pane-{self.side}.json")
+        if not path:
+            return
+        payload = {
+            "side": self.side, "role": self.title, "label": self.label,
+            "variant": self.variant, "harness": self.harness, "cwd": os.getcwd(),
+            "pricing_profile": self.pricing_profile,
+            "session_id": self.session_id,
+            "transcript_path": self.path,
+        }
+        try:
+            with open(path, "w") as handle:
+                json.dump(payload, handle)
+        except OSError:
+            pass
+
+    def _record_call(self, call):
+        # The readout's cost is this running total, so it accrues for every call. Writing the
+        # per-call log needs a state directory; a session started without one still shows cost.
+        context_tokens = call.fresh_input + call.cache_read + call.cache_write
+        cost = pricing.cost(
+            call.model, call.fresh_input, call.cache_read, call.output,
+            write_5m=call.write_5m, write_1h=call.write_1h,
+            profile=self.pricing_profile, context_tokens=context_tokens,
+        )
+        self.call_cost += cost
+
+        path = self._state_path(f"calls-{self.side}.jsonl")
+        if not path:
+            return
+        payload = asdict(call)
+        payload["total"] = call.total
+        payload["pricing_profile"] = self.pricing_profile
+        rate_names = ("input", "cache_read", "cache_write_5m", "cache_write_1h", "output")
+        payload["context_tokens"] = context_tokens
+        payload["rates_per_million"] = dict(zip(
+            rate_names,
+            (rate * 1_000_000 for rate in pricing.rates_for(
+                call.model, self.pricing_profile, context_tokens)),
+        ))
+        payload["cost"] = cost
+        try:
+            with open(path, "a") as handle:
+                handle.write(json.dumps(payload, separators=(",", ":")) + "\n")
+        except OSError:
+            pass
+
+    def resolve(self):
+        """Pin this pane to the transcript its own harness created.
+
+        Both panes run in the same working directory, so the newest session there is ambiguous
+        between them, and resolving by timing alone can label the two backwards.
+
+        Claude is launched with an explicit session id, making both its native and Penny
+        transcript paths deterministic. Codex has no equivalent launch flag, but records the
+        provider in its session header, so each Codex pane takes the session matching how it was
+        started.
+        """
+        if self.path and os.path.exists(self.path):
+            return self.path
+
+        if self.session_id:
+            self.path = transcript_for_session(self.session_id, os.getcwd())
+            if self.path:
+                self._publish(self.path)
+            return self.path
+
+        # A transcript that already existed when the pane launched belongs to some other
+        # session, so require one created afterwards rather than adopting the newest on disk.
+        taken = self._peers() | self._preexisting
+        candidate = newest_session(self.harness, os.getcwd(),
+                                   since=self.started, exclude=taken,
+                                   provider=self.provider)
+        if candidate and not self._claim(candidate):
+            return None
+        self.path = candidate
+        if candidate:
+            self._publish(candidate)
+        return self.path
+
+    def _existing_sessions(self) -> set:
+        """Transcripts already on disk, which cannot belong to a session starting now."""
+        root = session_root(self.harness)
+        if not root or not root.exists():
+            return set()
+        return {str(p) for p in root.rglob("*.jsonl")}
+
+    def _publish(self, path):
+        """Record which transcript this pane resolved, so the other can rule it out."""
+        state = os.environ.get("PENNY_COMPARE_STATE")
+        if not state:
+            return
+        try:
+            self.session_id = session_id_for_transcript(self.harness, path)
+            with open(os.path.join(state, f"session-{self.side}"), "w") as handle:
+                handle.write(path)
+            self._publish_metadata()
+        except OSError:
+            pass
+
+    def _peers(self) -> set:
+        """Transcripts the other pane has already resolved."""
+        state = os.environ.get("PENNY_COMPARE_STATE")
+        if not state:
+            return set()
+        found = set()
+        try:
+            for name in os.listdir(state):
+                if not name.startswith("session-") or name.endswith(self.side):
+                    continue
+                with open(os.path.join(state, name)) as handle:
+                    value = handle.read().strip()
+                if value:
+                    found.add(value)
+        except OSError:
+            pass
+        return found
+
+    def _claim(self, path):
+        """Record this transcript as taken, so the other pane looks past it."""
+        state = os.environ.get("PENNY_COMPARE_STATE")
+        if not state:
+            return True
+        marker = os.path.join(state, "claimed")
+        try:
+            with open(marker, "a+") as handle:
+                handle.seek(0)
+                taken = {line.strip() for line in handle if line.strip()}
+                if path in taken:
+                    return False
+                handle.write(path + "\n")
+            return True
+        except OSError:
+            return True
+
+    def _check_reset(self):
+        """Zero the counters when the launcher signals a reset.
+
+        A comparison is only fair once both sides have a warm cache, so the run that matters is
+        usually not the first one. Rebasing on the counts already recorded lets the readout start
+        from zero without restarting either session.
+        """
+        state = os.environ.get("PENNY_COMPARE_STATE")
+        if not state:
+            return
+        marker = os.path.join(state, "reset")
+        try:
+            stamp = os.stat(marker).st_mtime
+        except OSError:
+            return
+        if stamp <= self._reset_at:
+            return
+        self._reset_at = stamp
+        self.reset_pending = True
+        self.counter.reset()
+
+    def paint(self):
+        self._check_reset()
+        path = self.resolve()
+        transcripts = related_transcripts(self.harness, path)
+        usage = self.reader.read_many(transcripts)
+        if usage is None:
+            targets = {"cache_read": 0, "cache_write": 0, "fresh_input": 0, "output": 0}
+            cost, model = 0.0, ""
+        else:
+            if self.reset_pending:
+                # Snapshot what is already recorded so counting restarts from here.
+                self.baseline = (usage.cache_read, usage.cache_write,
+                                 usage.fresh_input, usage.output)
+                self.write_5m_baseline = usage.write_5m
+                self.write_1h_baseline = usage.write_1h
+                self.call_cost_baseline = self.call_cost
+                self.reset_pending = False
+            targets = {
+                "cache_read": usage.cache_read - self.baseline[0],
+                "cache_write": usage.cache_write - self.baseline[1],
+                "fresh_input": usage.fresh_input - self.baseline[2],
+                "output": usage.output - self.baseline[3],
+            }
+            # Cost reflects what is displayed, not the session lifetime total, so it drops to
+            # zero on reset along with the counts above it.
+            # Costs are accumulated per call because GovCloud GPT rates change when that call's
+            # input context exceeds 272K. Pricing cumulative session totals would select the
+            # expensive tier merely because several smaller calls added up past the threshold.
+            cost = self.call_cost - self.call_cost_baseline
+            model = usage.model
+
+        # A transcript exists from session start but carries no usage until a turn completes, so
+        # counting begins at the first usage block rather than inheriting what was on disk.
+        waiting = sum(targets.values()) == 0
+
+        now = time.monotonic()
+        rolled = self.counter.advance(targets, now - self.last)
+        self.last = now
+
+        delta = None
+        if self.delta_path:
+            try:
+                with open(self.delta_path) as handle:
+                    delta = handle.read().strip() or None
+            except OSError:
+                delta = None
+
+        lines = hud.render(self.label, model, rolled, cost, delta, waiting=waiting)
+        # Right-align to the pane. Erasing runs from the readout's own column to the right
+        # edge, so a value that shrinks leaves no digits behind while the harness output to
+        # the left of the block stays untouched.
+        size = shutil.get_terminal_size()
+        column = max(1, size.columns - hud.WIDTH - 1)
+
+        # The full-screen renderer repaints its entire frame while a turn streams. That can
+        # erase this overlay even when the HUD data has not changed, so the caller must repaint
+        # it continuously rather than treating an identical frame as already visible.
+        # DECSC/DECRC (ESC 7 / ESC 8) rather than ESC[s / ESC[u: the harness below is scrolling
+        # its own region continuously, and only the DEC pair restores the cursor relative to it.
+        # Reassert the scroll region first so the readout's rows stay excluded from that scroll
+        # even after the harness resets it.
+        out = ["\0337", f"\033[{hud.ROWS + 1};{size.lines}r"]
+        for index, line in enumerate(lines):
+            out.append(f"\033[{index + 1};{column}H\033[K{line}")
+        out.append("\0338")
+        sys.stdout.write("".join(out))
+        sys.stdout.flush()
+
+        # Publish this pane's numbers so the other pane can show a comparison.
+        if usage is not None:
+            state = os.environ.get("PENNY_COMPARE_STATE")
+            if state:
+                try:
+                    with open(os.path.join(state, f"{self.title.lower()}.txt"), "w") as handle:
+                        handle.write(f"{cost}\n{sum(targets.values())}\n{targets['cache_write']}\n")
+                except OSError:
+                    pass
+
+    def loop(self):
+        while not self.stop.wait(1 / 30):
+            try:
+                self.paint()
+            except Exception:
+                # A painting error must not take down the harness, but silently swallowing it
+                # leaves the readout stuck looking like it is merely waiting. Report it once.
+                if not self.reported_error:
+                    self.reported_error = True
+                    traceback.print_exc(file=sys.stderr)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("variant", choices=sorted(VARIANTS))
+    parser.add_argument("--delta-file", default=None)
+    parser.add_argument("--model", default=None,
+                        help="start both panes on one model so the comparison is like for like")
+    parser.add_argument("--label", default=None,
+                        help="name shown on this pane's readout (default: NATIVE or PENNY)")
+    parser.add_argument("--side", choices=("left", "right", "single"),
+                        help="stable pane position used by the final report")
+    parser.add_argument("--pricing-profile", choices=("standard", "govcloud"),
+                        default="standard")
+    parser.add_argument("--native-proxy", action="store_true",
+                        help=f"route the NATIVE pane through {NATIVE_PROXY_URL_ENV} / "
+                             f"{NATIVE_PROXY_TOKEN_ENV} instead of going direct")
+    args = parser.parse_args()
+
+    harness, command, title = VARIANTS[args.variant]
+    session_id = None
+    if harness == "claude-code":
+        session_id = str(uuid.uuid4())
+        command = command + ["--session-id", session_id]
+    if args.model:
+        # Passed to the harness rather than set in the environment, so Penny's own model
+        # configuration is not overridden on the penny side.
+        command = command + ["--model", args.model]
+
+    rows = shutil.get_terminal_size().lines
+    # Reserve the top rows for the readout and confine scrolling to what is left.
+    sys.stdout.write(f"\033[2J\033[{hud.ROWS + 1};{rows}r\033[{hud.ROWS + 1};1H")
+    sys.stdout.flush()
+
+    # Codex and OpenCode have no caller-assigned session id, so they identify their transcript by the provider
+    # recorded in the session header instead.
+    is_penny = command[0] == "penny"
+    readout = Readout(harness, title, args.delta_file,
+                      # OpenCode owns provider/model selection inside its own TUI. The
+                      # penny-opencode spelling is only a launch convenience, not a reliable
+                      # provider identity like Codex's session header is.
+                      provider=is_penny if harness == "codex" else None,
+                      label=args.label,
+                      session_id=session_id,
+                      side=args.side,
+                      variant=args.variant,
+                      pricing_profile=args.pricing_profile)
+    thread = threading.Thread(target=readout.loop, daemon=True)
+    thread.start()
+
+    # The alt-screen renderer owns a full-screen buffer rather than appending to the main
+    # screen, so a pane that is not the active one still repaints as output arrives. The
+    # classic renderer leaves the inactive side blank until its turn ends, which reads as
+    # that harness being broken rather than merely unfocused.
+    env = dict(os.environ)
+    env.setdefault("CLAUDE_CODE_NO_FLICKER", "1")
+
+    if title == "NATIVE":
+        env = {k: v for k, v in env.items() if k not in NATIVE_STRIP_ENV}
+        if args.native_proxy:
+            proxy_url = os.environ.get(NATIVE_PROXY_URL_ENV)
+            proxy_token = os.environ.get(NATIVE_PROXY_TOKEN_ENV)
+            if proxy_url and proxy_token:
+                env["ANTHROPIC_BASE_URL"] = proxy_url
+                env["ANTHROPIC_AUTH_TOKEN"] = proxy_token
+            else:
+                # Missing config must not take the pane down: a pane that exits here leaves
+                # a one-sided comparison with no visible reason. Say so and run direct.
+                missing = " and ".join(
+                    name for name, value in ((NATIVE_PROXY_URL_ENV, proxy_url),
+                                             (NATIVE_PROXY_TOKEN_ENV, proxy_token))
+                    if not value
+                )
+                sys.stderr.write(
+                    f"penny-compare: --native-proxy needs {missing}; "
+                    f"running the NATIVE pane direct instead\n",
+                )
+
+    try:
+        subprocess.call(command, env=env)
+    except FileNotFoundError:
+        sys.stderr.write(f"penny-compare: {command[0]} not found on PATH\n")
+        return 1
+    finally:
+        readout.stop.set()
+        # Release the scroll region so the pane is left usable.
+        sys.stdout.write("\033[r\033[2J\033[H")
+        sys.stdout.flush()
+
+
+if __name__ == "__main__":
+    sys.exit(main() or 0)