pcode 0.1.0__py3-none-any.whl

pcode/file_refs.py

@@ -0,0 +1,265 @@
+"""Workspace file references: type `@` in the prompt to insert a relative path.
+
+Completing `@ui.py` inserts `./src/pcode/ui.py`, the same workspace-relative
+form the file tools accept, so the model can read or search it without guessing
+where the file lives. The path is all pcode sends; whether the file is worth
+reading is the model's call.
+"""
+
+import os
+import re
+import subprocess
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from pathlib import Path
+from time import monotonic
+
+from prompt_toolkit.completion import CompleteEvent, Completer, Completion
+from prompt_toolkit.document import Document
+from prompt_toolkit.lexers import Lexer
+
+# Directories that never hold source worth referencing, used only when the
+# workspace is not a Git checkout (Git supplies its own ignore rules).
+# Ripgrep skips the hidden ones on its own; the rest it would happily list.
+IGNORED_DIRECTORIES = frozenset(
+    {
+        ".git",
+        ".hg",
+        ".svn",
+        ".mypy_cache",
+        ".pytest_cache",
+        ".ruff_cache",
+        ".tox",
+        ".venv",
+        "__pycache__",
+        "build",
+        "dist",
+        "node_modules",
+        "venv",
+    }
+)
+MAX_FILES = 20000
+# Paths read from a listing before the menu's own cap applies; see _ripgrep_files.
+LISTING_CEILING = MAX_FILES * 10
+MAX_COMPLETIONS = 50
+CACHE_SECONDS = 10.0
+
+# A reference is a whitespace-delimited token: an unfinished `@fragment`, a
+# `./` or `../` path, or a quoted path whose name contains spaces.
+REFERENCE_PATTERN = re.compile(r'(?<!\S)(?:"\.{1,2}/[^"\n]+"|@\S*|\.{1,2}/\S+)')
+
+
+def _git_files(root: Path) -> list[str] | None:
+    """Tracked and untracked files below root, or None outside a Git checkout."""
+    try:
+        result = subprocess.run(
+            ["git", "-C", str(root), "ls-files", "--cached", "--others", "--exclude-standard"],
+            capture_output=True,
+            text=True,
+            timeout=5,
+        )
+    except (OSError, subprocess.SubprocessError):
+        return None
+    if result.returncode != 0:
+        return None
+    # A file staged for deletion is still cached; listing it would complete to a
+    # path that no longer exists, so drop duplicates and keep insertion order.
+    return list(dict.fromkeys(result.stdout.splitlines()))
+
+
+def _ripgrep_files(root: Path) -> list[str] | None:
+    """Files below root via ripgrep, or None when `rg` cannot be run.
+
+    Used where there is no Git listing to borrow. Ripgrep walks in parallel and
+    skips hidden entries itself, which beats `os.walk` several times over on a
+    large tree; completion runs on the UI event loop, so that time is a visible
+    stall between keystrokes.
+
+    Ripgrep's default `.gitignore` handling is left alone: outside a checkout it
+    reads those files only under `--no-require-git`, and tool-managed trees such
+    as a uv cache or a virtualenv carry a `.gitignore` of `*` that would leave
+    the menu empty. `.ignore` files still apply, as they do everywhere.
+
+    Output is read as it arrives and the process is stopped at LISTING_CEILING,
+    which bounds an unreasonable tree without truncating a merely large one:
+    ripgrep's parallel walk emits in a different order each run, so a cut below
+    the caller's own cap would reshuffle the menu between refreshes.
+    """
+    excluded = ",".join(sorted(IGNORED_DIRECTORIES))
+    arguments = ["rg", "--files", "--glob", f"!{{{excluded}}}"]
+    paths: list[str] = []
+    try:
+        with subprocess.Popen(
+            arguments, cwd=root, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True
+        ) as process:
+            for line in process.stdout:
+                paths.append(line.rstrip("\n"))
+                if len(paths) >= LISTING_CEILING:
+                    process.terminate()
+                    break
+            # A stopped run is killed mid-write, so only judge a complete one.
+            if process.wait() not in (0, 1) and len(paths) < LISTING_CEILING:
+                return None
+    except (OSError, subprocess.SubprocessError):
+        return None
+    # Ripgrep's parallel walk emits in thread completion order; sorting restores
+    # the stable menu `os.walk` gave, and costs far less than `--sort path`,
+    # which would make ripgrep search serially.
+    return sorted(paths)
+
+
+def _walked_files(root: Path) -> list[str]:
+    """Files below root in pure Python, for a machine without ripgrep."""
+    paths: list[str] = []
+    for directory, subdirectories, names in os.walk(root):
+        subdirectories[:] = sorted(
+            name
+            for name in subdirectories
+            if name not in IGNORED_DIRECTORIES and not name.startswith(".")
+        )
+        base = Path(directory)
+        for name in sorted(names):
+            if name.startswith("."):
+                continue
+            paths.append(str((base / name).relative_to(root)))
+            if len(paths) >= MAX_FILES:
+                return paths
+    return paths
+
+
+def human_size(size: int) -> str:
+    if size < 1024:
+        return f"{size} B"
+    if size < 1024 * 1024:
+        return f"{size / 1024:.1f} KB"
+    return f"{size / (1024 * 1024):.1f} MB"
+
+
+@dataclass
+class WorkspaceFiles:
+    """Workspace-relative file paths, refreshed at most every CACHE_SECONDS.
+
+    Completion runs on every keystroke, so the listing is cached rather than
+    re-read; new files show up within the refresh window.
+    """
+
+    root: Path
+    cache_seconds: float = CACHE_SECONDS
+    _paths: list[str] = field(default_factory=list)
+    _loaded_at: float | None = None
+
+    def paths(self) -> list[str]:
+        now = monotonic()
+        if self._loaded_at is None or now - self._loaded_at >= self.cache_seconds:
+            # An empty listing falls through as a missing one: a workspace that
+            # Git or ripgrep ignores wholesale (a gitignored directory, say) is
+            # still worth walking rather than offering nothing.
+            listed = _git_files(self.root) or _ripgrep_files(self.root)
+            self._paths = (listed or _walked_files(self.root))[:MAX_FILES]
+            self._loaded_at = now
+        return self._paths
+
+    def matches(self, fragment: str, limit: int = MAX_COMPLETIONS) -> list[str]:
+        """Paths containing fragment, closest match first.
+
+        A basename hit outranks a directory hit, and an earlier hit outranks a
+        later one, so `@ui.py` offers `src/pcode/ui.py` before `tests/test_ui.py`.
+        """
+        needle = fragment.lstrip("@").removeprefix("./").lower()
+        if not needle:
+            # Shallowest first: the interesting top-level files, not build output.
+            return sorted(self.paths(), key=lambda path: (path.count("/"), path))[:limit]
+        ranked = []
+        for path in self.paths():
+            position = path.lower().find(needle)
+            if position < 0:
+                continue
+            name_position = os.path.basename(path).lower().find(needle)
+            ranked.append(((0, name_position) if name_position >= 0 else (1, position), path))
+        ranked.sort(key=lambda item: (item[0], len(item[1]), item[1]))
+        return [path for _, path in ranked[:limit]]
+
+    def describe(self, path: str) -> str:
+        """The file's size, so a reference's cost to read is visible before picking."""
+        try:
+            return human_size((self.root / path).stat().st_size)
+        except OSError:
+            return "file"
+
+
+def reference_fragment(text: str) -> str | None:
+    """The path fragment of a trailing `@` reference, or None when absent.
+
+    `@` only starts a reference at a word boundary, so emails and decorators in
+    pasted text never open the menu.
+    """
+    token = text.rsplit(maxsplit=1)[-1] if text.split() else ""
+    if not text or text[-1].isspace() or not token.startswith("@"):
+        return None
+    return token[1:]
+
+
+class ReferenceLexer(Lexer):
+    """Style file references in the editor so they read as paths, not prose.
+
+    ``extra`` adds further ``(pattern, style)`` tokens, such as paste markers.
+    """
+
+    def __init__(self, extra: list[tuple[re.Pattern[str], str]] = ()) -> None:
+        self.patterns = [(REFERENCE_PATTERN, "class:reference"), *extra]
+
+    def lex_document(self, document: Document):
+        def line_fragments(number: int):
+            line = document.lines[number]
+            matches = sorted(
+                (match.start(), match.end(), style)
+                for pattern, style in self.patterns
+                for match in pattern.finditer(line)
+            )
+            fragments = []
+            position = 0
+            for start, end, style in matches:
+                if start < position:
+                    continue
+                fragments.append(("", line[position:start]))
+                fragments.append((style, line[start:end]))
+                position = end
+            fragments.append(("", line[position:]))
+            return fragments
+
+        return line_fragments
+
+
+class FileReferenceCompleter(Completer):
+    """Replace a trailing `@fragment` with a `./` workspace-relative path.
+
+    ``workspace`` may be a callable, read on every completion: a hosted session
+    learns its worktree from the host only after the prompt exists, and /switch
+    can move the terminal to a session in another checkout.
+    """
+
+    def __init__(self, workspace: Path | Callable[[], Path] | None = None) -> None:
+        self._workspace = workspace if callable(workspace) else (lambda: workspace)
+        self.files = WorkspaceFiles(self._root())
+
+    def _root(self) -> Path:
+        return (self._workspace() or Path.cwd()).resolve()
+
+    def get_completions(self, document: Document, complete_event: CompleteEvent):
+        fragment = reference_fragment(document.text_before_cursor)
+        if fragment is None:
+            return
+        if (root := self._root()) != self.files.root:
+            self.files = WorkspaceFiles(root)
+        for path in self.files.matches(fragment):
+            # The `@` is a trigger, not part of the reference: it is replaced so
+            # the model receives a path its file tools accept verbatim. The
+            # leading './' is what marks it as a path; only a name containing
+            # whitespace needs quotes to keep its end unambiguous in prose.
+            reference = f"./{path}"
+            yield Completion(
+                f'"{reference}" ' if any(char.isspace() for char in reference) else f"{reference} ",
+                start_position=-(len(fragment) + 1),
+                display=path,
+                display_meta=self.files.describe(path),
+            )

