pcode 0.1.0__py3-none-any.whl

pcode/claude_sdk/session_pool.py

@@ -0,0 +1,251 @@
+"""Live CLI processes for one event loop, matched to requests by history."""
+
+import asyncio
+import logging
+import time
+import weakref
+from contextlib import suppress
+from dataclasses import dataclass
+
+from pcode.claude_sdk.errors import ClaudeProcessError, ClaudeStartError
+from pcode.claude_sdk.messages import _blocks, _result_ids, _tool_use_ids, replay
+from pcode.claude_sdk.resume import ForkPoint, ResumeIndex
+from pcode.claude_sdk.session import ClaudeSession, SessionConfig, _Diverged
+
+logger = logging.getLogger(__name__)
+
+# Each process holds about 110-135 MB of its own beyond the ~200 MB binary the
+# processes share (measured with vmmap), and forking a transcript is warm anyway,
+# so keeping one only saves the ~0.8 s start. Finished ones are capped and expire
+# soon; one is enough for the next turn to continue on. A parked one belongs to a
+# run still executing its tools (often a parent waiting on delegations), so it is
+# never capped, only expired, and later. Under memory pressure none are kept.
+# The first two are defaults for `claude_idle_processes` / `claude_idle_minutes`.
+MAX_IDLE_SESSIONS = 1
+IDLE_SECONDS = 10 * 60
+PARKED_SECONDS = 30 * 60
+# How often idle processes are checked against memory pressure, and what counts:
+# less than this share of physical memory available.
+PRESSURE_CHECK_SECONDS = 60.0
+LOW_MEMORY_FRACTION = 0.10
+
+
+@dataclass
+class Checkout:
+    session: ClaudeSession
+    route: str  # "live", "fork" or "fresh", for logs and tests
+
+
+class SessionPool:
+    """Live processes for one event loop, matched to requests by history."""
+
+    def __init__(self, index: ResumeIndex, factory=None) -> None:
+        self.index = index
+        self.factory = factory
+        self.sessions: list[ClaudeSession] = []
+        self._lock = asyncio.Lock()
+        self._closing: set[asyncio.Task] = set()
+        self._expiry: asyncio.TimerHandle | None = None
+
+    async def checkout(
+        self, config: SessionConfig, messages: list[dict], chain: list[str]
+    ) -> Checkout:
+        """A session that has been sent `messages` and is producing the next response."""
+        for attempt in (1, 2):
+            try:
+                return await self._checkout(config, messages, chain)
+            except _Diverged as error:
+                # A live process that had gone its own way is retired: fork now.
+                if error.live and attempt == 1:
+                    continue
+                # Internal only: the runtime retries this name as transient.
+                raise ClaudeProcessError(config.model, error.message) from error
+        raise AssertionError("unreachable")
+
+    async def _checkout(self, config, messages, chain) -> Checkout:
+        async with self._lock:
+            for session in sorted(self.sessions, key=lambda s: s.last_used, reverse=True):
+                if session.accepts(config, messages, chain):
+                    session.busy = True
+                    route, point, start = "live", None, len(session.chain)
+                    break
+            else:
+                session = None
+                point, start = self._fork_point(config, messages, chain)
+                route = "fork" if point else "fresh"
+        if session is None:
+            session = await self._open(config, point)
+            if session is None:  # the transcript to fork is gone
+                route, point, start = "fresh", None, 0
+                session = await self._open(config, None)
+        delta = messages[start:]
+        try:
+            if route == "live":
+                await session.send(delta[0])
+            else:
+                open_ids = _tool_use_ids(messages[start - 1]) if start else ()
+                single = len(delta) == 1 and delta[0]["role"] == "user"
+                if single and _result_ids(delta[0]) == set(open_ids):
+                    await session.start(_blocks(delta[0]))
+                else:
+                    await session.start(replay(delta, open_ids))
+        except _Diverged as error:
+            error.live = route == "live"
+            self.release(session, ok=False)
+            raise
+        except BaseException:
+            self.release(session, ok=False)
+            raise
+        session.chain = chain[: len(messages)]
+        logger.debug("claude request via %s session (%d new messages)", route, len(delta))
+        return Checkout(session, route)
+
+    def _fork_point(self, config, messages, chain) -> tuple[ForkPoint | None, int]:
+        for position in range(len(messages) - 2, -1, -1):
+            if messages[position]["role"] != "assistant":
+                continue
+            point = self.index.get(chain[position])
+            if point is not None and point.cwd == config.cwd:
+                return point, position + 1
+        return None, 0
+
+    async def _open(self, config: SessionConfig, point: ForkPoint | None) -> ClaudeSession | None:
+        session = ClaudeSession(config, point)
+        session.busy = True
+        session.retire = lambda: self._drop(session)
+        self.sessions.append(session)
+        try:
+            await session.connect(self.factory)
+        except ClaudeStartError:
+            self._drop(session)
+            if point is not None:
+                logger.debug("could not resume Claude session %s", point.session_id)
+                self.index.forget(point.session_id)
+                return None
+            raise
+        except BaseException:
+            self._drop(session)
+            raise
+        return session
+
+    def record(self, session: ClaudeSession, assistant_hash: str) -> None:
+        """Note the response just produced, so the history can continue or fork from it."""
+        session.chain.append(assistant_hash)
+        if session.cli_session_id and session.response_uuid:
+            point = ForkPoint(session.cli_session_id, session.response_uuid, session.config.cwd)
+            self.index.add(assistant_hash, point)
+
+    def release(self, session: ClaudeSession, *, ok: bool) -> None:
+        session.busy = False
+        session.last_used = time.monotonic()
+        if not ok:
+            session.dead = True
+        low = memory_low()
+        for stale in [s for s in self.sessions if not s.busy and (s.dead or low)]:
+            self._drop(stale)
+        finished = sorted(
+            (s for s in self.sessions if not s.busy and not s.open_tool_ids),
+            key=lambda s: s.last_used,
+        )
+        for stale in finished[: max(0, len(finished) - _idle_limit())]:
+            self._drop(stale)
+        self._schedule_expiry()
+
+    @staticmethod
+    def _deadline(session: ClaudeSession) -> float:
+        return session.last_used + (PARKED_SECONDS if session.open_tool_ids else _idle_seconds())
+
+    def _schedule_expiry(self) -> None:
+        if self._expiry is not None:
+            self._expiry.cancel()
+            self._expiry = None
+        idle = [s for s in self.sessions if not s.busy]
+        if idle:
+            delay = max(0.0, min(map(self._deadline, idle)) - time.monotonic())
+            delay = min(delay, PRESSURE_CHECK_SECONDS)
+            self._expiry = asyncio.get_running_loop().call_later(delay, self._expire)
+
+    def _expire(self) -> None:
+        self._expiry = None
+        now = time.monotonic()
+        low = memory_low()
+        for session in [
+            s for s in self.sessions if not s.busy and (low or self._deadline(s) <= now)
+        ]:
+            self._drop(session)
+        self._schedule_expiry()
+
+    def _drop(self, session: ClaudeSession) -> None:
+        with suppress(ValueError):
+            self.sessions.remove(session)
+        task = asyncio.get_running_loop().create_task(session.close())
+        self._closing.add(task)
+        task.add_done_callback(self._closing.discard)
+
+    async def aclose(self) -> None:
+        if self._expiry is not None:
+            self._expiry.cancel()
+            self._expiry = None
+        for session in list(self.sessions):
+            self._drop(session)
+        if self._closing:
+            await asyncio.gather(*self._closing, return_exceptions=True)
+
+
+def _preference(key: str) -> int | None:
+    """A saved whole-number preference, read on use so `/config set` applies at once."""
+    from pcode.preferences import load_preferences
+
+    value = load_preferences().get(key, "")
+    return int(value) if value.isdecimal() else None
+
+
+def memory_low() -> bool:
+    """Whether the machine is short of memory, so no idle process is worth keeping.
+
+    Dropping one costs the next request a warm fork, even mid-round: the tool
+    results then arrive structured at the fork point.
+    """
+    try:
+        import psutil
+
+        memory = psutil.virtual_memory()
+    except Exception:
+        return False
+    return memory.available < memory.total * LOW_MEMORY_FRACTION
+
+
+def _idle_limit() -> int:
+    value = _preference("claude_idle_processes")
+    return MAX_IDLE_SESSIONS if value is None else value
+
+
+def _idle_seconds() -> float:
+    return (_preference("claude_idle_minutes") or 0) * 60 or IDLE_SECONDS
+
+
+_pools: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, SessionPool]" = (
+    weakref.WeakKeyDictionary()
+)
+_index: ResumeIndex | None = None
+
+
+def pool() -> SessionPool:
+    """This event loop's pool; processes cannot be shared across loops."""
+    global _index
+    loop = asyncio.get_running_loop()
+    current = _pools.get(loop)
+    if current is None:
+        _index = _index or ResumeIndex()
+        current = _pools[loop] = SessionPool(_index)
+    return current
+
+
+async def shutdown() -> None:
+    """Stop this event loop's CLI processes, interrupting any mid-turn first.
+
+    Called as pcode exits; the SDK's own atexit hook only sends SIGTERM.
+    """
+    current = _pools.pop(asyncio.get_running_loop(), None)
+    if current is not None:
+        await current.aclose()

