From b0e2475f2867a05d448b2e70f4b4202263ae18a8 Mon Sep 17 00:00:00 2001 From: elkaix Date: Mon, 20 Jul 2026 06:19:22 -0400 Subject: [PATCH 1/5] refactor(tui): unify prompt completion context --- CHANGELOG.md | 1 + src/pythinker_code/ui/shell/prompt.py | 521 ++---------------- .../ui/shell/prompting/completion/__init__.py | 21 + .../ui/shell/prompting/completion/context.py | 278 ++++++++++ .../ui/shell/prompting/completion/slash.py | 350 ++++++++++++ tests/ui_and_conv/test_file_completer.py | 55 ++ tests/ui_and_conv/test_slash_completer.py | 48 ++ tests/ui_and_conv/test_slash_highlight.py | 8 + 8 files changed, 795 insertions(+), 487 deletions(-) create mode 100644 src/pythinker_code/ui/shell/prompting/completion/__init__.py create mode 100644 src/pythinker_code/ui/shell/prompting/completion/context.py create mode 100644 src/pythinker_code/ui/shell/prompting/completion/slash.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a92066a..49968ca1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ GitHub Releases page; `0.8.0` is the new starting line. ## Unreleased - Freeze the public shell prompt compatibility contract with constructor and rendering coverage. +- Unify slash and file-mention completion behind one canonical completion context, adding quoted `@"path with spaces"` file mentions. - Keep the interactive prompt scene within the terminal height with a priority-ordered row allocator, and shut prompt background tasks/processes down with an awaited lifecycle. - Add xAI Grok OAuth login (browser loopback and device-code). - Add GitHub Copilot device-code OAuth login for individual github.com accounts. diff --git a/src/pythinker_code/ui/shell/prompt.py b/src/pythinker_code/ui/shell/prompt.py index 4f8b24d1..e29bf164 100644 --- a/src/pythinker_code/ui/shell/prompt.py +++ b/src/pythinker_code/ui/shell/prompt.py @@ -22,7 +22,6 @@ from prompt_toolkit import PromptSession from prompt_toolkit.application import Application from prompt_toolkit.application.current import get_app_or_none -from prompt_toolkit.auto_suggest import AutoSuggest, Suggestion from prompt_toolkit.buffer import Buffer from prompt_toolkit.clipboard.pyperclip import PyperclipClipboard from prompt_toolkit.completion import ( @@ -57,7 +56,6 @@ from prompt_toolkit.layout.dimension import Dimension from prompt_toolkit.layout.margins import Margin from prompt_toolkit.layout.menus import CompletionsMenu -from prompt_toolkit.lexers import Lexer from prompt_toolkit.patch_stdout import patch_stdout from prompt_toolkit.utils import get_cwidth from pydantic import BaseModel, ValidationError @@ -88,6 +86,17 @@ PromptSceneBudget, allocate_prompt_scene_rows, ) +from pythinker_code.ui.shell.prompting.completion.context import ( + CompletionKind, + parse_completion_context, +) +from pythinker_code.ui.shell.prompting.completion.slash import ( + InputHighlightLexer, + SlashCommandAutoSuggest, + SlashCommandCompleter, + command_name_set, + discard_slash_command, +) from pythinker_code.ui.shell.prompting.lifecycle import PromptLifecycle from pythinker_code.ui.shell.prompting.state import ( BufferObserved, @@ -180,463 +189,8 @@ class CwdLostError(OSError): """Raised when the working directory no longer exists (e.g. external drive unplugged).""" -def _command_name_set(commands: Sequence[SlashCommand[Any]]) -> frozenset[str]: - """Lowercased names and aliases for slash highlighting and completion.""" - names: set[str] = set() - for cmd in commands: - names.add(cmd.name.lower()) - names.update(alias.lower() for alias in cmd.aliases) - return frozenset(names) - - -def _is_known_slash_command_prefix(name: str, known: frozenset[str]) -> bool: - """True when ``name`` is a registered command or a prefix of one (e.g. ``/skill:py``).""" - lower = name.lower() - if lower in known: - return True - return any(command_name.startswith(lower) for command_name in known) - - -def _fuzzy_subsequence(needle: str, haystack: str) -> bool: - """True when ``needle`` is an (ordered, gap-tolerant) subsequence of ``haystack``. - - Mirrors the selector/settings filters (``selector.py``, ``settings_list.py``); - used as the lowest-priority slash match so a misspelled distinctive word like - ``gurd`` still surfaces ``/skill:pythinker-guard`` when no prefix matches. - """ - pos = 0 - for ch in needle: - found = haystack.find(ch, pos) - if found < 0: - return False - pos = found + 1 - return True - - -def _slash_first_arg_context( - document: Document, - known_names: frozenset[str], - arg_suggestions: dict[str, tuple[str, ...]], -) -> tuple[str, str] | None: - """When the cursor trails the first argument of a known slash command, return (cmd, partial).""" - if document.text_after_cursor.strip(): - return None - line = document.current_line_before_cursor - if not line.startswith("/"): - return None - match = re.match(r"^/([A-Za-z0-9][A-Za-z0-9_:.-]*)(?:\s+(\S*))?$", line) - if match is None: - return None - command = match.group(1).lower() - if command not in known_names or command not in arg_suggestions: - return None - return command, match.group(2) or "" - - -def _slash_command_token_before_cursor(document: Document) -> str | None: - """Return the active slash-command token, or ``None`` when completion should stay hidden.""" - text = document.text_before_cursor - - if document.text_after_cursor.strip(): - return None - - last_space = text.rfind(" ") - token = text[last_space + 1 :] - prefix = text[: last_space + 1] if last_space != -1 else "" - - if prefix.strip() or not token.startswith("/"): - return None - return token - - -def _slash_suggest_token_before_cursor(document: Document) -> str | None: - """Return the active slash token for inline ghost-text suggestion. - - Unlike :func:`_slash_command_token_before_cursor` (which gates the dropdown - menu and only fires when the slash command starts the line), this accepts a - ``/name`` token anywhere on the current line, so mid-sentence references such - as ``use /desi`` still ghost-complete. Ghost text only renders at the end of - the buffer, so we still require nothing typed after the cursor. - """ - if document.text_after_cursor.strip(): - return None - line = document.current_line_before_cursor - last_space = line.rfind(" ") - token = line[last_space + 1 :] - if not token.startswith("/"): - return None - return token - - -def _discard_slash_command(buffer: Buffer) -> bool: - """Cancel slash completion and remove the in-progress root slash command.""" - document = buffer.document - token = _slash_command_token_before_cursor(document) - if token is None: - return False - - buffer.cancel_completion() - prefix = document.text_before_cursor[: -len(token)] - new_text = prefix + document.text_after_cursor - if not new_text.strip(): - buffer.set_document(Document(), bypass_readonly=True) - return True - - buffer.set_document( - Document(new_text, cursor_position=min(len(prefix), len(new_text))), - bypass_readonly=True, - ) - return True - - -# A "/name" token that starts the input or follows whitespace. The name charset -# matches registered command names and aliases (including "skill:x" / "flow:x"). -_SLASH_TOKEN_RE = re.compile(r"(? None: - self._known_names = known_names - self._agent_mode = agent_mode - self._arg_suggestions = arg_suggestions or _no_arg_suggestions - - @override - def lex_document(self, document: Document) -> Callable[[int], StyleAndTextTuples]: - known = self._known_names() - arg_suggestions = self._arg_suggestions() - agent_mode = self._agent_mode() - lines = document.lines - - def spans(line: str, lineno: int) -> list[tuple[int, int, str]]: - out: list[tuple[int, int, str]] = [] - # Leading "!" bash prefix: first line, agent mode, command after it. - if agent_mode and lineno == 0 and line.startswith("!") and line[1:].strip(): - out.append((0, 1, "class:bash-prefix")) - # Slash commands anywhere on the line (registered names and prefixes). - for match in _SLASH_TOKEN_RE.finditer(line): - name = match.group(1) - if not _is_known_slash_command_prefix(name, known): - continue - # Path-like tokens ("/clear/subdir") are not commands. - if match.end() < len(line) and line[match.end()] == "/": - continue - out.append((match.start(), match.end(), "class:slash-command")) - # First argument after a line-start slash command with known subcommands. - if line.startswith("/"): - arg_match = re.match(r"^/([A-Za-z0-9][A-Za-z0-9_:.-]*)(?:\s+(\S+))", line) - if arg_match is not None: - command = arg_match.group(1).lower() - partial = arg_match.group(2) - options = arg_suggestions.get(command) - if options and any(option.startswith(partial.lower()) for option in options): - arg_start = arg_match.start(2) - out.append((arg_start, arg_match.end(2), "class:slash-arg")) - # "@path" file mentions at a word boundary (agent mode only). - if agent_mode: - for match in _MENTION_TOKEN_RE.finditer(line): - start = match.start() - if start > 0: - prev = line[start - 1] - if prev.isalnum() or prev in _MENTION_TRIGGER_GUARDS: - continue - out.append((start, match.end(), "class:file-mention")) - out.sort(key=lambda span: span[0]) - return out - - def get_line(lineno: int) -> StyleAndTextTuples: - try: - line = lines[lineno] - except IndexError: - return [] - fragments: StyleAndTextTuples = [] - pos = 0 - for start, end, style in spans(line, lineno): - if start < pos: - continue # defensive: drop overlapping spans - if start > pos: - fragments.append(("", line[pos:start])) - fragments.append((style, line[start:end])) - pos = end - if pos < len(line): - fragments.append(("", line[pos:])) - return fragments - - return get_line - - -def _no_exact_suggestions() -> dict[str, str]: - return {} - - -def _no_arg_suggestions() -> dict[str, tuple[str, ...]]: - return {} - - -class SlashCommandAutoSuggest(AutoSuggest): - """Inline ghost-text completion for a partially typed slash command. - - While the user types a ``/name`` token -- at the start of the line *or* - mid-sentence (e.g. ``use /desi``) -- the remainder of the best (alphabetically - first) matching command renders as dim ghost text after the cursor; Tab - accepts it word-for-word. After a command that declares fixed subcommands - (e.g. ``/theme cur``), the first argument is ghost-completed too. The dropdown - menu stays line-start-only, so mid-sentence typing never pops a completion - list. Rendering and the standard accept bindings (right-arrow / ctrl-e) come - from prompt_toolkit's auto-suggest plumbing; the Tab binding is added in - CustomPromptSession. - """ - - def __init__( - self, - known_names: Callable[[], frozenset[str]], - *, - exact_suggestions: Callable[[], dict[str, str]] | None = None, - arg_suggestions: Callable[[], dict[str, tuple[str, ...]]] | None = None, - ) -> None: - self._known_names = known_names - self._exact_suggestions = exact_suggestions or _no_exact_suggestions - self._arg_suggestions = arg_suggestions or _no_arg_suggestions - - @override - def get_suggestion(self, buffer: Buffer, document: Document) -> Suggestion | None: - arg_ctx = _slash_first_arg_context(document, self._known_names(), self._arg_suggestions()) - if arg_ctx is not None: - command, partial = arg_ctx - options = self._arg_suggestions()[command] - partial_lower = partial.lower() - if not partial: - return Suggestion(options[0]) - matches = [ - option - for option in options - if option.startswith(partial_lower) and len(option) > len(partial) - ] - if matches: - return Suggestion(matches[0][len(partial) :]) - return None - - token = _slash_suggest_token_before_cursor(document) - if token is None or len(token) < 2: - return None - typed = token[1:] - typed_lower = typed.lower() - exact = self._exact_suggestions().get(typed_lower) - if exact is not None and typed_lower in self._known_names(): - return Suggestion(exact) - matches = sorted( - name - for name in self._known_names() - if name.lower().startswith(typed_lower) and len(name) > len(typed) - ) - if not matches: - return None - return Suggestion(matches[0][len(typed) :]) - - -class SlashCommandCompleter(Completer): - """ - A completer that: - - Shows one line per slash command using the canonical "/name" - - Matches exact names first, then name/alias prefixes, while inserting the canonical "/name" - - Only activates when the current token starts with '/' - """ - - def __init__( - self, - available_commands: Sequence[SlashCommand[Any]], - *, - annotate_meta: bool = False, - command_scope: str = "command", - is_task_running: Callable[[], bool] | None = None, - arg_suggestions: Callable[[], dict[str, tuple[str, ...]]] | None = None, - ) -> None: - super().__init__() - self._available_commands = sorted(available_commands, key=lambda c: c.name) - self._command_names = _command_name_set(available_commands) - self._annotate_meta = annotate_meta - self._command_scope = command_scope - self._is_task_running = is_task_running - self._arg_suggestions = arg_suggestions or _no_arg_suggestions - - def completion_active(self, document: Document) -> bool: - """Return whether slash command or subcommand completion should be active.""" - if _slash_command_token_before_cursor(document) is not None: - return True - return ( - _slash_first_arg_context(document, self._command_names, self._arg_suggestions()) - is not None - ) - - @staticmethod - def should_complete(document: Document) -> bool: - """Return whether slash command completion should be active for the current buffer.""" - return _slash_command_token_before_cursor(document) is not None - - @override - def get_completions( - self, document: Document, complete_event: CompleteEvent - ) -> Iterable[Completion]: - if not self.completion_active(document): - return - - arg_ctx = _slash_first_arg_context(document, self._command_names, self._arg_suggestions()) - if arg_ctx is not None: - _, partial = arg_ctx - partial_lower = partial.lower() - for option in self._arg_suggestions()[arg_ctx[0]]: - if partial and not option.startswith(partial_lower): - continue - yield Completion( - text=option[len(partial) :] if partial else option, - start_position=-len(partial), - display=option, - ) - return - - token = _slash_command_token_before_cursor(document) - if token is None: - return - - typed = token[1:] - typed_lower = typed.lower() - seen: set[str] = set() - - def emit(cmd: SlashCommand[Any], label: str | None = None) -> Iterable[Completion]: - if cmd.name in seen: - return - seen.add(cmd.name) - shown = label or cmd.name - yield Completion( - text=f"/{shown}", - start_position=-len(token), - display=f"/{shown}", - display_meta=self._display_meta(cmd), - ) - - if not typed: - for cmd in self._available_commands: - yield from emit(cmd) - return - - def match_tier(cmd: SlashCommand[Any]) -> tuple[int, str] | None: - """Return ``(tier, label)`` or ``None``. Lower tier = stronger match. - Name matches rank above alias matches so typing toward a command name - (e.g. ``/report`` → ``/reports``) wins over a command that only matches - via an exact alias. For alias-only matches the label is the matched - alias, so the menu surfaces what the user typed toward (e.g. ``/res`` - → ``/resume``) rather than the differently-named command (``/sessions``).""" - name_lower = cmd.name.lower() - if name_lower == typed_lower: - return (0, cmd.name) - if name_lower.startswith(typed_lower): - return (1, cmd.name) - alias_prefix: str | None = None - for alias in cmd.aliases: - alias_lower = alias.lower() - if alias_lower == typed_lower: - return (2, alias) - if alias_prefix is None and alias_lower.startswith(typed_lower): - alias_prefix = alias - if alias_prefix is not None: - return (3, alias_prefix) - # Namespaced commands ("skill:designer-skill", "flow:build-api") also - # match on their bare segment after the prefix, so `/designer` or - # `/build` surfaces them. The label stays the canonical name so the - # accepted completion inserts "/skill:designer-skill", not the bare - # term -- one execution path, no duplicate command. - segment = name_lower.split(":", 1)[1] if ":" in name_lower else name_lower - if ":" in name_lower: - if segment == typed_lower: - return (4, cmd.name) - if segment.startswith(typed_lower): - return (5, cmd.name) - # Last resort: fuzzy subsequence on the bare segment, so the - # distinctive word -- even misspelled (``gurd`` -> ``guard``) -- - # surfaces a command whose shared prefix (``pythinker-``) makes - # plain prefix matching useless. Gated at 2+ chars to avoid a - # single keystroke matching nearly everything. Label is canonical. - if len(typed_lower) >= 2 and _fuzzy_subsequence(typed_lower, segment): - return (6, cmd.name) - return None - - # Rank by (match tier, command-name length, name): the closest, shortest - # command name surfaces first within each tier. - matched: list[tuple[int, int, str, str, SlashCommand[Any]]] = [] - for cmd in self._available_commands: - result = match_tier(cmd) - if result is not None: - tier, label = result - matched.append((tier, len(cmd.name), cmd.name, label, cmd)) - matched.sort(key=lambda item: (item[0], item[1], item[2])) - if matched and matched[0][0] < 6: - matched = [item for item in matched if item[0] < 6] - - for _, _, _, label, cmd in matched: - yield from emit(cmd, label) - - def _disabled_during_task(self, cmd: SlashCommand[Any]) -> bool: - """True when a running turn blocks this shell-level command.""" - if self._is_task_running is None or not self._is_task_running(): - return False - from pythinker_code.ui.shell.slash import registry as shell_registry - - shell_cmd = shell_registry.find_command(cmd.name) - return shell_cmd is not None and not shell_cmd.available_during_task - - def _display_meta(self, cmd: SlashCommand[Any]) -> str: - if self._disabled_during_task(cmd): - return "disabled while a task is in progress" - if not self._annotate_meta: - return cmd.description - - # Only surface a kind tag when it distinguishes the entry from a plain - # command. Skills and flows are interleaved with commands in the agent - # menu, so their tag carries information; the generic command/shell scope - # is already obvious from the menu itself, so tagging every row is noise. - if cmd.name.startswith("skill:"): - kind: str | None = "skill" - elif cmd.name.startswith("flow:"): - kind = "flow" - else: - kind = None - - parts: list[str] = [] - if kind is not None: - parts.append(f"[{kind}]") - parts.append(cmd.description) - if cmd.aliases: - parts.append(f"aliases: {', '.join('/' + alias for alias in cmd.aliases)}") - return " ".join(part for part in parts if part) +_command_name_set = command_name_set +_discard_slash_command = discard_slash_command def _card_side_padding() -> int: @@ -1240,10 +794,10 @@ def _match_prefix_len(self, app: Any) -> int: document = getattr(getattr(app, "current_buffer", None), "document", None) if not isinstance(document, Document): return 0 - token = _slash_command_token_before_cursor(document) - if token is None: + context = parse_completion_context(document, allow_file=False) + if context.kind is not CompletionKind.SLASH_COMMAND: return 0 - return len(token[1:]) + return len(context.token[1:]) def _selected_meta_lines(self, text: str, meta_width: int) -> list[str]: lines = _wrap_to_width( @@ -1618,7 +1172,6 @@ class LocalFileMentionCompleter(Completer): """ _FRAGMENT_PATTERN = re.compile(r"[^\s@]+") - _TRIGGER_GUARDS = _MENTION_TRIGGER_GUARDS def __init__( self, @@ -1723,30 +1276,11 @@ def _get_deep_paths(self) -> list[str]: self._cache_time = now return self._cached_paths - @staticmethod - def _extract_fragment(text: str) -> str | None: - index = text.rfind("@") - if index == -1: - return None - - if index > 0: - prev = text[index - 1] - if prev.isalnum() or prev in LocalFileMentionCompleter._TRIGGER_GUARDS: - return None - - fragment = text[index + 1 :] - if not fragment: - return "" - - if any(ch.isspace() for ch in fragment): - return None - - return fragment - @staticmethod def should_complete(document: Document) -> bool: """Return whether `@` file completion should be active for the buffer.""" - return LocalFileMentionCompleter._extract_fragment(document.text_before_cursor) is not None + context = parse_completion_context(document, allow_slash=False) + return context.kind is CompletionKind.FILE def _is_completed_file(self, fragment: str) -> bool: candidate = fragment.rstrip("/") @@ -1761,9 +1295,10 @@ def _is_completed_file(self, fragment: str) -> bool: def get_completions( self, document: Document, complete_event: CompleteEvent ) -> Iterable[Completion]: - fragment = self._extract_fragment(document.text_before_cursor) - if fragment is None: + context = parse_completion_context(document, allow_slash=False) + if context.kind is not CompletionKind.FILE: return + fragment = context.token if self._is_completed_file(fragment): return @@ -1790,7 +1325,19 @@ def _rank(c: Completion) -> tuple[int, ...]: return (cat, test_penalty) candidates.sort(key=_rank) - yield from candidates + if not context.quoted: + yield from candidates + return + for candidate in candidates: + escaped = candidate.text.replace("\\", "\\\\").replace('"', '\\"') + yield Completion( + text=f'"{escaped}"', + start_position=context.start_position, + display=candidate.display, + display_meta=candidate.display_meta, + style=candidate.style, + selected_style=candidate.selected_style, + ) finally: self._fragment_hint = None diff --git a/src/pythinker_code/ui/shell/prompting/completion/__init__.py b/src/pythinker_code/ui/shell/prompting/completion/__init__.py new file mode 100644 index 00000000..52c46e1d --- /dev/null +++ b/src/pythinker_code/ui/shell/prompting/completion/__init__.py @@ -0,0 +1,21 @@ +"""Internal prompt completion primitives.""" + +from pythinker_code.ui.shell.prompting.completion.context import ( + CompletionContext, + CompletionKind, + parse_completion_context, +) +from pythinker_code.ui.shell.prompting.completion.slash import ( + InputHighlightLexer, + SlashCommandAutoSuggest, + SlashCommandCompleter, +) + +__all__ = ( + "CompletionContext", + "CompletionKind", + "InputHighlightLexer", + "SlashCommandAutoSuggest", + "SlashCommandCompleter", + "parse_completion_context", +) diff --git a/src/pythinker_code/ui/shell/prompting/completion/context.py b/src/pythinker_code/ui/shell/prompting/completion/context.py new file mode 100644 index 00000000..59b754ce --- /dev/null +++ b/src/pythinker_code/ui/shell/prompting/completion/context.py @@ -0,0 +1,278 @@ +"""Canonical parsing for prompt completion tokens.""" + +from __future__ import annotations + +import re +from collections.abc import Iterator, Mapping, Set +from dataclasses import dataclass +from enum import StrEnum +from typing import Literal + +from prompt_toolkit.document import Document + + +class CompletionKind(StrEnum): + """The completion grammar selected at the cursor.""" + + SLASH_COMMAND = "slash_command" + SLASH_ARGUMENT = "slash_argument" + FILE = "file" + NONE = "none" + + +@dataclass(frozen=True, slots=True) +class CompletionContext: + """A parsed completion token and its replacement coordinates.""" + + kind: CompletionKind + token: str + start_position: int + command: str | None + argument_index: int | None + quoted: bool = False + + +_NONE = CompletionContext(CompletionKind.NONE, "", 0, None, None) +_COMMAND_CHARS = frozenset("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_:.-") +_MENTION_TRIGGER_GUARDS = frozenset((".", "-", "_", "`", "'", '"', ":", "@", "#", "~")) +_SLASH_TOKEN_RE = re.compile(r"(? bool: + if index == 0: + return True + previous = text[index - 1] + return not (previous.isalnum() or previous in _MENTION_TRIGGER_GUARDS) + + +def _unescape_quoted(value: str) -> str: + out: list[str] = [] + escaped = False + for character in value: + if escaped: + if character in {'"', "\\"}: + out.append(character) + else: + out.extend(("\\", character)) + escaped = False + elif character == "\\": + escaped = True + else: + out.append(character) + if escaped: + out.append("\\") + return "".join(out) + + +def _file_context(text: str) -> CompletionContext: + index = text.rfind("@") + if index < 0 or not _mention_boundary(text, index): + return _NONE + + raw = text[index + 1 :] + if raw.startswith('"'): + content = raw[1:] + escaped = False + closing_index: int | None = None + for position, character in enumerate(content): + if escaped: + escaped = False + elif character == "\\": + escaped = True + elif character == '"': + closing_index = position + break + if closing_index is not None and content[closing_index + 1 :]: + return _NONE + raw_token = content if closing_index is None else content[:closing_index] + return CompletionContext( + CompletionKind.FILE, + _unescape_quoted(raw_token), + -(len(raw_token) + 1 + int(closing_index is not None)), + None, + None, + quoted=True, + ) + + if any(character.isspace() for character in raw) or "@" in raw: + return _NONE + return CompletionContext(CompletionKind.FILE, raw, -len(raw), None, None) + + +def _slash_context( + document: Document, + *, + known_commands: Set[str], + argument_commands: Set[str], + slash_activation: Literal["root", "any"], +) -> CompletionContext: + if document.text_after_cursor.strip(): + return _NONE + + line = document.current_line_before_cursor + if slash_activation == "any": + token_start = len(line) + while token_start > 0 and not line[token_start - 1].isspace(): + token_start -= 1 + candidate = line[token_start:] + if candidate.startswith("/") and all( + character in _COMMAND_CHARS for character in candidate[1:] + ): + return CompletionContext( + CompletionKind.SLASH_COMMAND, + candidate, + -len(candidate), + None, + None, + ) + + if line.startswith("/"): + command_end = 1 + while command_end < len(line) and line[command_end] in _COMMAND_CHARS: + command_end += 1 + if command_end > 1 and command_end < len(line) and line[command_end].isspace(): + command = line[1:command_end].lower() + if command in known_commands: + remainder = line[command_end:] + leading = len(remainder) - len(remainder.lstrip()) + arguments = remainder[leading:] + if not arguments: + argument_index = 0 + token = "" + else: + parts = arguments.split() + trailing_space = arguments[-1].isspace() + argument_index = len(parts) if trailing_space else len(parts) - 1 + token = "" if trailing_space else parts[-1] + return CompletionContext( + CompletionKind.SLASH_ARGUMENT, + token, + -len(token), + command, + argument_index, + ) + + if slash_activation == "root": + stripped = line.lstrip(" ") + if len(stripped) != len(line) and not line[: -len(stripped)].isspace(): + return _NONE + token_start = len(line) - len(stripped) + else: + return _NONE + + candidate = line[token_start:] + if candidate.startswith("/") and all( + character in _COMMAND_CHARS for character in candidate[1:] + ): + if slash_activation == "root" and line[:token_start].strip(): + return _NONE + return CompletionContext( + CompletionKind.SLASH_COMMAND, + candidate, + -len(candidate), + None, + None, + ) + + return _NONE + + +def parse_completion_context( + document: Document, + *, + known_commands: Set[str] = frozenset(), + argument_commands: Set[str] = frozenset(), + slash_activation: Literal["root", "any"] = "root", + allow_slash: bool = True, + allow_file: bool = True, +) -> CompletionContext: + """Parse the active slash or file token at the cursor exactly once. + + ``slash_activation`` preserves the two user-interface facets: ``root`` gates + dropdown completion to a leading slash command, while ``any`` recognizes the + final slash token for inline ghost suggestions and highlighting. + """ + if allow_slash: + slash = _slash_context( + document, + known_commands=known_commands, + argument_commands=argument_commands, + slash_activation=slash_activation, + ) + if slash.kind is CompletionKind.SLASH_COMMAND: + return slash + if ( + slash.kind is CompletionKind.SLASH_ARGUMENT + and slash.command in argument_commands + and slash.argument_index == 0 + ): + return slash + if slash.kind is CompletionKind.SLASH_ARGUMENT: + return slash + if allow_file: + return _file_context(document.text_before_cursor) + return _NONE + + +def iter_completion_contexts( + document: Document, + *, + known_commands: Set[str], + argument_suggestions: Mapping[str, tuple[str, ...]], + include_files: bool, +) -> Iterator[tuple[int, int, CompletionContext]]: + """Yield parser-produced contexts for syntax highlighting.""" + for line_number, line in enumerate(document.lines): + for match in _SLASH_TOKEN_RE.finditer(line): + if match.end() < len(line) and line[match.end()] == "/": + continue + parsed = parse_completion_context( + Document(line[: match.end()], cursor_position=match.end()), + known_commands=known_commands, + argument_commands=argument_suggestions.keys(), + slash_activation="any", + allow_file=False, + ) + if parsed.kind is CompletionKind.SLASH_COMMAND: + yield line_number, match.end(), parsed + + argument = _ARGUMENT_RE.match(line) + if argument is not None: + parsed = parse_completion_context( + Document(line[: argument.end(2)], cursor_position=argument.end(2)), + known_commands=known_commands, + argument_commands=argument_suggestions.keys(), + allow_file=False, + ) + if parsed.kind is CompletionKind.SLASH_ARGUMENT: + yield line_number, argument.end(2), parsed + + if not include_files: + continue + for index, character in enumerate(line): + if character != "@" or not _mention_boundary(line, index): + continue + end = index + 1 + quoted = end < len(line) and line[end] == '"' + if quoted: + end += 1 + escaped = False + while end < len(line): + current = line[end] + end += 1 + if escaped: + escaped = False + elif current == "\\": + escaped = True + elif current == '"': + break + else: + while end < len(line) and not line[end].isspace() and line[end] != "@": + end += 1 + parsed = parse_completion_context( + Document(line[:end], cursor_position=end), + allow_slash=False, + ) + if parsed.kind is CompletionKind.FILE: + yield line_number, end, parsed diff --git a/src/pythinker_code/ui/shell/prompting/completion/slash.py b/src/pythinker_code/ui/shell/prompting/completion/slash.py new file mode 100644 index 00000000..26dba629 --- /dev/null +++ b/src/pythinker_code/ui/shell/prompting/completion/slash.py @@ -0,0 +1,350 @@ +"""Slash completion, suggestion, and input highlighting.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable, Sequence +from typing import Any, override + +from prompt_toolkit.auto_suggest import AutoSuggest, Suggestion +from prompt_toolkit.buffer import Buffer +from prompt_toolkit.completion import CompleteEvent, Completer, Completion +from prompt_toolkit.document import Document +from prompt_toolkit.formatted_text import StyleAndTextTuples +from prompt_toolkit.lexers import Lexer + +from pythinker_code.ui.shell.prompting.completion.context import ( + CompletionKind, + iter_completion_contexts, + parse_completion_context, +) +from pythinker_code.utils.slashcmd import SlashCommand + + +def command_name_set(commands: Sequence[SlashCommand[Any]]) -> frozenset[str]: + """Return lower-cased command names and aliases.""" + names: set[str] = set() + for command in commands: + names.add(command.name.lower()) + names.update(alias.lower() for alias in command.aliases) + return frozenset(names) + + +def _is_known_slash_command_prefix(name: str, known: frozenset[str]) -> bool: + lower = name.lower() + return lower in known or any(command_name.startswith(lower) for command_name in known) + + +def _fuzzy_subsequence(needle: str, haystack: str) -> bool: + position = 0 + for character in needle: + found = haystack.find(character, position) + if found < 0: + return False + position = found + 1 + return True + + +def _no_exact_suggestions() -> dict[str, str]: + return {} + + +def _no_arg_suggestions() -> dict[str, tuple[str, ...]]: + return {} + + +def discard_slash_command(buffer: Buffer) -> bool: + """Cancel slash completion and remove the in-progress root slash command.""" + document = buffer.document + context = parse_completion_context(document, allow_file=False) + if context.kind is not CompletionKind.SLASH_COMMAND: + return False + + buffer.cancel_completion() + prefix = document.text_before_cursor[: context.start_position] + new_text = prefix + document.text_after_cursor + if not new_text.strip(): + buffer.set_document(Document(), bypass_readonly=True) + return True + + buffer.set_document( + Document(new_text, cursor_position=min(len(prefix), len(new_text))), + bypass_readonly=True, + ) + return True + + +class InputHighlightLexer(Lexer): + """Highlight slash commands, slash arguments, file mentions, and ``!``.""" + + def __init__( + self, + known_names: Callable[[], frozenset[str]], + *, + agent_mode: Callable[[], bool], + arg_suggestions: Callable[[], dict[str, tuple[str, ...]]] | None = None, + ) -> None: + self._known_names = known_names + self._agent_mode = agent_mode + self._arg_suggestions = arg_suggestions or _no_arg_suggestions + + @override + def lex_document(self, document: Document) -> Callable[[int], StyleAndTextTuples]: + known = self._known_names() + arguments = self._arg_suggestions() + agent_mode = self._agent_mode() + lines = document.lines + parsed_by_line: dict[int, list[tuple[int, int, str]]] = {} + + for line_number, end, context in iter_completion_contexts( + document, + known_commands=known, + argument_suggestions=arguments, + include_files=agent_mode, + ): + start = end + context.start_position + if context.kind is CompletionKind.SLASH_COMMAND: + name = context.token[1:] + if not _is_known_slash_command_prefix(name, known): + continue + style = "class:slash-command" + elif context.kind is CompletionKind.SLASH_ARGUMENT: + options = arguments.get(context.command or "") + if not options or not any( + option.startswith(context.token.lower()) for option in options + ): + continue + style = "class:slash-arg" + elif context.kind is CompletionKind.FILE: + start -= 1 + style = "class:file-mention" + else: + continue + parsed_by_line.setdefault(line_number, []).append((start, end, style)) + + def get_line(line_number: int) -> StyleAndTextTuples: + try: + line = lines[line_number] + except IndexError: + return [] + spans = parsed_by_line.get(line_number, []) + if agent_mode and line_number == 0 and line.startswith("!") and line[1:].strip(): + spans.append((0, 1, "class:bash-prefix")) + spans.sort(key=lambda span: span[0]) + fragments: StyleAndTextTuples = [] + position = 0 + for start, end, style in spans: + if start < position: + continue + if start > position: + fragments.append(("", line[position:start])) + fragments.append((style, line[start:end])) + position = end + if position < len(line): + fragments.append(("", line[position:])) + return fragments + + return get_line + + +class SlashCommandAutoSuggest(AutoSuggest): + """Inline ghost-text completion for slash commands and first arguments.""" + + def __init__( + self, + known_names: Callable[[], frozenset[str]], + *, + exact_suggestions: Callable[[], dict[str, str]] | None = None, + arg_suggestions: Callable[[], dict[str, tuple[str, ...]]] | None = None, + ) -> None: + self._known_names = known_names + self._exact_suggestions = exact_suggestions or _no_exact_suggestions + self._arg_suggestions = arg_suggestions or _no_arg_suggestions + + @override + def get_suggestion(self, buffer: Buffer, document: Document) -> Suggestion | None: + known = self._known_names() + arguments = self._arg_suggestions() + context = parse_completion_context( + document, + known_commands=known, + argument_commands=arguments.keys(), + slash_activation="any", + allow_file=False, + ) + if context.kind is CompletionKind.SLASH_ARGUMENT: + if context.argument_index != 0 or context.command not in arguments: + return None + partial = context.token + options = arguments[context.command] + partial_lower = partial.lower() + if not partial: + return Suggestion(options[0]) + matches = [ + option + for option in options + if option.startswith(partial_lower) and len(option) > len(partial) + ] + return Suggestion(matches[0][len(partial) :]) if matches else None + + if context.kind is not CompletionKind.SLASH_COMMAND or len(context.token) < 2: + return None + typed = context.token[1:] + typed_lower = typed.lower() + exact = self._exact_suggestions().get(typed_lower) + if exact is not None and typed_lower in known: + return Suggestion(exact) + matches = sorted( + name + for name in known + if name.lower().startswith(typed_lower) and len(name) > len(typed) + ) + return Suggestion(matches[0][len(typed) :]) if matches else None + + +class SlashCommandCompleter(Completer): + """Complete and rank root slash commands and fixed first arguments.""" + + def __init__( + self, + available_commands: Sequence[SlashCommand[Any]], + *, + annotate_meta: bool = False, + command_scope: str = "command", + is_task_running: Callable[[], bool] | None = None, + arg_suggestions: Callable[[], dict[str, tuple[str, ...]]] | None = None, + ) -> None: + super().__init__() + self._available_commands = sorted(available_commands, key=lambda command: command.name) + self._command_names = command_name_set(available_commands) + self._annotate_meta = annotate_meta + self._command_scope = command_scope + self._is_task_running = is_task_running + self._arg_suggestions = arg_suggestions or _no_arg_suggestions + + def _context(self, document: Document): + arguments = self._arg_suggestions() + return parse_completion_context( + document, + known_commands=self._command_names, + argument_commands=arguments.keys(), + allow_file=False, + ) + + def completion_active(self, document: Document) -> bool: + context = self._context(document) + return context.kind is CompletionKind.SLASH_COMMAND or ( + context.kind is CompletionKind.SLASH_ARGUMENT + and context.argument_index == 0 + and context.command in self._arg_suggestions() + ) + + @staticmethod + def should_complete(document: Document) -> bool: + context = parse_completion_context(document, allow_file=False) + return context.kind is CompletionKind.SLASH_COMMAND + + @override + def get_completions( + self, document: Document, complete_event: CompleteEvent + ) -> Iterable[Completion]: + context = self._context(document) + if context.kind is CompletionKind.SLASH_ARGUMENT: + if context.argument_index != 0 or context.command not in self._arg_suggestions(): + return + partial = context.token + for option in self._arg_suggestions()[context.command]: + if partial and not option.startswith(partial.lower()): + continue + yield Completion( + text=option[len(partial) :] if partial else option, + start_position=context.start_position, + display=option, + ) + return + if context.kind is not CompletionKind.SLASH_COMMAND: + return + + typed = context.token[1:] + typed_lower = typed.lower() + seen: set[str] = set() + + def emit(command: SlashCommand[Any], label: str | None = None) -> Iterable[Completion]: + if command.name in seen: + return + seen.add(command.name) + shown = label or command.name + yield Completion( + text=f"/{shown}", + start_position=context.start_position, + display=f"/{shown}", + display_meta=self._display_meta(command), + ) + + if not typed: + for command in self._available_commands: + yield from emit(command) + return + + def match_tier(command: SlashCommand[Any]) -> tuple[int, str] | None: + name_lower = command.name.lower() + if name_lower == typed_lower: + return (0, command.name) + if name_lower.startswith(typed_lower): + return (1, command.name) + alias_prefix: str | None = None + for alias in command.aliases: + alias_lower = alias.lower() + if alias_lower == typed_lower: + return (2, alias) + if alias_prefix is None and alias_lower.startswith(typed_lower): + alias_prefix = alias + if alias_prefix is not None: + return (3, alias_prefix) + segment = name_lower.split(":", 1)[1] if ":" in name_lower else name_lower + if ":" in name_lower: + if segment == typed_lower: + return (4, command.name) + if segment.startswith(typed_lower): + return (5, command.name) + if len(typed_lower) >= 2 and _fuzzy_subsequence(typed_lower, segment): + return (6, command.name) + return None + + matched: list[tuple[int, int, str, str, SlashCommand[Any]]] = [] + for command in self._available_commands: + result = match_tier(command) + if result is not None: + tier, label = result + matched.append((tier, len(command.name), command.name, label, command)) + matched.sort(key=lambda item: (item[0], item[1], item[2])) + if matched and matched[0][0] < 6: + matched = [item for item in matched if item[0] < 6] + for _, _, _, label, command in matched: + yield from emit(command, label) + + def _disabled_during_task(self, command: SlashCommand[Any]) -> bool: + if self._is_task_running is None or not self._is_task_running(): + return False + from pythinker_code.ui.shell.slash import registry as shell_registry + + shell_command = shell_registry.find_command(command.name) + return shell_command is not None and not shell_command.available_during_task + + def _display_meta(self, command: SlashCommand[Any]) -> str: + if self._disabled_during_task(command): + return "disabled while a task is in progress" + if not self._annotate_meta: + return command.description + if command.name.startswith("skill:"): + kind: str | None = "skill" + elif command.name.startswith("flow:"): + kind = "flow" + else: + kind = None + parts: list[str] = [] + if kind is not None: + parts.append(f"[{kind}]") + parts.append(command.description) + if command.aliases: + parts.append(f"aliases: {', '.join('/' + alias for alias in command.aliases)}") + return " ".join(part for part in parts if part) diff --git a/tests/ui_and_conv/test_file_completer.py b/tests/ui_and_conv/test_file_completer.py index a3908675..fa1a4fa6 100644 --- a/tests/ui_and_conv/test_file_completer.py +++ b/tests/ui_and_conv/test_file_completer.py @@ -10,6 +10,10 @@ from prompt_toolkit.document import Document from pythinker_code.ui.shell.prompt import LocalFileMentionCompleter +from pythinker_code.ui.shell.prompting.completion.context import ( + CompletionKind, + parse_completion_context, +) def _completion_texts(completer: LocalFileMentionCompleter, text: str) -> list[str]: @@ -92,6 +96,57 @@ def test_at_guard_prevents_email_like_fragments(tmp_path: Path): assert not texts +def test_file_context_matrix_for_boundaries_quotes_and_cursor_position(): + quoted = parse_completion_context(Document('@"docs/design notes.md')) + punctuation = parse_completion_context(Document("see (@src/main.py")) + email = parse_completion_context(Document("email@example.com")) + escaped_whitespace = parse_completion_context(Document(r"@docs/design\ notes.md")) + mid_token = parse_completion_context(Document(text="@src/main.py", cursor_position=len("@src"))) + + assert (quoted.kind, quoted.token, quoted.quoted) == ( + CompletionKind.FILE, + "docs/design notes.md", + True, + ) + assert (punctuation.kind, punctuation.token) == (CompletionKind.FILE, "src/main.py") + assert email.kind is CompletionKind.NONE + assert escaped_whitespace.kind is CompletionKind.NONE + assert (mid_token.kind, mid_token.token) == (CompletionKind.FILE, "src") + + +def test_quoted_completion_retains_quotes_and_escapes_path(tmp_path: Path): + docs = tmp_path / "docs" + docs.mkdir() + path = docs / 'design "notes"\\draft.md' + path.write_text("notes\n") + completer = LocalFileMentionCompleter(tmp_path) + document = Document('@"docs/design') + event = CompleteEvent(completion_requested=True) + + completions = list(completer.get_completions(document, event)) + + assert [completion.text for completion in completions] == [ + '"docs/design \\"notes\\"\\\\draft.md"' + ] + assert completions[0].start_position == -len('"docs/design') + + +def test_unicode_and_cjk_paths_complete(tmp_path: Path): + (tmp_path / "café.md").write_text("accent\n") + (tmp_path / "设计说明.md").write_text("CJK\n") + completer = LocalFileMentionCompleter(tmp_path) + + assert "café.md" in _completion_texts(completer, "@caf") + assert "设计说明.md" in _completion_texts(completer, "@设计") + + +def test_completed_quoted_file_short_circuits_completions(tmp_path: Path): + (tmp_path / "design notes.md").write_text("done\n") + completer = LocalFileMentionCompleter(tmp_path) + + assert not _completion_texts(completer, '@"design notes.md"') + + def test_scoped_walk_finds_late_alphabetical_dirs(tmp_path: Path): """Directories that sort late alphabetically must still be reachable. diff --git a/tests/ui_and_conv/test_slash_completer.py b/tests/ui_and_conv/test_slash_completer.py index fbcc8524..a7992dc6 100644 --- a/tests/ui_and_conv/test_slash_completer.py +++ b/tests/ui_and_conv/test_slash_completer.py @@ -22,6 +22,10 @@ _find_prompt_float_container, _wrap_to_width, ) +from pythinker_code.ui.shell.prompting.completion.context import ( + CompletionKind, + parse_completion_context, +) from pythinker_code.ui.shell.slash import slash_command_arg_suggestions from pythinker_code.utils.slashcmd import SlashCommand @@ -204,6 +208,50 @@ def test_should_complete_only_for_root_slash_token(): assert not SlashCommandCompleter.should_complete(Document(text="/he next", cursor_position=8)) +def test_completion_context_preserves_root_and_mid_line_slash_facets(): + root = parse_completion_context(Document("/he")) + mid_line_root = parse_completion_context(Document("please /he")) + mid_line_suggest = parse_completion_context(Document("please /he"), slash_activation="any") + + assert (root.kind, root.token, root.start_position) == ( + CompletionKind.SLASH_COMMAND, + "/he", + -3, + ) + assert mid_line_root.kind is CompletionKind.NONE + assert mid_line_suggest.kind is CompletionKind.SLASH_COMMAND + assert mid_line_suggest.token == "/he" + + +def test_completion_context_reports_first_and_later_arguments(): + known = frozenset({"theme"}) + + first = parse_completion_context( + Document("/theme cur"), known_commands=known, argument_commands=known + ) + later = parse_completion_context( + Document("/theme current extra"), known_commands=known, argument_commands=known + ) + escaped_whitespace = parse_completion_context( + Document(r"/theme current\ value"), known_commands=known, argument_commands=known + ) + + assert (first.kind, first.command, first.argument_index, first.token) == ( + CompletionKind.SLASH_ARGUMENT, + "theme", + 0, + "cur", + ) + assert (later.argument_index, later.token) == (1, "extra") + assert (escaped_whitespace.argument_index, escaped_whitespace.token) == (1, "value") + + +def test_completion_context_rejects_cursor_mid_slash_token(): + context = parse_completion_context(Document(text="/help", cursor_position=3)) + + assert context.kind is CompletionKind.NONE + + def test_completion_active_for_theme_subcommand(): completer = _theme_completer() assert completer.completion_active(Document(text="/theme ", cursor_position=len("/theme "))) diff --git a/tests/ui_and_conv/test_slash_highlight.py b/tests/ui_and_conv/test_slash_highlight.py index d0e7ec68..8d06e86f 100644 --- a/tests/ui_and_conv/test_slash_highlight.py +++ b/tests/ui_and_conv/test_slash_highlight.py @@ -125,6 +125,14 @@ def test_at_mention_at_start_highlighted(): assert _mentions(_lex_line("@README.md")) == ["@README.md"] +def test_quoted_at_mention_with_spaces_is_highlighted(): + assert _mentions(_lex_line('read @"docs/design notes.md" now')) == ['@"docs/design notes.md"'] + + +def test_mention_after_allowed_punctuation_is_highlighted(): + assert _mentions(_lex_line("read (@src/main.py), now")) == ["@src/main.py),"] + + def test_email_like_at_not_highlighted(): # "@" glued to an alphanumeric is not a mention boundary. assert _mentions(_lex_line("ping foo@bar.com")) == [] From 8e6291a5bb1f5a1bce3e97df3df619878fa60429 Mon Sep 17 00:00:00 2001 From: elkaix Date: Mon, 20 Jul 2026 07:13:05 -0400 Subject: [PATCH 2/5] feat(tui): index workspace mentions asynchronously --- CHANGELOG.md | 1 + src/pythinker_code/ui/shell/prompt.py | 226 ++------ .../ui/shell/prompting/completion/__init__.py | 12 + .../shell/prompting/completion/workspace.py | 546 ++++++++++++++++++ tests/ui_and_conv/test_file_completer.py | 148 +++-- tests/ui_and_conv/test_workspace_index.py | 322 +++++++++++ 6 files changed, 1004 insertions(+), 251 deletions(-) create mode 100644 src/pythinker_code/ui/shell/prompting/completion/workspace.py create mode 100644 tests/ui_and_conv/test_workspace_index.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 49968ca1..5cfb3bc8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,7 @@ GitHub Releases page; `0.8.0` is the new starting line. - Freeze the public shell prompt compatibility contract with constructor and rendering coverage. - Unify slash and file-mention completion behind one canonical completion context, adding quoted `@"path with spaces"` file mentions. +- Index workspace file mentions asynchronously with a cwd-aware, generation-owned snapshot index so completion never blocks on disk or Git scans. - Keep the interactive prompt scene within the terminal height with a priority-ordered row allocator, and shut prompt background tasks/processes down with an awaited lifecycle. - Add xAI Grok OAuth login (browser loopback and device-code). - Add GitHub Copilot device-code OAuth login for individual github.com accounts. diff --git a/src/pythinker_code/ui/shell/prompt.py b/src/pythinker_code/ui/shell/prompt.py index e29bf164..8a19fba3 100644 --- a/src/pythinker_code/ui/shell/prompt.py +++ b/src/pythinker_code/ui/shell/prompt.py @@ -11,27 +11,20 @@ import sys import time from collections import deque -from collections.abc import Awaitable, Callable, Iterable, Sequence +from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass from datetime import datetime from enum import Enum from hashlib import md5 from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal, Protocol, cast, override, runtime_checkable +from typing import TYPE_CHECKING, Any, Literal, Protocol, cast, runtime_checkable from prompt_toolkit import PromptSession from prompt_toolkit.application import Application from prompt_toolkit.application.current import get_app_or_none from prompt_toolkit.buffer import Buffer from prompt_toolkit.clipboard.pyperclip import PyperclipClipboard -from prompt_toolkit.completion import ( - CompleteEvent, - Completer, - Completion, - FuzzyCompleter, - WordCompleter, - merge_completers, -) +from prompt_toolkit.completion import Completion, merge_completers from prompt_toolkit.data_structures import Point from prompt_toolkit.document import Document from prompt_toolkit.filters import Condition, has_completions @@ -59,6 +52,7 @@ from prompt_toolkit.patch_stdout import patch_stdout from prompt_toolkit.utils import get_cwidth from pydantic import BaseModel, ValidationError +from pythinker_host import get_current_host from pythinker_host.path import HostPath from pythinker_code.config import StatusLineConfig @@ -97,6 +91,10 @@ command_name_set, discard_slash_command, ) +from pythinker_code.ui.shell.prompting.completion.workspace import ( + HostFileMentionCompleter, + WorkspaceIndex, +) from pythinker_code.ui.shell.prompting.lifecycle import PromptLifecycle from pythinker_code.ui.shell.prompting.state import ( BufferObserved, @@ -1163,183 +1161,7 @@ def _render_count_line( return fragments -class LocalFileMentionCompleter(Completer): - """Offer fuzzy `@` path completion by indexing workspace files. - - File discovery and ignore rules are delegated to - :mod:`pythinker_code.utils.file_filter` so that the web backend can reuse - them. - """ - - _FRAGMENT_PATTERN = re.compile(r"[^\s@]+") - - def __init__( - self, - root: Path, - *, - refresh_interval: float = 2.0, - limit: int = 1000, - ) -> None: - self._root = root - self._refresh_interval = refresh_interval - self._limit = limit - self._cache_time: float = 0.0 - self._cached_paths: list[str] = [] - self._cache_scope: str | None = None - self._top_cache_time: float = 0.0 - self._top_cached_paths: list[str] = [] - self._fragment_hint: str | None = None - self._is_git: bool | None = None # lazily detected - self._git_index_mtime: float | None = None - - self._word_completer = WordCompleter( - self._get_paths, - WORD=False, - pattern=self._FRAGMENT_PATTERN, - ) - - self._fuzzy = FuzzyCompleter( - self._word_completer, - WORD=False, - pattern=r"^[^\s@]*", - ) - - def _get_paths(self) -> list[str]: - fragment = self._fragment_hint or "" - if "/" not in fragment and len(fragment) < 3: - return self._get_top_level_paths() - return self._get_deep_paths() - - def _get_top_level_paths(self) -> list[str]: - from pythinker_code.utils.file_filter import is_ignored - - now = time.monotonic() - if now - self._top_cache_time <= self._refresh_interval: - return self._top_cached_paths - - entries: list[str] = [] - try: - for entry in sorted(self._root.iterdir(), key=lambda p: p.name): - name = entry.name - if is_ignored(name): - continue - entries.append(f"{name}/" if entry.is_dir() else name) - if len(entries) >= self._limit: - break - except OSError: - return self._top_cached_paths - - self._top_cached_paths = entries - self._top_cache_time = now - return self._top_cached_paths - - def _get_deep_paths(self) -> list[str]: - from pythinker_code.utils.file_filter import ( - detect_git, - git_index_mtime, - list_files_git, - list_files_walk, - ) - - fragment = self._fragment_hint or "" - - scope: str | None = None - if "/" in fragment: - scope = fragment.rsplit("/", 1)[0] - - now = time.monotonic() - cache_valid = ( - now - self._cache_time <= self._refresh_interval and self._cache_scope == scope - ) - - # Invalidate on .git/index mtime change. - if cache_valid and self._is_git: - mtime = git_index_mtime(self._root) - if mtime != self._git_index_mtime: - cache_valid = False - - if cache_valid: - return self._cached_paths - - if self._is_git is None: - self._is_git = detect_git(self._root) - - paths: list[str] | None = None - if self._is_git: - paths = list_files_git(self._root, scope) - self._git_index_mtime = git_index_mtime(self._root) - if paths is None: - paths = list_files_walk(self._root, scope, limit=self._limit) - - self._cached_paths = paths - self._cache_scope = scope - self._cache_time = now - return self._cached_paths - - @staticmethod - def should_complete(document: Document) -> bool: - """Return whether `@` file completion should be active for the buffer.""" - context = parse_completion_context(document, allow_slash=False) - return context.kind is CompletionKind.FILE - - def _is_completed_file(self, fragment: str) -> bool: - candidate = fragment.rstrip("/") - if not candidate: - return False - try: - return (self._root / candidate).is_file() - except OSError: - return False - - @override - def get_completions( - self, document: Document, complete_event: CompleteEvent - ) -> Iterable[Completion]: - context = parse_completion_context(document, allow_slash=False) - if context.kind is not CompletionKind.FILE: - return - fragment = context.token - if self._is_completed_file(fragment): - return - - mention_doc = Document(text=fragment, cursor_position=len(fragment)) - self._fragment_hint = fragment - try: - # First, ask the fuzzy completer for candidates. - candidates = list(self._fuzzy.get_completions(mention_doc, complete_event)) - - # re-rank: prefer basename matches - frag_lower = fragment.lower() - - def _rank(c: Completion) -> tuple[int, ...]: - path = c.text - base = path.rstrip("/").split("/")[-1].lower() - if base.startswith(frag_lower): - cat = 0 - elif frag_lower in base: - cat = 1 - else: - cat = 2 - test_penalty = int(any("test" in segment.lower() for segment in path.split("/"))) - # preserve original FuzzyCompleter's order in the same category - return (cat, test_penalty) - - candidates.sort(key=_rank) - if not context.quoted: - yield from candidates - return - for candidate in candidates: - escaped = candidate.text.replace("\\", "\\\\").replace('"', '\\"') - yield Completion( - text=f'"{escaped}"', - start_position=context.start_position, - display=candidate.display, - display_meta=candidate.display_meta, - style=candidate.style, - selected_style=candidate.selected_style, - ) - finally: - self._fragment_hint = None +LocalFileMentionCompleter = HostFileMentionCompleter class _HistoryEntry(BaseModel): @@ -1959,11 +1781,19 @@ def __init__( is_task_running=lambda: self._running_prompt_delegate is not None, arg_suggestions=self._slash_arg_suggestions, ) + self._workspace_root = HostPath.cwd() + self._workspace_index = WorkspaceIndex( + get_current_host(), + self._lifecycle, + self._workspace_root, + on_publish=self._on_workspace_snapshot_published, + ) + self._lifecycle.register_closer("workspace index", self._workspace_index.aclose) + self._file_mention_completer = HostFileMentionCompleter(self._workspace_index) self._agent_mode_completer = merge_completers( [ self._agent_slash_completer, - # TODO(host): we need an async HostFileMentionCompleter - LocalFileMentionCompleter(HostPath.cwd().unsafe_to_local_path()), + self._file_mention_completer, ], deduplicate=True, ) @@ -3843,7 +3673,25 @@ def running_prompt_accepts_submission(self) -> bool: return False return delegate.running_prompt_accepts_submission() + def _on_workspace_snapshot_published(self) -> None: + """Re-run file completion when a fresh workspace snapshot lands mid-menu.""" + app = self._session.app + if not app.is_running: + return + buffer = self._session.default_buffer + if not HostFileMentionCompleter.should_complete(buffer.document): + return + buffer.start_completion(select_first=False) + app.invalidate() + async def _prompt_once(self, *, append_history: bool | None) -> UserInput: + workspace_index = getattr(self, "_workspace_index", None) + if workspace_index is not None: + workspace_root = HostPath.cwd() + if workspace_root != getattr(self, "_workspace_root", None): + self._workspace_root = workspace_root + workspace_index.set_root(workspace_root) + workspace_index.request_refresh("") placeholder = None if (delegate := self._active_prompt_delegate()) is not None: placeholder = delegate.running_prompt_placeholder() diff --git a/src/pythinker_code/ui/shell/prompting/completion/__init__.py b/src/pythinker_code/ui/shell/prompting/completion/__init__.py index 52c46e1d..4f71153c 100644 --- a/src/pythinker_code/ui/shell/prompting/completion/__init__.py +++ b/src/pythinker_code/ui/shell/prompting/completion/__init__.py @@ -10,12 +10,24 @@ SlashCommandAutoSuggest, SlashCommandCompleter, ) +from pythinker_code.ui.shell.prompting.completion.workspace import ( + HostFileMentionCompleter, + LocalFileMentionCompleter, + WorkspaceEntry, + WorkspaceIndex, + WorkspaceSnapshot, +) __all__ = ( "CompletionContext", "CompletionKind", + "HostFileMentionCompleter", "InputHighlightLexer", + "LocalFileMentionCompleter", "SlashCommandAutoSuggest", "SlashCommandCompleter", + "WorkspaceEntry", + "WorkspaceIndex", + "WorkspaceSnapshot", "parse_completion_context", ) diff --git a/src/pythinker_code/ui/shell/prompting/completion/workspace.py b/src/pythinker_code/ui/shell/prompting/completion/workspace.py new file mode 100644 index 00000000..eabda395 --- /dev/null +++ b/src/pythinker_code/ui/shell/prompting/completion/workspace.py @@ -0,0 +1,546 @@ +"""Asynchronous workspace indexing for file-mention completion.""" + +from __future__ import annotations + +import asyncio +import re +import time +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from stat import S_ISDIR +from typing import override + +from prompt_toolkit.completion import ( + CompleteEvent, + Completer, + Completion, + FuzzyCompleter, + WordCompleter, +) +from prompt_toolkit.document import Document +from pythinker_host import AsyncReadable, Host, HostProcess +from pythinker_host.path import HostPath + +from pythinker_code.ui.shell.prompting.completion.context import ( + CompletionKind, + parse_completion_context, +) +from pythinker_code.ui.shell.prompting.lifecycle import PromptLifecycle +from pythinker_code.utils.file_filter import is_ignored +from pythinker_code.utils.logging import logger + +_MAX_ENTRIES = 1000 +_MAX_DEPTH = 8 +_MAX_DIRECTORY_VISITS = 256 +_MAX_DIRECTORY_ENTRIES = 2000 +_MAX_GIT_OUTPUT_BYTES = 4 * 1024 * 1024 +_GIT_TIMEOUT_SECONDS = 5.0 + + +@dataclass(frozen=True, slots=True) +class WorkspaceEntry: + """One immutable path in a workspace completion snapshot.""" + + relative_path: str + is_directory: bool + + +@dataclass(frozen=True, slots=True) +class WorkspaceSnapshot: + """One generation-owned immutable workspace scan result.""" + + root: HostPath + generation: int + entries: tuple[WorkspaceEntry, ...] + refreshed_at: float + degraded: bool + + +class WorkspaceIndex: + """Publish non-blocking workspace snapshots discovered through a ``Host``.""" + + def __init__( + self, + host: Host, + lifecycle: PromptLifecycle, + root: HostPath, + *, + refresh_interval: float = 2.0, + limit: int = _MAX_ENTRIES, + on_publish: Callable[[], None] | None = None, + ) -> None: + self._host = host + self._lifecycle = lifecycle + self._root = root + self._on_publish = on_publish + self._generation = 1 + self._refresh_interval = max(0.0, refresh_interval) + self._limit = min(_MAX_ENTRIES, max(1, limit)) + self._snapshots: dict[tuple[bool, str | None], WorkspaceSnapshot] = {} + self._refresh_task: asyncio.Task[None] | None = None + self._refresh_key: tuple[bool, str | None] | None = None + self._tasks: set[asyncio.Task[None]] = set() + self._closed = False + + def set_root(self, root: HostPath) -> None: + """Start a new ownership generation rooted at ``root``.""" + self._generation += 1 + self._root = root + self._snapshots.clear() + self._refresh_key = None + if self._refresh_task is not None and not self._refresh_task.done(): + self._refresh_task.cancel() + self._refresh_task = None + + def snapshot(self, fragment: str) -> tuple[WorkspaceEntry, ...]: + """Return the current immutable entries for ``fragment`` without I/O.""" + snapshot = self._snapshots.get(self._snapshot_key(fragment)) + if snapshot is None: + return () + return snapshot.entries + + def request_refresh(self, fragment: str) -> None: + """Schedule replacement discovery when the relevant snapshot is stale.""" + if self._closed: + return + + key = self._snapshot_key(fragment) + if ( + self._refresh_task is not None + and not self._refresh_task.done() + and self._refresh_key == key + ): + return + if self._refresh_task is not None and not self._refresh_task.done(): + self._refresh_task.cancel() + + current = self._snapshots.get(key) + now = time.monotonic() + if current is not None and now - current.refreshed_at <= self._refresh_interval: + self._refresh_task = None + self._refresh_key = None + return + + generation = self._generation + root = self._root + refresh = self._refresh(root, generation, key) + try: + task = self._lifecycle.create_task(refresh) + except RuntimeError as exc: + refresh.close() + logger.debug("Workspace refresh was not scheduled: error={!r}", exc) + return + self._refresh_task = task + self._refresh_key = key + self._tasks.add(task) + task.add_done_callback(self._tasks.discard) + + async def aclose(self) -> None: + """Cancel and await all workspace discovery owned by this index.""" + if self._closed: + return + self._closed = True + tasks = tuple(self._tasks) + for task in tasks: + if not task.done(): + task.cancel() + if tasks: + results = await asyncio.gather(*tasks, return_exceptions=True) + for result in results: + if isinstance(result, BaseException) and not isinstance( + result, asyncio.CancelledError + ): + logger.warning("Workspace refresh failed during shutdown: error={!r}", result) + self._refresh_task = None + self._refresh_key = None + + @staticmethod + def _snapshot_key(fragment: str) -> tuple[bool, str | None]: + if "/" not in fragment and len(fragment) < 3: + return (False, None) + if "/" not in fragment: + return (True, None) + scope = fragment.rsplit("/", 1)[0].strip("/") + if not scope or ".." in scope.split("/"): + return (True, None) + return (True, scope) + + async def _refresh( + self, + root: HostPath, + generation: int, + key: tuple[bool, str | None], + ) -> None: + try: + deep, scope = key + if deep: + entries, degraded = await self._discover_deep(root, scope) + else: + entries, degraded = await self._discover_top_level(root) + except asyncio.CancelledError: + raise + except Exception as exc: + logger.warning( + "Workspace refresh failed; publishing a degraded snapshot: root={} error={!r}", + root, + exc, + ) + entries = () + degraded = True + + if self._closed or generation != self._generation or root != self._root: + return + self._snapshots[key] = WorkspaceSnapshot( + root=root, + generation=generation, + entries=entries[: self._limit], + refreshed_at=time.monotonic(), + degraded=degraded, + ) + if self._on_publish is not None: + try: + self._on_publish() + except Exception as exc: + logger.warning("Workspace publish notification failed: error={!r}", exc) + + async def _discover_top_level(self, root: HostPath) -> tuple[tuple[WorkspaceEntry, ...], bool]: + children: list[HostPath] = [] + truncated = False + work = 0 + async for child in self._host.iterdir(root): + work += 1 + if work > _MAX_DIRECTORY_ENTRIES: + truncated = True + break + if is_ignored(child.name): + continue + if len(children) >= self._limit: + truncated = True + break + children.append(child) + + entries: list[WorkspaceEntry] = [] + degraded = truncated + if truncated: + logger.warning("Workspace top-level discovery was truncated: root={}", root) + for child in sorted(children, key=lambda path: path.name): + try: + stat_result = await self._host.stat(child) + except OSError as exc: + degraded = True + logger.debug( + "Workspace top-level entry could not be inspected: path={} error={!r}", + child, + exc, + ) + continue + entries.append(WorkspaceEntry(child.name, S_ISDIR(stat_result.st_mode))) + return tuple(entries), degraded + + async def _discover_deep( + self, root: HostPath, scope: str | None + ) -> tuple[tuple[WorkspaceEntry, ...], bool]: + repository = await self._is_git_repository(root) + if repository: + git_entries = await self._discover_git(root, scope) + if git_entries is not None: + return git_entries[: self._limit], False + logger.warning( + "Workspace Git discovery failed; using bounded traversal: root={} scope={!r}", + root, + scope, + ) + else: + logger.debug( + "Workspace is not a Git repository; using bounded traversal: root={} scope={!r}", + root, + scope, + ) + return await self._discover_bounded(root, scope), True + + async def _is_git_repository(self, root: HostPath) -> bool: + result = await self._run_git(root, "rev-parse", "--is-inside-work-tree") + if result is None: + return False + return result.strip() == "true" + + async def _discover_git( + self, root: HostPath, scope: str | None + ) -> tuple[WorkspaceEntry, ...] | None: + scope_args = ("--", f"{scope}/") if scope else () + tracked = await self._run_git( + root, + "-c", + "core.quotepath=false", + "ls-files", + "-z", + "--recurse-submodules", + *scope_args, + ) + if tracked is None: + return None + deleted = await self._run_git( + root, + "-c", + "core.quotepath=false", + "ls-files", + "-z", + "--deleted", + *scope_args, + ) + if deleted is None: + return None + untracked = await self._run_git( + root, + "-c", + "core.quotepath=false", + "ls-files", + "-z", + "--others", + "--exclude-standard", + *scope_args, + ) + if untracked is None: + return None + + deleted_paths = {path for path in deleted.split("\0") if path} + paths = [path for path in tracked.split("\0") if path and path not in deleted_paths] + tracked_paths = set(paths) + paths.extend( + path + for path in untracked.split("\0") + if path and path not in tracked_paths and path not in deleted_paths + ) + return self._entries_from_git_paths(paths) + + def _entries_from_git_paths(self, paths: Iterable[str]) -> tuple[WorkspaceEntry, ...]: + entries: list[WorkspaceEntry] = [] + seen: set[str] = set() + for path in paths: + if path.startswith("/"): + continue + normalized = path.strip("/") + parts = normalized.split("/") + if ( + not normalized + or normalized.startswith("../") + or ".." in parts + or any(is_ignored(part) for part in parts) + ): + continue + for index in range(1, len(parts)): + directory = "/".join(parts[:index]) + if directory not in seen: + seen.add(directory) + entries.append(WorkspaceEntry(directory, True)) + if len(entries) >= self._limit: + return tuple(entries) + if normalized not in seen: + seen.add(normalized) + entries.append(WorkspaceEntry(normalized, False)) + if len(entries) >= self._limit: + return tuple(entries) + return tuple(entries) + + async def _discover_bounded( + self, root: HostPath, scope: str | None + ) -> tuple[WorkspaceEntry, ...]: + if scope and (scope.startswith("/") or ".." in scope.split("/")): + return () + start = root.joinpath(*scope.split("/")) if scope else root + entries: list[WorkspaceEntry] = [] + if scope: + entries.append(WorkspaceEntry(scope, True)) + + queue: list[tuple[HostPath, str, int]] = [(start, scope or "", 0)] + directory_visits = 0 + work = 0 + while queue and len(entries) < self._limit and directory_visits < _MAX_DIRECTORY_VISITS: + directory, relative_directory, depth = queue.pop(0) + directory_visits += 1 + children: list[HostPath] = [] + try: + async for child in self._host.iterdir(directory): + work += 1 + if work > _MAX_DIRECTORY_ENTRIES: + logger.warning( + "Workspace traversal work limit reached: root={} scope={!r}", + root, + scope, + ) + return tuple(entries[: self._limit]) + if not is_ignored(child.name): + children.append(child) + except OSError as exc: + logger.debug( + "Workspace directory could not be traversed: path={} error={!r}", + directory, + exc, + ) + continue + + for child in sorted(children, key=lambda path: path.name): + relative_path = ( + f"{relative_directory}/{child.name}" if relative_directory else child.name + ) + try: + stat_result = await self._host.stat(child) + except OSError as exc: + logger.debug( + "Workspace entry could not be inspected: path={} error={!r}", child, exc + ) + continue + is_directory = S_ISDIR(stat_result.st_mode) + entries.append(WorkspaceEntry(relative_path, is_directory)) + if len(entries) >= self._limit: + break + if is_directory and depth < _MAX_DEPTH: + queue.append((child, relative_path, depth + 1)) + if queue: + logger.warning( + "Workspace traversal was truncated: root={} scope={!r} entries={} directories={}", + root, + scope, + len(entries), + directory_visits, + ) + return tuple(entries[: self._limit]) + + async def _run_git(self, root: HostPath, *args: str) -> str | None: + process: HostProcess | None = None + try: + async with asyncio.timeout(_GIT_TIMEOUT_SECONDS): + process = await self._host.exec("git", *args, cwd=str(root)) + stdout = await self._read_bounded(process.stdout) + await self._read_bounded(process.stderr) + returncode = await process.wait() + except asyncio.CancelledError: + if process is not None: + await self._stop_process(process) + raise + except TimeoutError: + if process is not None: + await self._stop_process(process) + logger.warning("Workspace Git command timed out: root={} command={!r}", root, args) + return None + except Exception as exc: + if process is not None: + await self._stop_process(process) + logger.warning( + "Workspace Git command could not run: root={} command={!r} error={!r}", + root, + args, + exc, + ) + return None + + if returncode != 0: + return None + encoding = "utf-8" + return stdout.decode(encoding=encoding, errors="replace") + + @staticmethod + async def _read_bounded(stream: AsyncReadable) -> bytes: + chunks = bytearray() + while True: + chunk = await stream.read(65536) + if not chunk: + return bytes(chunks) + chunks.extend(chunk) + if len(chunks) > _MAX_GIT_OUTPUT_BYTES: + raise RuntimeError("Git output exceeded workspace indexing limit") + + @staticmethod + async def _stop_process(process: HostProcess) -> None: + try: + await process.kill() + except Exception as exc: + logger.warning("Workspace Git process could not be killed: error={!r}", exc) + try: + await process.wait() + except Exception as exc: + logger.warning("Workspace Git process cleanup failed: error={!r}", exc) + + +class HostFileMentionCompleter(Completer): + """Offer fuzzy ``@`` completion using only immutable workspace snapshots.""" + + _FRAGMENT_PATTERN = re.compile(r"[^\s@]+") + + def __init__(self, index: WorkspaceIndex) -> None: + self._index = index + + @staticmethod + def should_complete(document: Document) -> bool: + """Return whether ``@`` file completion should be active for the buffer.""" + context = parse_completion_context(document, allow_slash=False) + return context.kind is CompletionKind.FILE + + @override + def get_completions( + self, document: Document, complete_event: CompleteEvent + ) -> Iterable[Completion]: + context = parse_completion_context(document, allow_slash=False) + if context.kind is not CompletionKind.FILE: + return + fragment = context.token + entries = self._index.snapshot(fragment) + self._index.request_refresh(fragment) + if any( + entry.relative_path == fragment.rstrip("/") and not entry.is_directory + for entry in entries + ): + return + + paths = [ + f"{entry.relative_path}/" if entry.is_directory else entry.relative_path + for entry in entries + ] + word_completer = WordCompleter( + paths, + WORD=False, + pattern=self._FRAGMENT_PATTERN, + ) + fuzzy = FuzzyCompleter(word_completer, WORD=False, pattern=r"^[^\s@]*") + mention_doc = Document(text=fragment, cursor_position=len(fragment)) + candidates = list(fuzzy.get_completions(mention_doc, complete_event)) + + frag_lower = fragment.lower() + + def _rank(completion: Completion) -> tuple[int, int]: + path = completion.text + basename = path.rstrip("/").split("/")[-1].lower() + if basename.startswith(frag_lower): + category = 0 + elif frag_lower in basename: + category = 1 + else: + category = 2 + test_penalty = int(any("test" in part.lower() for part in path.split("/"))) + return (category, test_penalty) + + candidates.sort(key=_rank) + if not context.quoted: + yield from candidates + return + for candidate in candidates: + escaped = candidate.text.replace("\\", "\\\\").replace('"', '\\"') + yield Completion( + text=f'"{escaped}"', + start_position=context.start_position, + display=candidate.display, + display_meta=candidate.display_meta, + style=candidate.style, + selected_style=candidate.selected_style, + ) + + +LocalFileMentionCompleter = HostFileMentionCompleter + +__all__ = ( + "HostFileMentionCompleter", + "LocalFileMentionCompleter", + "WorkspaceEntry", + "WorkspaceIndex", + "WorkspaceSnapshot", +) diff --git a/tests/ui_and_conv/test_file_completer.py b/tests/ui_and_conv/test_file_completer.py index fa1a4fa6..77860361 100644 --- a/tests/ui_and_conv/test_file_completer.py +++ b/tests/ui_and_conv/test_file_completer.py @@ -2,36 +2,74 @@ from __future__ import annotations +import asyncio import subprocess +from collections.abc import Coroutine from pathlib import Path +from typing import Any, override +import pytest from inline_snapshot import snapshot -from prompt_toolkit.completion import CompleteEvent +from prompt_toolkit.completion import CompleteEvent, Completion from prompt_toolkit.document import Document +from pythinker_host import get_current_host +from pythinker_host.path import HostPath from pythinker_code.ui.shell.prompt import LocalFileMentionCompleter from pythinker_code.ui.shell.prompting.completion.context import ( CompletionKind, parse_completion_context, ) - - -def _completion_texts(completer: LocalFileMentionCompleter, text: str) -> list[str]: +from pythinker_code.ui.shell.prompting.completion.workspace import WorkspaceIndex +from pythinker_code.ui.shell.prompting.lifecycle import PromptLifecycle + + +class _Lifecycle(PromptLifecycle): + def __init__(self) -> None: + super().__init__() + self.created: list[asyncio.Task[None]] = [] + + @override + def create_task(self, coro: Coroutine[Any, Any, None]) -> asyncio.Task[None]: + task = super().create_task(coro) + self.created.append(task) + return task + + +async def _completion_results(root: Path, text: str, *, limit: int = 1000) -> list[Completion]: + lifecycle = _Lifecycle() + index = WorkspaceIndex( + get_current_host(), + lifecycle, + HostPath.unsafe_from_local_path(root), + refresh_interval=60, + limit=limit, + ) + completer = LocalFileMentionCompleter(index) document = Document(text=text, cursor_position=len(text)) + context = parse_completion_context(document, allow_slash=False) + index.request_refresh(context.token) + if lifecycle.created: + await lifecycle.created[-1] event = CompleteEvent(completion_requested=True) - return [completion.text for completion in completer.get_completions(document, event)] + completions = list(completer.get_completions(document, event)) + await lifecycle.aclose() + return completions + + +async def _completion_texts(root: Path, text: str, *, limit: int = 1000) -> list[str]: + return [completion.text for completion in await _completion_results(root, text, limit=limit)] -def test_top_level_paths_skip_ignored_names(tmp_path: Path): +@pytest.mark.asyncio +async def test_top_level_paths_skip_ignored_names(tmp_path: Path): """Only surface non-ignored entries when completing the top level.""" (tmp_path / "src").mkdir() (tmp_path / "node_modules").mkdir() (tmp_path / ".DS_Store").write_text("") (tmp_path / "README.md").write_text("hello") - completer = LocalFileMentionCompleter(tmp_path) - - texts = _completion_texts(completer, "@") + texts = await _completion_texts(tmp_path, "@") assert "src/" in texts assert "README.md" in texts @@ -39,22 +77,22 @@ def test_top_level_paths_skip_ignored_names(tmp_path: Path): assert ".DS_Store" not in texts -def test_directory_completion_continues_after_slash(tmp_path: Path): +@pytest.mark.asyncio +async def test_directory_completion_continues_after_slash(tmp_path: Path): """Continue descending when the fragment ends with a slash.""" src = tmp_path / "src" src.mkdir() nested = src / "module.py" nested.write_text("print('hi')\n") - completer = LocalFileMentionCompleter(tmp_path) - - texts = _completion_texts(completer, "@src/") + texts = await _completion_texts(tmp_path, "@src/") assert "src/" in texts assert "src/module.py" in texts -def test_completed_file_short_circuits_completions(tmp_path: Path): +@pytest.mark.asyncio +async def test_completed_file_short_circuits_completions(tmp_path: Path): """Stop offering fuzzy matches once the fragment resolves to an existing file.""" agents = tmp_path / "AGENTS.md" agents.write_text("# Agents\n") @@ -63,14 +101,13 @@ def test_completed_file_short_circuits_completions(tmp_path: Path): nested_dir.mkdir(parents=True) (nested_dir / "README.md").write_text("nested\n") - completer = LocalFileMentionCompleter(tmp_path) - - texts = _completion_texts(completer, "@AGENTS.md") + texts = await _completion_texts(tmp_path, "@AGENTS.md") assert not texts -def test_limit_is_enforced(tmp_path: Path): +@pytest.mark.asyncio +async def test_limit_is_enforced(tmp_path: Path): """Respect the configured limit when building top-level candidates.""" for index in range(10): (tmp_path / f"dir{index}").mkdir() @@ -78,20 +115,17 @@ def test_limit_is_enforced(tmp_path: Path): (tmp_path / f"file{index}.txt").write_text("x") limit = 8 - completer = LocalFileMentionCompleter(tmp_path, limit=limit) - - texts = _completion_texts(completer, "@") + texts = await _completion_texts(tmp_path, "@", limit=limit) assert len(set(texts)) == limit -def test_at_guard_prevents_email_like_fragments(tmp_path: Path): +@pytest.mark.asyncio +async def test_at_guard_prevents_email_like_fragments(tmp_path: Path): """Ignore `@` that are embedded inside identifiers (e.g. emails).""" (tmp_path / "example.py").write_text("") - completer = LocalFileMentionCompleter(tmp_path) - - texts = _completion_texts(completer, "email@example.com") + texts = await _completion_texts(tmp_path, "email@example.com") assert not texts @@ -114,16 +148,13 @@ def test_file_context_matrix_for_boundaries_quotes_and_cursor_position(): assert (mid_token.kind, mid_token.token) == (CompletionKind.FILE, "src") -def test_quoted_completion_retains_quotes_and_escapes_path(tmp_path: Path): +@pytest.mark.asyncio +async def test_quoted_completion_retains_quotes_and_escapes_path(tmp_path: Path): docs = tmp_path / "docs" docs.mkdir() path = docs / 'design "notes"\\draft.md' path.write_text("notes\n") - completer = LocalFileMentionCompleter(tmp_path) - document = Document('@"docs/design') - event = CompleteEvent(completion_requested=True) - - completions = list(completer.get_completions(document, event)) + completions = await _completion_results(tmp_path, '@"docs/design') assert [completion.text for completion in completions] == [ '"docs/design \\"notes\\"\\\\draft.md"' @@ -131,23 +162,22 @@ def test_quoted_completion_retains_quotes_and_escapes_path(tmp_path: Path): assert completions[0].start_position == -len('"docs/design') -def test_unicode_and_cjk_paths_complete(tmp_path: Path): +@pytest.mark.asyncio +async def test_unicode_and_cjk_paths_complete(tmp_path: Path): (tmp_path / "café.md").write_text("accent\n") (tmp_path / "设计说明.md").write_text("CJK\n") - completer = LocalFileMentionCompleter(tmp_path) + assert "café.md" in await _completion_texts(tmp_path, "@caf") + assert "设计说明.md" in await _completion_texts(tmp_path, "@设计") - assert "café.md" in _completion_texts(completer, "@caf") - assert "设计说明.md" in _completion_texts(completer, "@设计") - -def test_completed_quoted_file_short_circuits_completions(tmp_path: Path): +@pytest.mark.asyncio +async def test_completed_quoted_file_short_circuits_completions(tmp_path: Path): (tmp_path / "design notes.md").write_text("done\n") - completer = LocalFileMentionCompleter(tmp_path) - - assert not _completion_texts(completer, '@"design notes.md"') + assert not await _completion_texts(tmp_path, '@"design notes.md"') -def test_scoped_walk_finds_late_alphabetical_dirs(tmp_path: Path): +@pytest.mark.asyncio +async def test_scoped_walk_finds_late_alphabetical_dirs(tmp_path: Path): """Directories that sort late alphabetically must still be reachable. Regression test for #1375: in large repos, ``os.walk`` exhausted the @@ -168,14 +198,13 @@ def test_scoped_walk_finds_late_alphabetical_dirs(tmp_path: Path): (target / "important.py").write_text("# find me") # With a low limit, the old os.walk approach would never reach zzz_target. - completer = LocalFileMentionCompleter(tmp_path, limit=50) - - texts = _completion_texts(completer, "@zzz_target/") + texts = await _completion_texts(tmp_path, "@zzz_target/", limit=50) assert "zzz_target/important.py" in texts -def test_basename_prefix_is_ranked_first(tmp_path: Path): +@pytest.mark.asyncio +async def test_basename_prefix_is_ranked_first(tmp_path: Path): """Prefer basename prefix matches over cross-segment fuzzy matches. For query 'fetch', we want '.../fetch.py' to appear before paths that only @@ -190,9 +219,7 @@ def test_basename_prefix_is_ranked_first(tmp_path: Path): patch_py = tmp_path / "src" / "pythinker_code" / "tools" / "file" / "patch.py" patch_py.write_text("# patch\n") - completer = LocalFileMentionCompleter(tmp_path) - - texts = _completion_texts(completer, "@fetch") + texts = await _completion_texts(tmp_path, "@fetch") # Snapshot the full candidate list to keep order/content deterministic assert texts == snapshot( @@ -203,16 +230,15 @@ def test_basename_prefix_is_ranked_first(tmp_path: Path): ) -def test_test_paths_are_ranked_after_source_matches(tmp_path: Path): +@pytest.mark.asyncio +async def test_test_paths_are_ranked_after_source_matches(tmp_path: Path): """Prefer source files over equally relevant test paths.""" (tmp_path / "tests").mkdir() (tmp_path / "zzz_src").mkdir() (tmp_path / "tests" / "foo.py").write_text("# test\n") (tmp_path / "zzz_src" / "foo.py").write_text("# source\n") - completer = LocalFileMentionCompleter(tmp_path) - - texts = _completion_texts(completer, "@foo") + texts = await _completion_texts(tmp_path, "@foo") assert texts[:2] == ["zzz_src/foo.py", "tests/foo.py"] @@ -229,7 +255,8 @@ def _init_git_repo(work_dir: Path) -> None: subprocess.run(cmd, cwd=work_dir, capture_output=True, check=True) -def test_tracked_ignored_dirs_filtered_in_git_mode(tmp_path: Path): +@pytest.mark.asyncio +async def test_tracked_ignored_dirs_filtered_in_git_mode(tmp_path: Path): """Tracked ``node_modules/`` and ``vendor/`` must still be filtered. Regression test: ``git ls-files`` returns all tracked paths, so @@ -247,20 +274,19 @@ def test_tracked_ignored_dirs_filtered_in_git_mode(tmp_path: Path): _init_git_repo(tmp_path) - completer = LocalFileMentionCompleter(tmp_path) - - texts = _completion_texts(completer, "@nod") + texts = await _completion_texts(tmp_path, "@nod") assert not any("node_modules" in t for t in texts), ( f"node_modules should be filtered even if tracked, got: {texts}" ) - texts = _completion_texts(completer, "@ven") + texts = await _completion_texts(tmp_path, "@ven") assert not any("vendor" in t for t in texts), ( f"vendor should be filtered even if tracked, got: {texts}" ) -def test_unstaged_rename_hides_deleted_path(tmp_path: Path): +@pytest.mark.asyncio +async def test_unstaged_rename_hides_deleted_path(tmp_path: Path): """After ``mv old.py new.py`` without staging, old.py must not appear. Regression test: ``git ls-files`` reads the index, so a file that was @@ -275,14 +301,12 @@ def test_unstaged_rename_hides_deleted_path(tmp_path: Path): # Rename without staging. (tmp_path / "src" / "old.py").rename(tmp_path / "src" / "new.py") - completer = LocalFileMentionCompleter(tmp_path) - - texts = _completion_texts(completer, "@old") + texts = await _completion_texts(tmp_path, "@old") assert not any("old.py" in t for t in texts), ( f"Deleted old.py should not appear in completion, got: {texts}" ) - texts = _completion_texts(completer, "@new") + texts = await _completion_texts(tmp_path, "@new") assert any("new.py" in t for t in texts), ( f"Renamed new.py should appear via --others, got: {texts}" ) diff --git a/tests/ui_and_conv/test_workspace_index.py b/tests/ui_and_conv/test_workspace_index.py new file mode 100644 index 00000000..0bdeb3ae --- /dev/null +++ b/tests/ui_and_conv/test_workspace_index.py @@ -0,0 +1,322 @@ +"""Tests for host-backed workspace file-mention indexing.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncGenerator, Coroutine +from stat import S_IFDIR, S_IFREG +from typing import Any, cast, override + +import pytest +from prompt_toolkit.completion import CompleteEvent +from prompt_toolkit.document import Document +from pythinker_host import Host, StatResult +from pythinker_host.path import HostPath + +from pythinker_code.ui.shell.prompt import LocalFileMentionCompleter +from pythinker_code.ui.shell.prompting.completion.workspace import ( + HostFileMentionCompleter, + WorkspaceIndex, +) +from pythinker_code.ui.shell.prompting.lifecycle import PromptLifecycle + + +class _Lifecycle(PromptLifecycle): + def __init__(self) -> None: + super().__init__() + self.created: list[asyncio.Task[None]] = [] + + @override + def create_task(self, coro: Coroutine[Any, Any, None]) -> asyncio.Task[None]: + task = super().create_task(coro) + self.created.append(task) + return task + + async def drain(self) -> None: + if self.created: + await asyncio.gather(*self.created, return_exceptions=True) + + +class _Stream: + def __init__(self, data: bytes) -> None: + self._data = data + self._read = False + + async def read(self, _n: int = -1) -> bytes: + if self._read: + return b"" + self._read = True + return self._data + + +class _Process: + def __init__(self, stdout: str, returncode: int) -> None: + encoding = "utf-8" + self.stdout = _Stream(stdout.encode(encoding=encoding)) + self.stderr = _Stream(b"") + self.stdin = cast(Any, None) + self.pid = 1 + self.returncode: int | None = None + self._exit_code = returncode + + async def wait(self) -> int: + self.returncode = self._exit_code + return self._exit_code + + async def kill(self) -> None: + self.returncode = -1 + + +class _FakeHost: + name = "workspace-test" + + def __init__(self) -> None: + self.children: dict[str, list[tuple[str, bool]]] = {} + self.iterdir_calls = 0 + self.stat_calls = 0 + self.exec_calls = 0 + self.git_repository = False + self.git_listing_fails = False + + def add_directory(self, path: str, children: list[tuple[str, bool]]) -> None: + self.children[path] = children + + async def _iterdir(self, path: HostPath) -> AsyncGenerator[HostPath]: + self.iterdir_calls += 1 + for name, _is_directory in self.children.get(str(path), []): + yield path / name + + def iterdir(self, path: HostPath) -> AsyncGenerator[HostPath]: + return self._iterdir(path) + + async def stat(self, path: HostPath, *, follow_symlinks: bool = True) -> StatResult: + del follow_symlinks + self.stat_calls += 1 + parent_children = self.children.get(str(path.parent), []) + is_directory = next( + (directory for name, directory in parent_children if name == path.name), False + ) + mode = (S_IFDIR if is_directory else S_IFREG) | 0o755 + return StatResult(mode, 0, 0, 0, 0, 0, 0, 0.0, 0.0, 0.0) + + async def exec(self, *args: str, env: object = None, cwd: str | None = None) -> _Process: + del env, cwd + self.exec_calls += 1 + if "rev-parse" in args: + output = "true\n" if self.git_repository else "" + return _Process(output, 0 if self.git_repository else 1) + if self.git_listing_fails: + return _Process("", 1) + return _Process("", 0) + + +class _ControlledHost(_FakeHost): + def __init__(self) -> None: + super().__init__() + self.started: dict[str, asyncio.Event] = {} + self.release: dict[str, asyncio.Event] = {} + self.ignore_cancellation: set[str] = set() + self.cancelled: set[str] = set() + + async def _iterdir(self, path: HostPath) -> AsyncGenerator[HostPath]: + root = str(path) + self.iterdir_calls += 1 + self.started.setdefault(root, asyncio.Event()).set() + gate = self.release.setdefault(root, asyncio.Event()) + try: + await gate.wait() + except asyncio.CancelledError: + self.cancelled.add(root) + if root not in self.ignore_cancellation: + raise + await gate.wait() + for name, _is_directory in self.children.get(root, []): + yield path / name + + +class _InspectableWorkspaceIndex(WorkspaceIndex): + def degraded(self) -> bool: + return next(iter(self._snapshots.values())).degraded + + +def _index( + host: _FakeHost, root: HostPath, *, limit: int = 1000 +) -> tuple[_InspectableWorkspaceIndex, _Lifecycle]: + lifecycle = _Lifecycle() + return ( + _InspectableWorkspaceIndex( + cast(Host, host), lifecycle, root, refresh_interval=60, limit=limit + ), + lifecycle, + ) + + +def _completion_texts(completer: HostFileMentionCompleter, text: str) -> list[str]: + document = Document(text=text, cursor_position=len(text)) + event = CompleteEvent(completion_requested=True) + return [completion.text for completion in completer.get_completions(document, event)] + + +def test_local_completer_name_is_a_facade_alias() -> None: + assert LocalFileMentionCompleter is HostFileMentionCompleter + + +@pytest.mark.asyncio +async def test_late_old_generation_never_publishes() -> None: + host = _ControlledHost() + root_a = HostPath("/workspace-a") + root_b = HostPath("/workspace-b") + host.add_directory(str(root_a), [("from-a.py", False)]) + host.add_directory(str(root_b), [("from-b.py", False)]) + host.ignore_cancellation.add(str(root_a)) + index, lifecycle = _index(host, root_a) + + index.request_refresh("") + await host.started.setdefault(str(root_a), asyncio.Event()).wait() + index.set_root(root_b) + index.request_refresh("") + await host.started.setdefault(str(root_b), asyncio.Event()).wait() + + host.release[str(root_b)].set() + await lifecycle.created[-1] + host.release[str(root_a)].set() + await lifecycle.drain() + + assert [entry.relative_path for entry in index.snapshot("")] == ["from-b.py"] + await lifecycle.aclose() + + +@pytest.mark.asyncio +async def test_obsolete_refresh_is_cancelled() -> None: + host = _ControlledHost() + root = HostPath("/workspace") + host.add_directory(str(root), [("file.py", False)]) + index, lifecycle = _index(host, root) + + index.request_refresh("") + await host.started.setdefault(str(root), asyncio.Event()).wait() + index.request_refresh("source") + await asyncio.sleep(0) + + assert str(root) in host.cancelled + await index.aclose() + await lifecycle.aclose() + + +@pytest.mark.asyncio +async def test_git_failure_uses_degraded_fallback() -> None: + host = _FakeHost() + root = HostPath("/workspace") + host.git_repository = True + host.git_listing_fails = True + host.add_directory(str(root), [("source.py", False)]) + index, lifecycle = _index(host, root) + + index.request_refresh("sou") + await lifecycle.drain() + + assert [entry.relative_path for entry in index.snapshot("sou")] == ["source.py"] + assert index.degraded() is True + await lifecycle.aclose() + + +@pytest.mark.asyncio +async def test_results_are_truncated_at_one_thousand_entries() -> None: + host = _FakeHost() + root = HostPath("/workspace") + host.add_directory(str(root), [(f"file-{index:04}.py", False) for index in range(1100)]) + index, lifecycle = _index(host, root) + + index.request_refresh("") + await lifecycle.drain() + + assert len(index.snapshot("")) == 1000 + await lifecycle.aclose() + + +@pytest.mark.asyncio +async def test_bounded_discovery_filters_ignored_paths() -> None: + host = _FakeHost() + root = HostPath("/workspace") + host.add_directory( + str(root), + [("node_modules", True), ("src", True), (".DS_Store", False)], + ) + host.add_directory(str(root / "src"), [("app.py", False), ("__pycache__", True)]) + index, lifecycle = _index(host, root) + + index.request_refresh("source") + await lifecycle.drain() + + paths = [entry.relative_path for entry in index.snapshot("source")] + assert "src/app.py" in paths + assert not any("node_modules" in path or "__pycache__" in path for path in paths) + assert ".DS_Store" not in paths + await lifecycle.aclose() + + +@pytest.mark.asyncio +async def test_quoted_path_with_spaces_uses_snapshot() -> None: + host = _FakeHost() + root = HostPath("/workspace") + host.add_directory(str(root), [("docs", True)]) + host.add_directory(str(root / "docs"), [("design notes.md", False)]) + index, lifecycle = _index(host, root) + index.request_refresh("docs/design") + await lifecycle.drain() + + completer = HostFileMentionCompleter(index) + + assert _completion_texts(completer, '@"docs/design') == ['"docs/design notes.md"'] + await lifecycle.aclose() + + +@pytest.mark.asyncio +async def test_snapshot_and_completion_reads_perform_no_host_io() -> None: + host = _FakeHost() + root = HostPath("/workspace") + host.add_directory(str(root), [("source.py", False)]) + index, lifecycle = _index(host, root) + index.request_refresh("") + await lifecycle.drain() + completer = HostFileMentionCompleter(index) + calls_before = (host.iterdir_calls, host.stat_calls, host.exec_calls) + + assert index.snapshot("") + assert _completion_texts(completer, "@s") == ["source.py"] + + assert (host.iterdir_calls, host.stat_calls, host.exec_calls) == calls_before + await lifecycle.aclose() + + +@pytest.mark.asyncio +async def test_on_publish_fires_for_fresh_snapshot_only() -> None: + host = _ControlledHost() + root_a = HostPath("/workspace-a") + root_b = HostPath("/workspace-b") + host.add_directory(str(root_a), [("from-a.py", False)]) + host.add_directory(str(root_b), [("from-b.py", False)]) + host.ignore_cancellation.add(str(root_a)) + lifecycle = _Lifecycle() + published: list[int] = [] + index = WorkspaceIndex( + cast(Host, host), + lifecycle, + root_a, + refresh_interval=60, + on_publish=lambda: published.append(1), + ) + + index.request_refresh("") + await host.started.setdefault(str(root_a), asyncio.Event()).wait() + index.set_root(root_b) + index.request_refresh("") + await host.started.setdefault(str(root_b), asyncio.Event()).wait() + + host.release[str(root_b)].set() + await lifecycle.created[-1] + host.release[str(root_a)].set() + await lifecycle.drain() + + assert len(published) == 1 + await lifecycle.aclose() From 14d58a5a744e403f3043af428d60601c698f2c35 Mon Sep 17 00:00:00 2001 From: elkaix Date: Mon, 20 Jul 2026 12:02:28 -0400 Subject: [PATCH 3/5] ci(typos): allow @caf completion-prefix test token --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 6362adec..113f0a34 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -194,6 +194,8 @@ iterm = "iterm" informations = "informations" # Project/domain-specific names and generated IDs. reviewr = "reviewr" +# ASCII prefix used in a file-completion test (@caf -> "café.md"). +caf = "caf" fnd = "fnd" edn = "edn" Encrypter = "Encrypter" From 42179a6ef53f3f72bf8e948e430af3639ee3e6cb Mon Sep 17 00:00:00 2001 From: elkaix Date: Mon, 20 Jul 2026 16:58:30 -0400 Subject: [PATCH 4/5] fix(tui): address review findings on prompt completion - Escape now dismisses the slash-argument menu (cancel_completion fallback) instead of doing nothing when there is no draft command to strip. - Strip trailing prose punctuation from unquoted @-mention highlight spans. - Rank scoped @dir/name completions by the fragment's basename so prefix priority applies to path-scoped fragments. - Drain git stdout/stderr concurrently to avoid a pipe deadlock on verbose output. - Guard the mid-session HostPath.cwd() lookup so a removed working directory no longer crashes the prompt turn. - Collapse a dead SLASH_ARGUMENT branch in the completion parser. - De-duplicate the PromptLifecycle test double, make the degraded-state test helper multi-key-safe, and re-frame the backslash-whitespace parser test to document the real (non-escape-aware) contract. --- CHANGELOG.md | 1 + src/pythinker_code/ui/shell/prompt.py | 26 +++++++++-- .../ui/shell/prompting/completion/context.py | 12 ++--- .../shell/prompting/completion/workspace.py | 17 +++++-- tests/ui_and_conv/_prompt_lifecycle.py | 27 +++++++++++ tests/ui_and_conv/test_file_completer.py | 17 +------ tests/ui_and_conv/test_slash_completer.py | 22 ++++++++- tests/ui_and_conv/test_slash_highlight.py | 3 +- tests/ui_and_conv/test_workspace_index.py | 46 +++++++++++-------- 9 files changed, 117 insertions(+), 54 deletions(-) create mode 100644 tests/ui_and_conv/_prompt_lifecycle.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5cfb3bc8..75556432 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,7 @@ GitHub Releases page; `0.8.0` is the new starting line. - Freeze the public shell prompt compatibility contract with constructor and rendering coverage. - Unify slash and file-mention completion behind one canonical completion context, adding quoted `@"path with spaces"` file mentions. - Index workspace file mentions asynchronously with a cwd-aware, generation-owned snapshot index so completion never blocks on disk or Git scans. +- Harden prompt completion: Escape now dismisses the slash-argument menu (not just draft commands), file-mention highlighting no longer swallows trailing prose punctuation, scoped `@dir/name` completions rank basename-prefix matches first, workspace Git scans drain stdout/stderr concurrently to avoid a pipe deadlock, and a mid-session working-directory removal no longer crashes the prompt turn. - Keep the interactive prompt scene within the terminal height with a priority-ordered row allocator, and shut prompt background tasks/processes down with an awaited lifecycle. - Add xAI Grok OAuth login (browser loopback and device-code). - Add GitHub Copilot device-code OAuth login for individual github.com accounts. diff --git a/src/pythinker_code/ui/shell/prompt.py b/src/pythinker_code/ui/shell/prompt.py index 8a19fba3..d230dc2d 100644 --- a/src/pythinker_code/ui/shell/prompt.py +++ b/src/pythinker_code/ui/shell/prompt.py @@ -1864,9 +1864,16 @@ def _(event: KeyPressEvent) -> None: @_kb.add("escape", eager=True, filter=_slash_completion_filter) def _(event: KeyPressEvent) -> None: - """Slash command completion: Escape discards the draft command.""" - if _discard_slash_command(event.current_buffer): - event.app.invalidate() + """Slash completion: Escape discards a draft command, or dismisses + the argument menu when there is no draft command to remove.""" + buffer = event.current_buffer + if not _discard_slash_command(buffer): + # Slash-argument completion (e.g. "/model gpt"): the eager + # binding swallowed Escape but there is no root command to + # strip, so dismiss the completion menu explicitly instead of + # leaving it open. + buffer.cancel_completion() + event.app.invalidate() @_kb.add("enter", filter=_non_slash_completion_filter) def _(event: KeyPressEvent) -> None: @@ -3687,8 +3694,17 @@ def _on_workspace_snapshot_published(self) -> None: async def _prompt_once(self, *, append_history: bool | None) -> UserInput: workspace_index = getattr(self, "_workspace_index", None) if workspace_index is not None: - workspace_root = HostPath.cwd() - if workspace_root != getattr(self, "_workspace_root", None): + try: + workspace_root: HostPath | None = HostPath.cwd() + except OSError: + # CWD was removed mid-session (e.g. an external drive was + # unplugged). Keep the last known root instead of crashing the + # prompt turn; the statusline render raises CwdLostError on the + # same turn to exit gracefully. + workspace_root = None + if workspace_root is not None and workspace_root != getattr( + self, "_workspace_root", None + ): self._workspace_root = workspace_root workspace_index.set_root(workspace_root) workspace_index.request_refresh("") diff --git a/src/pythinker_code/ui/shell/prompting/completion/context.py b/src/pythinker_code/ui/shell/prompting/completion/context.py index 59b754ce..3a4bc5b3 100644 --- a/src/pythinker_code/ui/shell/prompting/completion/context.py +++ b/src/pythinker_code/ui/shell/prompting/completion/context.py @@ -35,6 +35,10 @@ class CompletionContext: _NONE = CompletionContext(CompletionKind.NONE, "", 0, None, None) _COMMAND_CHARS = frozenset("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_:.-") _MENTION_TRIGGER_GUARDS = frozenset((".", "-", "_", "`", "'", '"', ":", "@", "#", "~")) +# Trailing prose punctuation that should not be pulled into an unquoted mention +# span (e.g. "(@src/main.py), now" -> "@src/main.py"). "." and "/" are excluded +# because they are common in real filenames and paths. +_MENTION_TRAILING_PUNCT = frozenset(")]}>,;:!?`'\"") _SLASH_TOKEN_RE = re.compile(r"(? index and line[end - 1] in _MENTION_TRAILING_PUNCT: + end -= 1 parsed = parse_completion_context( Document(line[:end], cursor_position=end), allow_slash=False, diff --git a/src/pythinker_code/ui/shell/prompting/completion/workspace.py b/src/pythinker_code/ui/shell/prompting/completion/workspace.py index eabda395..61354875 100644 --- a/src/pythinker_code/ui/shell/prompting/completion/workspace.py +++ b/src/pythinker_code/ui/shell/prompting/completion/workspace.py @@ -411,8 +411,13 @@ async def _run_git(self, root: HostPath, *args: str) -> str | None: try: async with asyncio.timeout(_GIT_TIMEOUT_SECONDS): process = await self._host.exec("git", *args, cwd=str(root)) - stdout = await self._read_bounded(process.stdout) - await self._read_bounded(process.stderr) + # Drain stdout and stderr concurrently: reading them + # sequentially can deadlock if git fills its stderr pipe buffer + # (e.g. submodule warnings) while we are still draining stdout. + stdout, _stderr = await asyncio.gather( + self._read_bounded(process.stdout), + self._read_bounded(process.stderr), + ) returncode = await process.wait() except asyncio.CancelledError: if process is not None: @@ -506,13 +511,17 @@ def get_completions( candidates = list(fuzzy.get_completions(mention_doc, complete_event)) frag_lower = fragment.lower() + # The typed fragment often carries a directory prefix (e.g. "src/mai"), + # but candidates are ranked by basename, so compare against the + # fragment's basename or prefix priority never applies to scoped paths. + frag_basename = frag_lower.rsplit("/", 1)[-1] def _rank(completion: Completion) -> tuple[int, int]: path = completion.text basename = path.rstrip("/").split("/")[-1].lower() - if basename.startswith(frag_lower): + if basename.startswith(frag_basename): category = 0 - elif frag_lower in basename: + elif frag_basename in basename: category = 1 else: category = 2 diff --git a/tests/ui_and_conv/_prompt_lifecycle.py b/tests/ui_and_conv/_prompt_lifecycle.py new file mode 100644 index 00000000..0ec777ec --- /dev/null +++ b/tests/ui_and_conv/_prompt_lifecycle.py @@ -0,0 +1,27 @@ +"""Shared PromptLifecycle test double for prompt-completion tests.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Coroutine +from typing import Any, override + +from pythinker_code.ui.shell.prompting.lifecycle import PromptLifecycle + + +class RecordingLifecycle(PromptLifecycle): + """PromptLifecycle that records created tasks so tests can await them.""" + + def __init__(self) -> None: + super().__init__() + self.created: list[asyncio.Task[None]] = [] + + @override + def create_task(self, coro: Coroutine[Any, Any, None]) -> asyncio.Task[None]: + task = super().create_task(coro) + self.created.append(task) + return task + + async def drain(self) -> None: + if self.created: + await asyncio.gather(*self.created, return_exceptions=True) diff --git a/tests/ui_and_conv/test_file_completer.py b/tests/ui_and_conv/test_file_completer.py index 77860361..4247dfbc 100644 --- a/tests/ui_and_conv/test_file_completer.py +++ b/tests/ui_and_conv/test_file_completer.py @@ -2,11 +2,8 @@ from __future__ import annotations -import asyncio import subprocess -from collections.abc import Coroutine from pathlib import Path -from typing import Any, override import pytest from inline_snapshot import snapshot @@ -21,19 +18,7 @@ parse_completion_context, ) from pythinker_code.ui.shell.prompting.completion.workspace import WorkspaceIndex -from pythinker_code.ui.shell.prompting.lifecycle import PromptLifecycle - - -class _Lifecycle(PromptLifecycle): - def __init__(self) -> None: - super().__init__() - self.created: list[asyncio.Task[None]] = [] - - @override - def create_task(self, coro: Coroutine[Any, Any, None]) -> asyncio.Task[None]: - task = super().create_task(coro) - self.created.append(task) - return task +from tests.ui_and_conv._prompt_lifecycle import RecordingLifecycle as _Lifecycle async def _completion_results(root: Path, text: str, *, limit: int = 1000) -> list[Completion]: diff --git a/tests/ui_and_conv/test_slash_completer.py b/tests/ui_and_conv/test_slash_completer.py index a7992dc6..e322752a 100644 --- a/tests/ui_and_conv/test_slash_completer.py +++ b/tests/ui_and_conv/test_slash_completer.py @@ -232,7 +232,10 @@ def test_completion_context_reports_first_and_later_arguments(): later = parse_completion_context( Document("/theme current extra"), known_commands=known, argument_commands=known ) - escaped_whitespace = parse_completion_context( + # A backslash before a space is NOT an escape: the parser splits on raw + # whitespace, so "current\ value" is two arguments and the cursor token is + # the trailing "value" at argument index 1 (escape sequences unsupported). + backslash_before_space = parse_completion_context( Document(r"/theme current\ value"), known_commands=known, argument_commands=known ) @@ -243,7 +246,10 @@ def test_completion_context_reports_first_and_later_arguments(): "cur", ) assert (later.argument_index, later.token) == (1, "extra") - assert (escaped_whitespace.argument_index, escaped_whitespace.token) == (1, "value") + assert (backslash_before_space.argument_index, backslash_before_space.token) == ( + 1, + "value", + ) def test_completion_context_rejects_cursor_mid_slash_token(): @@ -380,6 +386,18 @@ def test_discard_slash_command_ignores_non_root_slash_text(): assert buffer.text == "ask /theme" +def test_discard_slash_command_declines_slash_argument(): + # A slash *argument* ("/theme cur") has no root command draft to strip, so + # discard declines it and leaves the text intact. The Escape keybinding + # relies on this False return to fall back to buffer.cancel_completion(), + # which dismisses the argument menu instead of doing nothing. + buffer = Buffer() + buffer.set_document(Document(text="/theme cur", cursor_position=10), bypass_readonly=True) + + assert _discard_slash_command(buffer) is False + assert buffer.text == "/theme cur" + + def test_completion_display_uses_canonical_command_name(): completer = SlashCommandCompleter( [ diff --git a/tests/ui_and_conv/test_slash_highlight.py b/tests/ui_and_conv/test_slash_highlight.py index 8d06e86f..1b17265c 100644 --- a/tests/ui_and_conv/test_slash_highlight.py +++ b/tests/ui_and_conv/test_slash_highlight.py @@ -130,7 +130,8 @@ def test_quoted_at_mention_with_spaces_is_highlighted(): def test_mention_after_allowed_punctuation_is_highlighted(): - assert _mentions(_lex_line("read (@src/main.py), now")) == ["@src/main.py),"] + # Trailing prose punctuation (")", ",") must stay out of the mention span. + assert _mentions(_lex_line("read (@src/main.py), now")) == ["@src/main.py"] def test_email_like_at_not_highlighted(): diff --git a/tests/ui_and_conv/test_workspace_index.py b/tests/ui_and_conv/test_workspace_index.py index 0bdeb3ae..92153c7d 100644 --- a/tests/ui_and_conv/test_workspace_index.py +++ b/tests/ui_and_conv/test_workspace_index.py @@ -3,9 +3,9 @@ from __future__ import annotations import asyncio -from collections.abc import AsyncGenerator, Coroutine +from collections.abc import AsyncGenerator from stat import S_IFDIR, S_IFREG -from typing import Any, cast, override +from typing import Any, cast import pytest from prompt_toolkit.completion import CompleteEvent @@ -18,23 +18,7 @@ HostFileMentionCompleter, WorkspaceIndex, ) -from pythinker_code.ui.shell.prompting.lifecycle import PromptLifecycle - - -class _Lifecycle(PromptLifecycle): - def __init__(self) -> None: - super().__init__() - self.created: list[asyncio.Task[None]] = [] - - @override - def create_task(self, coro: Coroutine[Any, Any, None]) -> asyncio.Task[None]: - task = super().create_task(coro) - self.created.append(task) - return task - - async def drain(self) -> None: - if self.created: - await asyncio.gather(*self.created, return_exceptions=True) +from tests.ui_and_conv._prompt_lifecycle import RecordingLifecycle as _Lifecycle class _Stream: @@ -135,8 +119,11 @@ async def _iterdir(self, path: HostPath) -> AsyncGenerator[HostPath]: class _InspectableWorkspaceIndex(WorkspaceIndex): + # White-box accessor for the index's internal degradation bookkeeping, which + # has no production consumer to observe. Aggregates across every primed + # snapshot so the assertion holds regardless of how many keys are populated. def degraded(self) -> bool: - return next(iter(self._snapshots.values())).degraded + return any(snapshot.degraded for snapshot in self._snapshots.values()) def _index( @@ -271,6 +258,25 @@ async def test_quoted_path_with_spaces_uses_snapshot() -> None: await lifecycle.aclose() +@pytest.mark.asyncio +async def test_scoped_completion_ranks_basename_prefix_first() -> None: + # With a directory prefix in the fragment ("src/mai"), ranking must compare + # against the fragment's basename ("mai"): "main.py" (basename prefix) should + # sort before "domain.py" (basename only contains "mai"). + host = _FakeHost() + root = HostPath("/workspace") + host.add_directory(str(root), [("src", True)]) + host.add_directory(str(root / "src"), [("main.py", False), ("domain.py", False)]) + index, lifecycle = _index(host, root) + index.request_refresh("src/mai") + await lifecycle.drain() + + completer = HostFileMentionCompleter(index) + + assert _completion_texts(completer, "@src/mai")[:2] == ["src/main.py", "src/domain.py"] + await lifecycle.aclose() + + @pytest.mark.asyncio async def test_snapshot_and_completion_reads_perform_no_host_io() -> None: host = _FakeHost() From 7627f669fa323c8aecd5035aceb7c1adc023cb96 Mon Sep 17 00:00:00 2001 From: elkaix Date: Mon, 20 Jul 2026 18:41:30 -0400 Subject: [PATCH 5/5] fix(tui): cancel sibling stream reader when a git reader fails MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The concurrent stdout/stderr drain in _run_git used asyncio.gather, which leaves the sibling _read_bounded coroutine running after the first reader raises — only the process was killed, so the other reader could keep draining a closing pipe. Move the concurrent read into _read_streams, which uses explicit tasks and cancels+awaits both readers on any failure (or cancellation) before propagating. Adds a regression test for a reader-exception path. --- .../shell/prompting/completion/workspace.py | 23 ++++++++-- tests/ui_and_conv/test_workspace_index.py | 42 +++++++++++++++++++ 2 files changed, 61 insertions(+), 4 deletions(-) diff --git a/src/pythinker_code/ui/shell/prompting/completion/workspace.py b/src/pythinker_code/ui/shell/prompting/completion/workspace.py index 61354875..67127a89 100644 --- a/src/pythinker_code/ui/shell/prompting/completion/workspace.py +++ b/src/pythinker_code/ui/shell/prompting/completion/workspace.py @@ -414,10 +414,7 @@ async def _run_git(self, root: HostPath, *args: str) -> str | None: # Drain stdout and stderr concurrently: reading them # sequentially can deadlock if git fills its stderr pipe buffer # (e.g. submodule warnings) while we are still draining stdout. - stdout, _stderr = await asyncio.gather( - self._read_bounded(process.stdout), - self._read_bounded(process.stderr), - ) + stdout = await self._read_streams(process) returncode = await process.wait() except asyncio.CancelledError: if process is not None: @@ -444,6 +441,24 @@ async def _run_git(self, root: HostPath, *args: str) -> str | None: encoding = "utf-8" return stdout.decode(encoding=encoding, errors="replace") + async def _read_streams(self, process: HostProcess) -> bytes: + """Drain stdout and stderr concurrently and return stdout. + + Uses explicit tasks so that if either reader raises (or this coroutine + is cancelled), the sibling reader is cancelled and awaited rather than + left running against a closing pipe. + """ + stdout_reader = asyncio.create_task(self._read_bounded(process.stdout)) + stderr_reader = asyncio.create_task(self._read_bounded(process.stderr)) + try: + stdout, _stderr = await asyncio.gather(stdout_reader, stderr_reader) + except BaseException: + for reader in (stdout_reader, stderr_reader): + reader.cancel() + await asyncio.gather(stdout_reader, stderr_reader, return_exceptions=True) + raise + return stdout + @staticmethod async def _read_bounded(stream: AsyncReadable) -> bytes: chunks = bytearray() diff --git a/tests/ui_and_conv/test_workspace_index.py b/tests/ui_and_conv/test_workspace_index.py index 92153c7d..2faa00d2 100644 --- a/tests/ui_and_conv/test_workspace_index.py +++ b/tests/ui_and_conv/test_workspace_index.py @@ -207,6 +207,48 @@ async def test_git_failure_uses_degraded_fallback() -> None: await lifecycle.aclose() +@pytest.mark.asyncio +async def test_read_streams_cancels_sibling_when_one_reader_fails() -> None: + # When one stream reader raises, the sibling reader must be cancelled and + # awaited rather than left draining a closing pipe. + host = _FakeHost() + index, lifecycle = _index(host, HostPath("/workspace")) + + class _Blocking: + def __init__(self) -> None: + self.started = asyncio.Event() + self.cancelled = False + + async def read(self, _n: int = -1) -> bytes: + self.started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.cancelled = True + raise + return b"" + + class _Raising: + def __init__(self, gate: asyncio.Event) -> None: + self._gate = gate + + async def read(self, _n: int = -1) -> bytes: + await self._gate.wait() # ensure the sibling is mid-read first + raise RuntimeError("stdout reader boom") + + blocking = _Blocking() + + class _Proc: + stdout = _Raising(blocking.started) + stderr = blocking + + with pytest.raises(RuntimeError, match="boom"): + await index._read_streams(cast(Any, _Proc())) + + assert blocking.cancelled is True + await lifecycle.aclose() + + @pytest.mark.asyncio async def test_results_are_truncated_at_one_thousand_entries() -> None: host = _FakeHost()