pcode/filesystem.py

@@ -0,0 +1,138 @@
+"""Mutation evidence layered on Harness's filesystem operations.
+
+Harness announces each write and edit through `_request` while it holds the
+file's lock and before it writes, so the before snapshot is taken there, inside
+the operation rather than from a later reread. The after side is what the write
+put there: the written content, or the same replacements applied to the snapshot.
+"""
+
+import posixpath
+from contextvars import ContextVar
+from dataclasses import dataclass, field
+
+from pydantic_ai import CapabilityEvent
+from pydantic_ai_harness.filesystem._toolset import (
+    FileSystemToolset,
+    _apply_replacements,
+    _is_binary,
+)
+
+from pcode.edits import MAX_SOURCE, completed_change, sensitive_path
+from pcode.runtime import EditCompleted
+from pcode.workspace_filesystem import WorkspaceFileSystem
+
+
+@dataclass(kw_only=True)
+class FileChangeEvent(CapabilityEvent, namespace="pcode_files", name="change"):
+    change: EditCompleted
+
+
+@dataclass
+class _Snapshot:
+    """What `_request` saw of one write or edit, for the event sent after it."""
+
+    display_path: str = ""
+    before: str | None = None
+    existed: bool = True
+    omitted: str = ""
+    taken: bool = field(default=False)
+
+
+# One per tool call: concurrent calls run in separate tasks with their own copy.
+_snapshot: ContextVar[_Snapshot | None] = ContextVar("pcode_file_snapshot", default=None)
+
+
+class DisplayFileSystem(WorkspaceFileSystem):
+    def _toolset_type(self):
+        return DisplayFileSystemToolset
+
+
+class DisplayFileSystemToolset(FileSystemToolset):
+    async def _emit_change(self, ctx, path, before, after, **kwargs):
+        change = completed_change(path, before, after, call_id=ctx.tool_call_id or "", **kwargs)
+        await ctx.emit(FileChangeEvent(change=change))
+
+    @staticmethod
+    def _display_path(scope, path: str, resolved: str) -> str:
+        """Workspace-relative inside the workspace, absolute outside it; never a sensitive name."""
+        if sensitive_path(path):
+            return path
+        relative = posixpath.relpath(resolved, scope.cwd)
+        return resolved if relative == ".." or relative.startswith("../") else relative
+
+    async def _request(self, scope, ctx, change, *, path, resolved):
+        snapshot = _snapshot.get()
+        if ctx is not None and snapshot is not None:
+            await self._take(snapshot, scope, path, resolved)
+        return await super()._request(scope, ctx, change, path=path, resolved=resolved)
+
+    async def _take(self, snapshot: _Snapshot, scope, path: str, resolved: str) -> None:
+        snapshot.taken = True
+        # Harness writes through the path as written; name the file it reaches,
+        # so an alias to a sensitive file is hidden as that file would be.
+        try:
+            real = await scope.workspace.realpath(resolved)
+        except OSError:
+            real = resolved
+        snapshot.display_path = (
+            path if sensitive_path(path) else self._display_path(scope, path, real)
+        )
+        if sensitive_path(snapshot.display_path):
+            return
+        try:
+            raw = await scope.workspace.read_bytes(real)
+        except FileNotFoundError:
+            snapshot.before, snapshot.existed = "", False
+            return
+        except PermissionError:
+            # Display evidence must not require permissions the write itself
+            # doesn't need; leave the diff out instead.
+            snapshot.omitted = "Before snapshot unavailable"
+            return
+        if len(raw) > MAX_SOURCE * 4:
+            snapshot.omitted = "File exceeds preview size limit"
+        elif _is_binary(raw):
+            snapshot.omitted = "Binary content"
+        else:
+            try:
+                snapshot.before = raw.decode("utf-8")
+            except UnicodeDecodeError:
+                snapshot.omitted = "Binary or non-UTF-8 content"
+
+    async def _recorded(self, ctx, operation, after) -> str:
+        """Run `operation`, then report the change it made, if it made one."""
+        snapshot = _Snapshot()
+        token = _snapshot.set(snapshot)
+        try:
+            result = await operation()
+        finally:
+            _snapshot.reset(token)
+        succeeded = isinstance(result, str) and result.startswith(("Wrote ", "Edited "))
+        if ctx is not None and snapshot.taken and succeeded:
+            await self._emit_change(
+                ctx,
+                snapshot.display_path,
+                snapshot.before,
+                after(snapshot.before) if snapshot.before is not None else "",
+                existed=snapshot.existed,
+                omitted=snapshot.omitted,
+            )
+        return result
+
+    async def _write_file(self, scope, ctx, path, content, *, expected_hash=None):
+        return await self._recorded(
+            ctx,
+            lambda: super(DisplayFileSystemToolset, self)._write_file(
+                scope, ctx, path, content, expected_hash=expected_hash
+            ),
+            lambda before: content,
+        )
+
+    async def _edit_file(self, scope, ctx, path, replacements, *, expected_hash=None):
+        return await self._recorded(
+            ctx,
+            lambda: super(DisplayFileSystemToolset, self)._edit_file(
+                scope, ctx, path, replacements, expected_hash=expected_hash
+            ),
+            lambda before: _apply_replacements(before, replacements, path),
+        )