pcode/claude_sdk/workspace.py

@@ -0,0 +1,34 @@
+"""`ClaudeWorkspace`: the directory a `claude:` request's CLI runs in."""
+
+from dataclasses import replace
+from pathlib import Path
+
+from pydantic_ai import RunContext
+from pydantic_ai.capabilities import AbstractCapability
+from pydantic_ai.models import ModelRequestContext
+
+# Model setting carrying the workspace the CLI runs in (see `ClaudeWorkspace`).
+CWD_SETTING = "pcode_claude_cwd"
+
+
+class ClaudeWorkspace(AbstractCapability):
+    """Run `claude:` requests' CLI in the agent's workspace.
+
+    The CLI tells the model its working directory, and keeps transcripts per
+    directory; pcode's process directory is not the workspace in worktree mode.
+    A `fallback` one yields to any other: sub-agents get the parent's as a
+    fallback, and an isolated worker's own checkout must win whatever the order.
+    """
+
+    def __init__(self, workspace: Path, *, fallback: bool = False) -> None:
+        self.workspace = Path(workspace)
+        self.fallback = fallback
+
+    async def before_model_request(
+        self, ctx: RunContext, request_context: ModelRequestContext
+    ) -> ModelRequestContext:
+        settings = request_context.model_settings or {}
+        if request_context.model.system != "claude" or (self.fallback and CWD_SETTING in settings):
+            return request_context
+        settings = {**settings, CWD_SETTING: str(self.workspace)}
+        return replace(request_context, model_settings=settings)

