feat(main): replace readline with prompt_toolkit for file completions

lzwjava · lzwjava · commit bf0b5c8f12a9 · 2026-03-15T14:27:52.000+08:00
diff --git a/iclaw/main.py b/iclaw/main.py
@@ -1,29 +1,14 @@
 #!/usr/bin/env python3
+import glob
 import json
+import os
+import re
 import sys
 import time
 from pathlib import Path
 
-try:
-    import readline
-
-    COMMANDS = [
-        "/model_provider",
-        "/model",
-        "/search_provider",
-        "/copy",
-        "/help",
-        ".exit",
-    ]
-
-    def completer(text, state):
-        matches = [c for c in COMMANDS if c.startswith(text)]
-        return matches[state] if state < len(matches) else None
-
-    readline.set_completer(completer)
-    readline.parse_and_bind("tab: complete")
-except ImportError:
-    pass
+from prompt_toolkit import PromptSession
+from prompt_toolkit.completion import Completer, Completion
 
 from iclaw.github_api import chat, get_copilot_token
 from iclaw.web_search import web_search
@@ -33,6 +18,58 @@ def completer(text, state):
 from iclaw.commands.search_provider import handle_search_provider_command
 from iclaw.commands.utils import handle_copy_command
 
+COMMANDS = [
+    "/model_provider",
+    "/model",
+    "/search_provider",
+    "/copy",
+    "/help",
+    ".exit",
+]
+
+
+class IclawCompleter(Completer):
+    """Handles both / command completion and @ file mention completion."""
+
+    def get_completions(self, document, complete_event):
+        text = document.text_before_cursor
+
+        # @ file mention: find the last @ not preceded by a non-space char
+        at_pos = text.rfind("@")
+        if at_pos != -1:
+            prefix = text[at_pos + 1 :]
+            if " " not in prefix:
+                pattern = f"{prefix}*"
+                matches = glob.glob(pattern) + glob.glob(
+                    os.path.join("**", pattern), recursive=True
+                )
+                seen = set()
+                for path in sorted(matches)[:20]:
+                    if path in seen:
+                        continue
+                    seen.add(path)
+                    meta = "dir" if os.path.isdir(path) else "file"
+                    yield Completion(
+                        path,
+                        start_position=-len(prefix),
+                        display=path,
+                        display_meta=meta,
+                    )
+                return
+
+        # / command completion at start of input
+        stripped = text.lstrip()
+        if stripped.startswith("/") or stripped == ".":
+            word = stripped
+            for cmd in COMMANDS:
+                if cmd.startswith(word):
+                    yield Completion(
+                        cmd,
+                        start_position=-len(stripped),
+                        display=cmd,
+                    )
+
+
 COMMANDS_HELP = [
     ("/model_provider", "Select and authenticate with the model provider"),
     ("/model", "Select specific model from your provider"),
@@ -120,6 +157,24 @@ def load_github_token():
         return None
 
 
+def resolve_at_mentions(text):
+    """Extract @file references, return augmented text with file contents prepended."""
+    mentions = re.findall(r"@(\S+)", text)
+    if not mentions:
+        return text
+    parts = []
+    for path in mentions:
+        if os.path.isfile(path):
+            try:
+                contents = Path(path).read_text()
+                parts.append(f'<file path="{path}">\n{contents}\n</file>')
+            except OSError:
+                pass
+    if parts:
+        return "\n".join(parts) + "\n\n" + text
+    return text
+
+
 def main():
     github_token = load_github_token()
     copilot_token = None
@@ -145,9 +200,11 @@ def main():
         print(f"  {cmd:<20} {desc}")
     print()
 
+    session = PromptSession(completer=IclawCompleter(), complete_while_typing=True)
+
     while True:
         try:
-            user_input = input("> ").strip()
+            user_input = session.prompt("> ").strip()
         except (EOFError, KeyboardInterrupt):
             print("\nGoodbye!")
             break
@@ -190,7 +247,9 @@ def main():
                 copilot_token = get_copilot_token(github_token)
                 token_expiry = time.monotonic() + TOKEN_REFRESH_INTERVAL
 
-            messages.append({"role": "user", "content": user_input})
+            messages.append(
+                {"role": "user", "content": resolve_at_mentions(user_input)}
+            )
             response_message = chat(messages, copilot_token, current_model, tools=TOOLS)
 
             while response_message.get("tool_calls"):
diff --git a/pyproject.toml b/pyproject.toml
@@ -8,7 +8,7 @@ version = "0.1.0"
 description = "Interactive CLI REPL for chatting with GitHub Copilot"
 readme = "README.md"
 requires-python = ">=3.8"
-dependencies = ["requests>=2.32.0", "pyperclip>=1.8.0", "beautifulsoup4>=4.12.0", "readability-lxml>=0.8.1", "lxml>=5.0.0", "tavily-python>=0.3.0"]
+dependencies = ["requests>=2.32.0", "pyperclip>=1.8.0", "beautifulsoup4>=4.12.0", "readability-lxml>=0.8.1", "lxml>=5.0.0", "tavily-python>=0.3.0", "prompt-toolkit>=3.0.0"]
 
 [project.scripts]
 iclaw = "iclaw.main:main"
diff --git a/tests/test_at_mention.py b/tests/test_at_mention.py
@@ -0,0 +1,165 @@
+"""Integration tests for @ file mention and IclawCompleter."""
+
+import os
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+from prompt_toolkit.completion import CompleteEvent
+from prompt_toolkit.document import Document
+
+from iclaw.main import IclawCompleter, resolve_at_mentions
+
+
+class TestResolveAtMentions(unittest.TestCase):
+    def setUp(self):
+        self.tmpdir = tempfile.mkdtemp()
+        self.file1 = os.path.join(self.tmpdir, "hello.txt")
+        Path(self.file1).write_text("hello world")
+        self.file2 = os.path.join(self.tmpdir, "other.py")
+        Path(self.file2).write_text("print('hi')")
+
+    def test_no_mentions_returns_original(self):
+        text = "just a plain message"
+        self.assertEqual(resolve_at_mentions(text), text)
+
+    def test_mention_nonexistent_file_returns_original(self):
+        text = "look at @nonexistent_file.txt"
+        self.assertEqual(resolve_at_mentions(text), text)
+
+    def test_mention_existing_file_prepends_contents(self):
+        text = f"explain @{self.file1}"
+        result = resolve_at_mentions(text)
+        self.assertIn("hello world", result)
+        self.assertIn(f'<file path="{self.file1}">', result)
+        self.assertIn(text, result)
+
+    def test_mention_multiple_files(self):
+        text = f"compare @{self.file1} and @{self.file2}"
+        result = resolve_at_mentions(text)
+        self.assertIn("hello world", result)
+        self.assertIn("print('hi')", result)
+        self.assertIn(text, result)
+
+    def test_file_contents_come_before_message(self):
+        text = f"explain @{self.file1}"
+        result = resolve_at_mentions(text)
+        file_tag_pos = result.index("<file")
+        msg_pos = result.index(text)
+        self.assertLess(file_tag_pos, msg_pos)
+
+    def test_mention_directory_ignored(self):
+        text = f"look at @{self.tmpdir}"
+        # directories are not files, so no injection
+        result = resolve_at_mentions(text)
+        self.assertEqual(result, text)
+
+    def test_unreadable_file_skipped(self):
+        text = f"explain @{self.file1}"
+        with patch("iclaw.main.Path.read_text", side_effect=OSError("perm denied")):
+            result = resolve_at_mentions(text)
+        self.assertEqual(result, text)
+
+
+class TestIclawCompleter(unittest.TestCase):
+    def setUp(self):
+        self.completer = IclawCompleter()
+        self.tmpdir = tempfile.mkdtemp()
+        self.orig_cwd = os.getcwd()
+        os.chdir(self.tmpdir)
+        # Create test files in the temp dir
+        Path("alpha.py").write_text("")
+        Path("alpha_test.py").write_text("")
+        Path("beta.txt").write_text("")
+        os.makedirs("subdir", exist_ok=True)
+        Path("subdir/gamma.py").write_text("")
+
+    def tearDown(self):
+        os.chdir(self.orig_cwd)
+
+    def _completions(self, text):
+        doc = Document(text)
+        return list(self.completer.get_completions(doc, CompleteEvent()))
+
+    # --- @ file mention ---
+
+    def test_at_with_no_prefix_returns_files(self):
+        completions = self._completions("@")
+        paths = [c.text for c in completions]
+        self.assertIn("alpha.py", paths)
+        self.assertIn("beta.txt", paths)
+
+    def test_at_with_partial_prefix_filters(self):
+        completions = self._completions("@alph")
+        paths = [c.text for c in completions]
+        self.assertTrue(all("alpha" in p for p in paths))
+        self.assertNotIn("beta.txt", paths)
+
+    def test_at_completion_replaces_prefix(self):
+        completions = self._completions("@alph")
+        for c in completions:
+            # start_position should be negative len of "alph"
+            self.assertEqual(c.start_position, -len("alph"))
+
+    def test_at_meta_shows_file_or_dir(self):
+        completions = self._completions("@")
+        meta_map = {c.text: c.display_meta for c in completions}
+        # display_meta is a FormattedText; convert to string for assertion
+        for path, meta in meta_map.items():
+            expected = "dir" if os.path.isdir(path) else "file"
+            self.assertIn(expected, str(meta))
+
+    def test_at_mid_sentence(self):
+        completions = self._completions("review the file @alph")
+        paths = [c.text for c in completions]
+        self.assertTrue(any("alpha" in p for p in paths))
+
+    def test_at_with_space_after_at_returns_nothing(self):
+        # "@foo bar" — space in prefix means we're past the mention word
+        completions = self._completions("@alpha.py bar")
+        self.assertEqual(completions, [])
+
+    def test_at_limits_to_20_results(self):
+        # Create 25 files
+        for i in range(25):
+            Path(f"zfile{i:02d}.py").write_text("")
+        completions = self._completions("@zfile")
+        self.assertLessEqual(len(completions), 20)
+
+    # --- / command completion ---
+
+    def test_slash_alone_returns_all_commands(self):
+        completions = self._completions("/")
+        texts = [c.text for c in completions]
+        # "/" prefix only matches slash-prefixed commands, not ".exit"
+        for cmd in ["/model_provider", "/model", "/search_provider", "/copy", "/help"]:
+            self.assertIn(cmd, texts)
+        self.assertNotIn(".exit", texts)
+
+    def test_slash_partial_filters_commands(self):
+        completions = self._completions("/mod")
+        texts = [c.text for c in completions]
+        self.assertIn("/model", texts)
+        self.assertIn("/model_provider", texts)
+        self.assertNotIn("/copy", texts)
+        self.assertNotIn("/help", texts)
+
+    def test_dot_exit_completion(self):
+        completions = self._completions(".")
+        texts = [c.text for c in completions]
+        self.assertIn(".exit", texts)
+
+    def test_no_trigger_returns_nothing(self):
+        completions = self._completions("hello world")
+        self.assertEqual(completions, [])
+
+    def test_at_takes_priority_over_slash_in_same_input(self):
+        # Input has both / at start and @ later — @ should win
+        completions = self._completions("/help @alph")
+        paths = [c.text for c in completions]
+        self.assertTrue(any("alpha" in p for p in paths))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_main.py b/tests/test_main.py