pcode/frame.py

@@ -0,0 +1,149 @@
+"""pcode's framed boxes: a heading on the top rule and an optional footer below.
+
+prompt_toolkit's own ``Frame`` centers its title between bars
+(``┌──| Title |──┐``) and has no footer. These read ``┌─ Title ────┐``, the
+way the task panel above the prompt draws its heading, and can carry a short
+footer such as the keys that close an overlay: ``└──── Esc cancel ─┘``.
+"""
+
+from prompt_toolkit.filters import Condition, has_completions
+from prompt_toolkit.formatted_text import AnyFormattedText, to_formatted_text
+from prompt_toolkit.formatted_text.utils import fragment_list_to_text
+from prompt_toolkit.key_binding import KeyBindings
+from prompt_toolkit.key_binding.bindings.focus import focus_next, focus_previous
+from prompt_toolkit.layout import (
+    AnyContainer,
+    ConditionalContainer,
+    DynamicContainer,
+    HSplit,
+    VSplit,
+    Window,
+)
+from prompt_toolkit.layout.dimension import AnyDimension
+from prompt_toolkit.utils import get_cwidth
+from prompt_toolkit.widgets import Box, Label, Shadow
+from prompt_toolkit.widgets.base import Border
+
+# Columns a heading or footer takes beyond its text: the corners, the rule
+# stub beside the text, and a space either side. See ``Frame``.
+TITLE_CHROME = 5
+
+
+def text_width(text: AnyFormattedText) -> int:
+    return get_cwidth(fragment_list_to_text(to_formatted_text(text)))
+
+
+class Frame:
+    """A box around ``body``, drop-in for prompt_toolkit's ``Frame``.
+
+    ``title`` and ``footer`` may be callables, read each render; an empty one
+    leaves a plain rule.
+    """
+
+    def __init__(
+        self,
+        body: AnyContainer,
+        title: AnyFormattedText = "",
+        style: str = "",
+        width: AnyDimension = None,
+        height: AnyDimension = None,
+        key_bindings: KeyBindings | None = None,
+        modal: bool = False,
+        footer: AnyFormattedText = "",
+    ) -> None:
+        self.title = title
+        self.footer = footer
+        self.body = body
+
+        def border(char: str = Border.HORIZONTAL, width: int | None = None) -> Window:
+            return Window(char=char, width=width, height=1, style="class:frame.border")
+
+        def text(get, style: str) -> Label:
+            # Reads the attribute each render, so a caller may reassign it.
+            return Label(
+                lambda: [("", " "), *to_formatted_text(get()), ("", " ")],
+                style=style,
+                dont_extend_width=True,
+            )
+
+        def shown(get) -> Condition:
+            return Condition(lambda: bool(fragment_list_to_text(to_formatted_text(get()))))
+
+        top = ConditionalContainer(
+            VSplit(
+                [
+                    border(Border.TOP_LEFT, 1),
+                    border(width=1),
+                    text(lambda: self.title, "class:frame.label"),
+                    border(),
+                    border(Border.TOP_RIGHT, 1),
+                ],
+                height=1,
+            ),
+            filter=shown(lambda: self.title),
+            alternative_content=VSplit(
+                [border(Border.TOP_LEFT, 1), border(), border(Border.TOP_RIGHT, 1)], height=1
+            ),
+        )
+        bottom = ConditionalContainer(
+            VSplit(
+                [
+                    border(Border.BOTTOM_LEFT, 1),
+                    border(),
+                    text(lambda: self.footer, "class:frame.footer"),
+                    border(width=1),
+                    border(Border.BOTTOM_RIGHT, 1),
+                ],
+                height=1,
+            ),
+            filter=shown(lambda: self.footer),
+            alternative_content=VSplit(
+                [border(Border.BOTTOM_LEFT, 1), border(), border(Border.BOTTOM_RIGHT, 1)],
+                height=1,
+            ),
+        )
+        middle = VSplit(
+            [
+                Window(width=1, char=Border.VERTICAL, style="class:frame.border"),
+                DynamicContainer(lambda: self.body),
+                Window(width=1, char=Border.VERTICAL, style="class:frame.border"),
+            ],
+            padding=0,
+        )
+        self.container = HSplit(
+            [top, middle, bottom],
+            width=width,
+            height=height,
+            style="class:frame " + style,
+            key_bindings=key_bindings,
+            modal=modal,
+        )
+
+    def __pt_container__(self) -> AnyContainer:
+        return self.container
+
+
+class Dialog:
+    """A modal framed box over a background, as pcode's popups use it.
+
+    The subset of prompt_toolkit's ``Dialog`` used here (no buttons, always on
+    a background), drawn with this module's ``Frame``.
+    """
+
+    def __init__(self, body: AnyContainer, title: AnyFormattedText = "") -> None:
+        self.body = body
+        self.title = title
+        keys = KeyBindings()
+        keys.add("tab", filter=~has_completions)(focus_next)
+        keys.add("s-tab", filter=~has_completions)(focus_previous)
+        frame = Frame(
+            DynamicContainer(lambda: self.body),
+            title=lambda: self.title,
+            style="class:dialog.body",
+            key_bindings=keys,
+            modal=True,
+        )
+        self.container = Box(body=Shadow(body=frame), style="class:dialog")
+
+    def __pt_container__(self) -> AnyContainer:
+        return self.container