pcode/cli.py

@@ -0,0 +1,73 @@
+"""Lightweight entry point: silence terminal echo before importing the frontend."""
+
+import sys
+from contextlib import ExitStack
+
+_startup: ExitStack | None = None
+
+
+def restore_stdin() -> None:
+    """Hand terminal settings back before the editor or a launch question takes over."""
+    if _startup is not None:
+        _startup.close()
+
+
+def ask(prompt: str) -> str:
+    restore_stdin()
+    return input(prompt)
+
+
+def _quiet_stdin(stack: ExitStack) -> None:
+    if not sys.stdin.isatty() or not sys.stdout.isatty():
+        return
+    try:
+        import termios
+    except ImportError:
+        return
+    try:
+        fd = sys.stdin.fileno()
+        original = termios.tcgetattr(fd)
+        quiet = termios.tcgetattr(fd)
+        # The kernel queues typeahead; no reader thread or input flushing is
+        # needed. Canonical mode goes too: canonical without echo is how
+        # terminals such as Ghostty recognize a password prompt. Signals stay.
+        quiet[3] &= ~(termios.ECHO | termios.ECHONL | termios.ICANON)
+        quiet[6][termios.VMIN] = 1
+        quiet[6][termios.VTIME] = 0
+        termios.tcsetattr(fd, termios.TCSANOW, quiet)
+    except (OSError, ValueError, termios.error):
+        return
+
+    def restore():
+        try:
+            termios.tcsetattr(fd, termios.TCSANOW, original)
+        except (OSError, termios.error):
+            pass
+
+    stack.callback(restore)
+
+
+def main() -> None:
+    global _startup
+    from pcode.completion import COMPLETE_COMMAND
+
+    if sys.argv[1:2] == [COMPLETE_COMMAND]:
+        # Shell completion calls back on every Tab: skip the frontend import.
+        from pcode.completion import print_candidates
+
+        sys.exit(print_candidates(sys.argv[2:]))
+    with ExitStack() as stack:
+        _startup = stack
+        _quiet_stdin(stack)
+        try:
+            from pcode.app import main as run
+
+            run()
+        finally:
+            _startup = None
+
+
+if __name__ == "__main__":
+    from pcode.cli import main
+
+    main()

