pcode 0.1.0__py3-none-any.whl

pcode/usage.py

@@ -0,0 +1,293 @@
+"""`/usage`: plan limits and spend for the Claude Code and Codex logins, on demand.
+
+Both sources are undocumented endpoints the vendors' own CLIs call, read with
+the logins pcode already runs on, so no admin key is involved:
+
+- Claude: `GET api.anthropic.com/api/oauth/usage` with Claude Code's OAuth
+  token. Subscription seats get `limits` (session and weekly percentages);
+  Enterprise seats billed at API rates get `spend` (used of the monthly cap).
+- Codex: `GET chatgpt.com/backend-api/wham/usage` through the Codex provider's
+  auth, which injects the token, refreshes it on a 401, and never writes the
+  Codex CLI's `auth.json`.
+
+Claude Code's token is read, never refreshed: its refresh tokens are single
+use, so rotating one here would sign Claude Code out.
+"""
+
+import asyncio
+import hashlib
+import json
+import os
+import subprocess
+import sys
+import time
+from collections.abc import Callable
+from pathlib import Path
+
+import httpx2
+
+CLAUDE_USAGE_URL = "https://api.anthropic.com/api/oauth/usage"
+CODEX_USAGE_URL = "https://chatgpt.com/backend-api/wham/usage"
+TIMEOUT = httpx2.Timeout(15, connect=5)
+KEYCHAIN_SERVICE = "Claude Code-credentials"
+
+CLAUDE_LIMIT_LABELS = {
+    "session": "Session (5h)",
+    "weekly_all": "Weekly",
+    "weekly_scoped": "Weekly",
+}
+CODEX_WINDOW_LABELS = {5 * 3600: "Session (5h)", 7 * 86400: "Weekly"}
+
+
+class UsageUnavailable(Exception):
+    """A provider's usage could not be read; the message says why."""
+
+
+def usage_report() -> list[str]:
+    """Transcript lines for every provider; one failing never hides the other."""
+    lines: list[str] = []
+    sources: tuple[tuple[str, Callable[[], list[str]]], ...] = (
+        ("Claude Code", claude_usage),
+        ("Codex", codex_usage),
+    )
+    for title, fetch in sources:
+        try:
+            body = fetch()
+        except UsageUnavailable as error:
+            body = [str(error)]
+        lines.append(title)
+        lines.extend(f"  {line}" for line in body)
+    return lines
+
+
+# Claude Code
+
+
+def keychain_service(config_dir: str | None) -> str:
+    """Claude Code's Keychain item for a config dir.
+
+    Setting `CLAUDE_CONFIG_DIR` suffixes the name with a hash of its value,
+    even when the value is the default `~/.claude`, so the unsuffixed item can
+    belong to a different (often stale) login.
+    """
+    if not config_dir:
+        return KEYCHAIN_SERVICE
+    return f"{KEYCHAIN_SERVICE}-{hashlib.sha256(config_dir.encode()).hexdigest()[:8]}"
+
+
+def claude_login(config_dir: str | None = None) -> dict:
+    """Claude Code's `claudeAiOauth` record: the macOS Keychain, else its credentials file."""
+    if config_dir is None:
+        config_dir = os.environ.get("CLAUDE_CONFIG_DIR", "").strip() or None
+    raw = ""
+    if sys.platform == "darwin":
+        try:
+            raw = subprocess.run(
+                ["security", "find-generic-password", "-s", keychain_service(config_dir), "-w"],
+                capture_output=True,
+                text=True,
+                timeout=10,
+            ).stdout
+        except (OSError, subprocess.SubprocessError):
+            raw = ""
+    if not raw.strip():
+        path = Path(config_dir).expanduser() if config_dir else Path.home() / ".claude"
+        try:
+            raw = (path / ".credentials.json").read_text()
+        except OSError:
+            raw = ""
+    try:
+        login = json.loads(raw).get("claudeAiOauth") if raw.strip() else None
+    except (ValueError, AttributeError):
+        login = None
+    if not isinstance(login, dict) or not login.get("accessToken"):
+        where = f" for {config_dir}" if config_dir else ""
+        raise UsageUnavailable(f"No Claude Code login{where}; run `claude` and sign in.")
+    return login
+
+
+def claude_usage() -> list[str]:
+    from pcode.auth import oauth_user_agent
+
+    login = claude_login()
+    expires = login.get("expiresAt")
+    if isinstance(expires, int | float) and expires / 1000 < time.time():
+        raise UsageUnavailable(
+            "Claude Code's token has expired; run `claude` (or a claude: model) to refresh it."
+        )
+    try:
+        response = httpx2.get(
+            CLAUDE_USAGE_URL,
+            timeout=TIMEOUT,
+            headers={
+                "Authorization": f"Bearer {login['accessToken']}",
+                "anthropic-beta": "oauth-2025-04-20",
+                # A generic client lands in a much stricter rate-limit bucket.
+                "User-Agent": oauth_user_agent(),
+            },
+        )
+    except httpx2.HTTPError as error:
+        raise UsageUnavailable(f"Request failed: {type(error).__name__}") from None
+    if response.status_code in (401, 403):
+        raise UsageUnavailable("Claude Code's login was rejected; run `claude` to sign in again.")
+    if response.status_code == 429:
+        raise UsageUnavailable("The usage endpoint is rate limiting; try again in a minute.")
+    if response.status_code != 200:
+        raise UsageUnavailable(f"The usage endpoint answered HTTP {response.status_code}.")
+    plan = login.get("subscriptionType")
+    return ([f"Plan: {plan}"] if plan else []) + claude_lines(response.json())
+
+
+def claude_lines(data: dict, now: float | None = None) -> list[str]:
+    now = time.time() if now is None else now
+    lines = []
+    limits = data.get("limits")
+    if limits:
+        for limit in limits:
+            kind = limit.get("kind", "")
+            label = CLAUDE_LIMIT_LABELS.get(kind) or kind.replace("_", " ").capitalize()
+            model = ((limit.get("scope") or {}).get("model") or {}).get("display_name")
+            if model:
+                label = f"{label} ({model})"
+            resets = iso_time(limit.get("resets_at"))
+            lines.append(percent_line(label, limit.get("percent"), resets, now))
+    else:
+        # Older responses carry only the windows, with no `limits` list.
+        for key, label in (("five_hour", "Session (5h)"), ("seven_day", "Weekly")):
+            window = data.get(key)
+            if window:
+                resets = iso_time(window.get("resets_at"))
+                lines.append(percent_line(label, window.get("utilization"), resets, now))
+    spend = data.get("spend") or {}
+    used, limit = money(spend.get("used")), money(spend.get("limit"))
+    if used is not None and limit is not None:
+        line = f"Spend: {used} of {limit}"
+        if spend.get("percent") is not None:
+            line += f" ({spend['percent']:.0f}%)"
+        if not spend.get("enabled"):
+            reason = spend.get("disabled_reason")
+            line += f", off ({reason.replace('_', ' ')})" if reason else ", off"
+        lines.append(line)
+    return lines or ["No limits reported for this login."]
+
+
+def money(amount: dict | None) -> str | None:
+    if not amount or amount.get("amount_minor") is None:
+        return None
+    value = amount["amount_minor"] / 10 ** amount.get("exponent", 2)
+    currency = amount.get("currency", "USD")
+    return f"${value:,.2f}" if currency == "USD" else f"{value:,.2f} {currency}"
+
+
+# Codex
+
+
+def codex_usage() -> list[str]:
+    try:
+        data = asyncio.run(_fetch_codex())
+    except UsageUnavailable:
+        raise
+    except httpx2.HTTPError as error:
+        raise UsageUnavailable(f"Request failed: {type(error).__name__}") from None
+    except Exception as error:  # noqa: BLE001 - credential errors are the provider's own types.
+        raise UsageUnavailable(str(error).split("\n")[0]) from None
+    return codex_lines(data)
+
+
+async def _fetch_codex() -> dict:
+    from pydantic_ai.providers.openai_codex import OpenAICodexProvider
+
+    from pcode.codex_login import credential_source
+
+    # The same route Codex model requests take: PCODE_LLM_PROXY when set.
+    proxy = os.environ.get("PCODE_LLM_PROXY", "").strip()
+    options = {"proxy": proxy, "trust_env": False} if proxy else {}
+    async with httpx2.AsyncClient(timeout=TIMEOUT, **options) as client:
+        source = credential_source()
+        # Attaches the provider's auth to `client`; without a pcode login it
+        # reads the Codex CLI's auth.json read-only.
+        if source is not None:
+            OpenAICodexProvider(credential_source=source, http_client=client)
+        else:
+            OpenAICodexProvider(http_client=client)
+        try:
+            response = await client.get(CODEX_USAGE_URL, headers={"Accept": "application/json"})
+        except (httpx2.ConnectError, httpx2.ConnectTimeout, httpx2.ProxyError) as error:
+            route = " through PCODE_LLM_PROXY" if proxy else ""
+            raise UsageUnavailable(
+                f"Could not reach chatgpt.com{route} ({type(error).__name__})."
+            ) from None
+    if response.status_code == 403 and "json" not in response.headers.get("content-type", ""):
+        hint = "" if proxy else "; set PCODE_LLM_PROXY if chatgpt.com needs a proxy here"
+        raise UsageUnavailable(f"chatgpt.com blocked the request (HTTP 403){hint}.")
+    if response.status_code != 200:
+        raise UsageUnavailable(f"The usage endpoint answered HTTP {response.status_code}.")
+    return response.json()
+
+
+def codex_lines(data: dict, now: float | None = None) -> list[str]:
+    now = time.time() if now is None else now
+    lines = []
+    if data.get("plan_type"):
+        lines.append(f"Plan: {data['plan_type']}")
+    lines.extend(codex_window_lines(data.get("rate_limit"), "", now))
+    for extra in data.get("additional_rate_limits") or ():
+        name = extra.get("limit_name") or extra.get("metered_feature") or "Other"
+        lines.extend(codex_window_lines(extra.get("rate_limit"), f"{name} ", now))
+    credits = data.get("credits") or {}
+    if credits.get("unlimited"):
+        lines.append("Credits: unlimited")
+    elif credits.get("has_credits") and credits.get("balance") is not None:
+        lines.append(f"Credits: {credits['balance']}")
+    spend = data.get("spend_control") or {}
+    cap = spend.get("individual_limit")
+    if cap:
+        line = f"Spend: {cap.get('used')} of {cap.get('limit')} ({cap.get('used_percent', 0)}%)"
+        lines.append(line + (", limit reached" if spend.get("reached") else ""))
+    return lines or ["No limits reported for this login."]
+
+
+def codex_window_lines(details: dict | None, prefix: str, now: float) -> list[str]:
+    lines = []
+    for key in ("primary_window", "secondary_window"):
+        window = (details or {}).get(key)
+        if not window:
+            continue
+        seconds = window.get("limit_window_seconds") or 0
+        label = CODEX_WINDOW_LABELS.get(seconds) or f"{duration(seconds)} window"
+        lines.append(
+            percent_line(prefix + label, window.get("used_percent"), window.get("reset_at"), now)
+        )
+    return lines
+
+
+# Formatting
+
+
+def percent_line(label: str, percent: float | None, resets: float | None, now: float) -> str:
+    line = f"{label}: {percent:.0f}% used" if percent is not None else f"{label}: unknown"
+    if resets is not None and resets > now:
+        line += f", resets in {duration(resets - now)}"
+    return line
+
+
+def iso_time(value: str | None) -> float | None:
+    from datetime import datetime
+
+    if not value:
+        return None
+    try:
+        return datetime.fromisoformat(value).timestamp()
+    except ValueError:
+        return None
+
+
+def duration(seconds: float) -> str:
+    minutes = int(seconds // 60)
+    days, minutes = divmod(minutes, 1440)
+    hours, minutes = divmod(minutes, 60)
+    if days:
+        return f"{days}d {hours}h" if hours else f"{days}d"
+    if hours:
+        return f"{hours}h {minutes}m" if minutes else f"{hours}h"
+    return f"{max(minutes, 1)}m"

pcode/word_wrap.py

@@ -0,0 +1,107 @@
+"""Word-aware soft wrapping for the editor buffer.
+
+prompt_toolkit wraps a buffer line strictly by column: the character that no
+longer fits moves to the next row, so a word straddles the wrap point. There is
+no word-wrap option on ``Window``/``BufferControl`` -- the break is decided
+character by character deep inside ``Window._copy_body``.
+
+Rather than reimplement that rendering, this pads the *display* text: when a
+word would straddle the wrap point, enough spaces are inserted before it to
+fill the current row, so prompt_toolkit's own column wrapping lands on the word
+boundary. The buffer text is untouched; the cursor mappings below keep editing
+positions correct, the same technique ``TabsProcessor`` uses to expand tabs.
+"""
+
+from __future__ import annotations
+
+from prompt_toolkit.layout.processors import Processor, Transformation, TransformationInput
+from prompt_toolkit.layout.utils import explode_text_fragments
+from prompt_toolkit.utils import get_cwidth
+
+__all__ = ["WordWrapProcessor", "wrap_padding"]
+
+
+def wrap_padding(text: str, width: int) -> dict[int, int]:
+    """Map index -> number of spaces to insert before ``text[index]``.
+
+    ``width`` is the number of columns available for the text itself (the line
+    prefix is already subtracted). Words wider than a full row are left alone:
+    they have to be split somewhere, and padding them would only waste a row.
+    """
+    if width < 2 or not text:
+        return {}
+
+    padding: dict[int, int] = {}
+    column = 0
+    index = 0
+    length = len(text)
+    while index < length:
+        char = text[index]
+        starts_word = not char.isspace() and (index == 0 or text[index - 1].isspace())
+        # A word starting at column 0, or one whose row is already full, wraps
+        # correctly on its own.
+        if starts_word and 0 < column < width:
+            end = index
+            while end < length and not text[end].isspace():
+                end += 1
+            word = get_cwidth(text[index:end])
+            if column + word > width >= word:
+                padding[index] = width - column
+                column = 0
+        char_width = get_cwidth(char)
+        if column + char_width > width:
+            column = 0
+        column += char_width
+        index += 1
+    return padding
+
+
+class WordWrapProcessor(Processor):
+    """Pad display lines so soft wraps fall between words.
+
+    :param prefix_width: Columns taken by the window's line prefix (the prompt
+        and its continuation), which ``TransformationInput.width`` includes.
+    """
+
+    def __init__(self, prefix_width: int = 0) -> None:
+        self.prefix_width = prefix_width
+
+    def apply_transformation(self, ti: TransformationInput) -> Transformation:
+        fragments = explode_text_fragments(ti.fragments)
+        padding = wrap_padding(
+            "".join(text for _, text, *_ in fragments), ti.width - self.prefix_width
+        )
+        if not padding:
+            return Transformation(ti.fragments)
+
+        result = []
+        positions = {}
+        position = 0
+        for index, fragment in enumerate(fragments):
+            pad = padding.get(index)
+            if pad:
+                result.append(("", " " * pad))
+                position += pad
+            positions[index] = position
+            result.append(fragment)
+            position += 1
+        # The cursor can also sit right after the line, and one past that.
+        positions[len(fragments)] = position
+        positions[len(fragments) + 1] = position + 1
+
+        def source_to_display(from_position: int) -> int:
+            return positions.get(from_position, from_position)
+
+        def display_to_source(display_position: int) -> int:
+            reversed_positions = {v: k for k, v in positions.items()}
+            while display_position >= 0:
+                if display_position in reversed_positions:
+                    return reversed_positions[display_position]
+                display_position -= 1
+            return 0
+
+        return Transformation(
+            result,
+            source_to_display=source_to_display,
+            display_to_source=display_to_source,
+        )

pcode/worker_ui.py

@@ -0,0 +1,235 @@
+"""Read-only popup following delegated agents' own streams; never steers them."""
+
+import asyncio
+
+from prompt_toolkit.application import Application, get_app
+from prompt_toolkit.document import Document
+from prompt_toolkit.filters import Always, has_focus
+from prompt_toolkit.key_binding import KeyBindings
+from prompt_toolkit.key_binding.bindings.focus import focus_next, focus_previous
+from prompt_toolkit.layout import DynamicContainer, HSplit, Layout, VSplit
+from prompt_toolkit.layout.dimension import Dimension
+from prompt_toolkit.widgets import Label, TextArea
+from rich.markdown import Markdown
+from rich.text import Text
+from rich.theme import Theme
+
+from pcode.frame import Frame
+from pcode.popup_ui import (
+    RichPane,
+    bind_list_paging,
+    list_pane_height,
+    popup_container,
+    popup_mouse,
+    popup_style,
+)
+from pcode.prefix_keys import PrefixKeys
+from pcode.runtime import ToolSummary
+from pcode.session_ui import literal
+from pcode.task_prompt import TaskPrompt
+from pcode.tool_display import label, plain
+from pcode.tool_panel import ACTIVE_TASK_ICON, AGENT_ICON, plan_row
+from pcode.workers import Worker, Workers
+
+# Streaming should look live without repainting the pane every token.
+REFRESH_SECONDS = 0.3
+
+
+def row(worker: Worker, width: int = 90) -> str:
+    """One list line: who, how long, what state, and the assignment."""
+    name = worker.agent[:1].upper() + worker.agent[1:]
+    task = plain(" ".join(worker.task.split()), width)
+    return f"{AGENT_ICON} {name} · {worker.elapsed():.0f}s · {worker.state()} · {task}"
+
+
+def tool_line(event) -> Text:
+    if isinstance(event, ToolSummary):
+        icon, style = ("✗", "red") if event.failed else ("✓", "dim")
+        elapsed = f" · {event.elapsed_seconds:.1f}s" if event.elapsed_seconds is not None else ""
+    else:
+        icon, style, elapsed = "⟳", "bold", ""
+    detail = plain(event.detail, limit=None)
+    return Text(f"{icon} {label(event.name)} · {detail}{elapsed}", style=style)
+
+
+def details(worker: Worker | None, *, code_theme: str, show_thinking: bool) -> list:
+    """The worker's assignment, plan, then everything it said and did, in order."""
+    if worker is None:
+        return [Text("Agents appear here once the model delegates a task.", style="dim")]
+    blocks: list = [
+        TaskPrompt(literal(worker.task)),
+        Text(f"  {worker.agent} · {worker.elapsed():.0f}s · {worker.state()}", style="dim"),
+    ]
+    if worker.plan:
+        done = sum(item.get("status") == "completed" for item in worker.plan)
+        blocks += [Text(""), Text(f"Tasks {done}/{len(worker.plan)}", style="bold")]
+        for item in worker.plan:
+            style, line = plan_row(item, ACTIVE_TASK_ICON)
+            active = style == "class:plan.in_progress"
+            blocks.append(Text(f"  {line}", style="bold" if active else "dim"))
+    for entry in worker.entries:
+        if entry.kind == "tool":
+            blocks.append(tool_line(entry.tool))
+        elif entry.kind == "thinking":
+            if show_thinking and entry.text.strip():
+                blocks += [Text(""), Text(entry.text.strip(), style="dim italic"), Text("")]
+        elif text := literal(entry.text).strip():
+            blocks += [Text(""), Markdown(text, code_theme=code_theme), Text("")]
+    if worker.running:
+        blocks += [Text(""), Text(f"  ({worker.state()}…)", style="dim")]
+    return blocks
+
+
+class WorkerBrowser:
+    """Follow delegated workers beside a conversation that may still be running."""
+
+    def __init__(
+        self,
+        workers: Workers,
+        *,
+        rich_theme: Theme | None = None,
+        code_theme: str = "ansi_dark",
+        color_system: str | None = "truecolor",
+        show_thinking: bool = False,
+        key_prefix: str | None = None,
+        **app_options,
+    ) -> None:
+        self.workers = workers
+        self.code_theme = code_theme
+        self.show_thinking = show_thinking
+        self.items = list(workers.items)
+        latest = workers.latest()
+        self.selected = latest.call_id if latest else None
+        self._rendered: tuple | None = None
+        self._refreshing = False
+        self.list = TextArea(read_only=True, wrap_lines=False, scrollbar=True)
+        self.list.window.cursorline = Always()
+        self.detail = RichPane(theme=rich_theme, color_system=color_system)
+        self.list.buffer.on_cursor_position_changed += lambda _: self.select()
+        keys = KeyBindings()
+        self.detail.bind_scrolling(keys)
+        bind_list_paging(keys, self.list, has_focus(self.list))
+
+        @keys.add("escape", eager=True)
+        @keys.add("c-c")
+        @keys.add("enter")
+        def close(event):
+            event.app.exit(result=None)
+
+        self.prefix_keys = shortcuts = PrefixKeys(key_prefix)
+        shortcuts.set_help(
+            lambda: [
+                ("↑/↓", "Select agent / scroll output"),
+                ("PgUp/PgDn", "Page"),
+                ("Ctrl+U/D", "Half page"),
+                ("Tab/Shift+Tab", "Change focus"),
+                ("Enter/Esc/Ctrl+C", "Close"),
+            ]
+        )
+
+        @shortcuts.add("t", "Thinking")
+        def toggle_thinking(event):
+            self.show_thinking = not self.show_thinking
+            self.select(force=True)
+
+        keys.add("tab")(focus_next)
+        keys.add("s-tab")(focus_previous)
+
+        header = Label(
+            lambda: (
+                f"Agents · {len(self.items)} this session · {self.workers.running()} running"
+                f" · read-only · thinking {'on' if self.show_thinking else 'off'}"
+            )
+        )
+        wide = VSplit(
+            [
+                Frame(self.list, title="Agents", width=Dimension(weight=2)),
+                Frame(self.detail, title="Output", width=Dimension(weight=3)),
+            ]
+        )
+        narrow = HSplit(
+            [
+                Frame(self.list, title="Agents", height=lambda: list_pane_height(len(self.items))),
+                Frame(self.detail, title="Output"),
+            ]
+        )
+        body = DynamicContainer(
+            lambda: wide if get_app().output.get_size().columns >= 100 else narrow
+        )
+        root_container = HSplit(
+            [
+                header,
+                body,
+                Label(shortcuts.summary),
+            ]
+        )
+        self.app = Application(
+            layout=Layout(popup_container(root_container, shortcuts), focused_element=self.list),
+            key_bindings=shortcuts.key_bindings(keys),
+            full_screen=True,
+            mouse_support=popup_mouse(shortcuts),
+            style=popup_style(app_options.pop("style", None)),
+            **app_options,
+        )
+        self.refresh()
+
+    def refresh(self, *, force: bool = True) -> None:
+        """Rebuild the list from the current records, keeping the selection."""
+        self.items = list(self.workers.items)
+        lines = [row(worker) for worker in self.items] or ["No agents yet"]
+        index = next(
+            (i for i, worker in enumerate(self.items) if worker.call_id == self.selected),
+            max(0, len(self.items) - 1),
+        )
+        position = sum(len(line) + 1 for line in lines[:index])
+        self._refreshing = True
+        self.list.buffer.set_document(Document("\n".join(lines), position), bypass_readonly=True)
+        self._refreshing = False
+        self.select(force=force)
+
+    def current(self) -> Worker | None:
+        return next((w for w in self.items if w.call_id == self.selected), None)
+
+    def select(self, force: bool = False) -> None:
+        if self._refreshing:
+            return
+        index = self.list.document.cursor_position_row
+        if index < len(self.items):
+            self.selected = self.items[index].call_id
+        worker = self.current()
+        # Elapsed time moves every second even when nothing else changes.
+        state = (
+            (worker.call_id, worker.version, int(worker.elapsed()), self.show_thinking)
+            if worker
+            else (None,)
+        )
+        if state == self._rendered and not force:
+            return
+        switched = self._rendered is None or state[0] != self._rendered[0]
+        self._rendered = state
+        blocks = details(worker, code_theme=self.code_theme, show_thinking=self.show_thinking)
+        if switched and worker is not None and worker.running:
+            # A running worker opens on its newest output, then follows it.
+            self.detail.set(blocks, anchor=len(blocks) - 1)
+        elif switched:
+            self.detail.set(blocks)
+        else:
+            self.detail.follow(blocks)
+
+    async def run(self) -> None:
+        async def follow():
+            while True:
+                await asyncio.sleep(REFRESH_SECONDS)
+                if [w.call_id for w in self.workers.items] != [w.call_id for w in self.items]:
+                    self.refresh()
+                elif any(w.running for w in self.items):
+                    # List rows carry elapsed time and state; keep them moving too.
+                    self.refresh(force=False)
+                else:
+                    self.select()
+                self.app.invalidate()
+
+        def start():
+            self.app.create_background_task(follow())
+
+        await self.app.run_async(pre_run=start)