pcode/clipboard.py

@@ -0,0 +1,76 @@
+"""Best-effort system clipboard for popups; never raises at the call site.
+
+A local helper (`pbcopy`, `wl-copy`, `xclip`) is tried first because it needs no
+cooperation from the terminal emulator. OSC 52 is the fallback, and goes first
+over ssh, where the helper would target the clipboard of the wrong machine.
+tmux only forwards OSC 52 with `set -g set-clipboard on`, so neither path alone
+is enough.
+"""
+
+import base64
+import os
+import shutil
+import subprocess
+
+_HELPERS = (
+    ("pbcopy",),
+    ("wl-copy",),
+    ("xclip", "-selection", "clipboard"),
+    ("xsel", "--clipboard", "--input"),
+    ("clip.exe",),
+)
+
+# Terminals drop oversized OSC 52 payloads, and a clipboard is not a transfer
+# channel: truncate loudly rather than copy something silently incomplete.
+LIMIT = 64 * 1024
+
+
+def osc52(text: str) -> str:
+    payload = base64.b64encode(text.encode("utf-8", "replace")).decode("ascii")
+    return f"\x1b]52;c;{payload}\x07"
+
+
+def _helper(text: str, output=None) -> bool:
+    for helper in _HELPERS:
+        if shutil.which(helper[0]) is None:
+            continue
+        try:
+            # The caller owns the alternate screen: a chatty helper must not
+            # write over it, so its streams are discarded rather than inherited.
+            subprocess.run(
+                helper,
+                input=text.encode("utf-8", "replace"),
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+                check=True,
+                timeout=5,
+            )
+            return True
+        except (OSError, subprocess.SubprocessError):
+            continue
+    return False
+
+
+def _terminal(text: str, output) -> bool:
+    if output is None:
+        return False
+    try:
+        output.write_raw(osc52(text))
+        output.flush()
+        return True
+    except Exception:
+        return False
+
+
+def copy(text: str, output=None) -> tuple[bool, bool]:
+    """Copy `text`, returning whether it worked and whether it was truncated.
+
+    `output` is a prompt_toolkit output to write OSC 52 to; without one, only a
+    local helper can copy.
+    """
+    truncated = len(text) > LIMIT
+    if truncated:
+        text = text[:LIMIT] + "\n[truncated by pcode at 64 KiB]"
+    remote = bool(os.environ.get("SSH_CONNECTION") or os.environ.get("SSH_TTY"))
+    order = (_terminal, _helper) if remote else (_helper, _terminal)
+    return any(attempt(text, output) for attempt in order), truncated

pcode/code_mode.py

@@ -0,0 +1,43 @@
+"""Opt-in sandboxed batching for the read-only tools, keeping mutations native.
+
+Code mode collapses many dependent tool calls into one `run_code` snippet, which
+pays off for investigation (fan out greps and reads, filter in Python). Tools
+whose terminal display is the point -- edits with their diffs, plan updates, the
+persistent shell, delegation -- stay native, so the transcript keeps showing what
+changed rather than an opaque snippet that changed it.
+"""
+
+from typing import TYPE_CHECKING
+
+from pcode.preferences import SETTINGS, load_preferences
+
+if TYPE_CHECKING:
+    from pydantic_ai_harness.code_mode import CodeMode
+
+# Read-only lookups worth batching. Shell tools are already kept native by
+# CodeMode itself; naming the sandboxed set explicitly also keeps MCP and future
+# tools native until they are reviewed for it.
+SANDBOXED_TOOLS = (
+    "read_file",
+    "list_files",
+    "grep",
+    "read_tool_result",
+    "web_search",
+    "get_page",
+)
+
+
+def code_mode_enabled() -> bool:
+    preferences = load_preferences()
+    return preferences.get("code_mode", SETTINGS["code_mode"].default) == "on"
+
+
+def create_code_mode() -> "CodeMode | None":
+    """Return the capability when enabled, or None to leave tool calling alone."""
+    if not code_mode_enabled():
+        return None
+    # Imported here so the display modules can share SANDBOXED_TOOLS without
+    # loading the Monty sandbox.
+    from pydantic_ai_harness.code_mode import CodeMode
+
+    return CodeMode(tools=list(SANDBOXED_TOOLS))

pcode/codex_login.py

@@ -0,0 +1,189 @@
+"""pcode's own OpenAI Codex sign-in: `/login openai-codex` and `/logout openai-codex`.
+
+The flow itself is Pydantic AI's `OpenAICodexOAuthFlow` (PKCE authorization
+code against the public Codex client, callback pinned by OpenAI's registration
+to `http://localhost:1455/auth/callback`). pcode owns the browser and the
+credential file, an owner-only JSON file in `config_dir()`, and hands that
+file to the provider as an `OpenAICodexCredentialSource`, so refreshed tokens
+are written back and a login survives restarts.
+
+The Codex CLI's `$CODEX_HOME/auth.json` remains a read-only fallback when pcode
+has no login of its own (see `credential_source`). pcode never writes that
+file: the CLI requires an `id_token` there that this flow does not keep, so a
+pcode-written `auth.json` would be unreadable by the CLI.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import tempfile
+import webbrowser
+from pathlib import Path
+
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.providers.openai_codex import (
+    OpenAICodexCredentials,
+    OpenAICodexCredentialSource,
+    OpenAICodexOAuthFlow,
+)
+
+from pcode.auth import LoginError
+from pcode.preferences import config_dir
+
+LOGIN_TIMEOUT_SECONDS = 300.0
+_MAX_FILE_BYTES = 1024 * 1024
+
+
+def credentials_path() -> Path:
+    override = os.environ.get("PCODE_CODEX_CREDENTIALS_FILE", "").strip()
+    if override:
+        return Path(override).expanduser()
+    return config_dir() / "codex-credentials.json"
+
+
+def have_credentials(path: Path | None = None) -> bool:
+    """Existence check only; the picker must never read a stored token."""
+    try:
+        return (path if path is not None else credentials_path()).is_file()
+    except OSError:
+        return False
+
+
+def _token(value: object) -> str | None:
+    if not isinstance(value, str) or not value or any(char.isspace() for char in value):
+        return None
+    return value
+
+
+def read_credentials(path: Path) -> OpenAICodexCredentials:
+    """Load pcode's own stored Codex credential, never the CLI's auth file."""
+    try:
+        with path.open("r", encoding="utf-8-sig") as source:
+            contents = source.read(_MAX_FILE_BYTES + 1)
+        if len(contents) > _MAX_FILE_BYTES:
+            raise ValueError
+        data = json.loads(contents)
+        entry = data.get("openai-codex") if isinstance(data, dict) else None
+        if not isinstance(entry, dict) or entry.get("type") != "oauth":
+            raise ValueError
+        access = _token(entry.get("access"))
+        refresh = _token(entry.get("refresh"))
+        account = _token(entry.get("account_id"))
+        if access is None or refresh is None or account is None:
+            raise ValueError
+        return OpenAICodexCredentials(
+            access_token=access, refresh_token=refresh, account_id=account
+        )
+    except Exception:
+        raise LoginError(
+            "No usable pcode OpenAI Codex login. Run /login openai-codex to sign in "
+            "with your ChatGPT account. Remove pcode's login with /logout openai-codex "
+            "to fall back to the CLI's credential."
+        ) from None
+
+
+def write_credentials(path: Path, credentials: OpenAICodexCredentials) -> None:
+    """Replace the credential file atomically, owner-readable only."""
+    path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
+    payload = {
+        "openai-codex": {
+            "type": "oauth",
+            "access": credentials.access_token,
+            "refresh": credentials.refresh_token,
+            "account_id": credentials.account_id,
+        }
+    }
+    name = None
+    try:
+        with tempfile.NamedTemporaryFile(mode="w", dir=path.parent, delete=False) as file:
+            name = file.name
+            # Restrict before any token bytes reach the filesystem.
+            os.chmod(file.fileno(), 0o600)
+            json.dump(payload, file, indent=2)
+            file.write("\n")
+        os.replace(name, path)
+        name = None
+    finally:
+        if name is not None:
+            Path(name).unlink(missing_ok=True)
+
+
+def delete_credentials(path: Path) -> bool:
+    try:
+        path.unlink()
+        return True
+    except FileNotFoundError:
+        return False
+    except OSError:
+        raise LoginError(f"Could not remove the stored OpenAI Codex login: {path}") from None
+
+
+class CredentialStore(OpenAICodexCredentialSource):
+    """pcode's credential file as the provider's storage, so refreshes persist."""
+
+    def __init__(self, path: Path | None = None) -> None:
+        self.path = path if path is not None else credentials_path()
+
+    async def load(self) -> OpenAICodexCredentials:
+        return await asyncio.to_thread(read_credentials, self.path)
+
+    async def save(self, credentials: OpenAICodexCredentials) -> None:
+        await asyncio.to_thread(write_credentials, self.path, credentials)
+
+
+def credential_source(path: Path | None = None) -> CredentialStore | None:
+    """pcode's own store when a login exists, else `None` for the CLI's `auth.json`."""
+    path = path if path is not None else credentials_path()
+    return CredentialStore(path) if have_credentials(path) else None
+
+
+def _open_browser(url: str) -> None:
+    try:
+        webbrowser.open(url)
+    except Exception:
+        # A headless machine is expected; the URL was already displayed.
+        pass
+
+
+async def login(
+    *,
+    path: Path | None = None,
+    notify=None,
+    open_browser: bool = True,
+    timeout: float = LOGIN_TIMEOUT_SECONDS,
+    flow: OpenAICodexOAuthFlow | None = None,
+) -> OpenAICodexCredentials:
+    """Run the browser sign-in and store the resulting credential."""
+    path = path if path is not None else credentials_path()
+    flow = flow if flow is not None else OpenAICodexOAuthFlow()
+    url = flow.authorization_url()
+    if notify is not None:
+        notify(url)
+    if open_browser:
+        # Opening a browser can block on some desktops; never stall the terminal.
+        await asyncio.to_thread(_open_browser, url)
+    try:
+        async with asyncio.timeout(timeout):
+            credentials = await flow.exchange_code_from_callback()
+    except TimeoutError:
+        raise LoginError(
+            "OpenAI Codex sign-in timed out. Run /login openai-codex to try again."
+        ) from None
+    except OSError:
+        # The callback port is pinned by OpenAI's client registration.
+        raise LoginError(
+            "Port 1455 is already in use, so the sign-in callback cannot start. "
+            "Close the other login (for example a running `codex login`) and try again."
+        ) from None
+    except UserError:
+        raise LoginError(
+            "OpenAI Codex authorization failed. Try /login openai-codex again."
+        ) from None
+    except Exception:
+        raise LoginError(
+            "OpenAI Codex sign-in failed. No credential details were logged."
+        ) from None
+    await asyncio.to_thread(write_credentials, path, credentials)
+    return credentials