kimi-cli 0.44__py3-none-any.whl → 0.78__py3-none-any.whl

kimi_cli/ui/shell/prompt.py

@@ -1,59 +1,76 @@
+from __future__ import annotations
+
 import asyncio
 import base64
-import contextlib
 import getpass
 import json
 import os
 import re
-import sys
 import time
-from collections.abc import Callable
+from collections import deque
+from collections.abc import Callable, Iterable, Sequence
+from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum
 from hashlib import md5
 from io import BytesIO
 from pathlib import Path
-from typing import override
+from typing import Any, Literal, override
 
-from kosong.base.message import ContentPart, ImageURLPart, TextPart
+from kaos.path import KaosPath
 from PIL import Image, ImageGrab
 from prompt_toolkit import PromptSession
 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,
-    DummyCompleter,
     FuzzyCompleter,
     WordCompleter,
     merge_completers,
 )
 from prompt_toolkit.document import Document
-from prompt_toolkit.filters import Always, Never, has_completions
+from prompt_toolkit.filters import has_completions
 from prompt_toolkit.formatted_text import FormattedText
 from prompt_toolkit.history import InMemoryHistory
 from prompt_toolkit.key_binding import KeyBindings, KeyPressEvent
 from prompt_toolkit.patch_stdout import patch_stdout
+from prompt_toolkit.styles import Style
 from pydantic import BaseModel, ValidationError
 
+from kimi_cli.llm import ModelCapability
 from kimi_cli.share import get_share_dir
 from kimi_cli.soul import StatusSnapshot
-from kimi_cli.ui.shell.metacmd import get_meta_commands
+from kimi_cli.ui.shell.console import console
+from kimi_cli.utils.clipboard import is_clipboard_available
 from kimi_cli.utils.logging import logger
+from kimi_cli.utils.slashcmd import SlashCommand
 from kimi_cli.utils.string import random_string
+from kimi_cli.wire.types import ContentPart, ImageURLPart, TextPart
 
 PROMPT_SYMBOL = "✨"
+PROMPT_SYMBOL_SHELL = "$"
+PROMPT_SYMBOL_THINKING = "💫"
 
 
-class MetaCommandCompleter(Completer):
-    """A completer that:
-    - Shows one line per meta command in the form: "/name (alias1, alias2)"
+class SlashCommandCompleter(Completer):
+    """
+    A completer that:
+    - Shows one line per slash command in the form: "/name (alias1, alias2)"
     - Matches by primary name or any alias while inserting the canonical "/name"
     - Only activates when the current token starts with '/'
     """
 
+    def __init__(self, available_commands: Sequence[SlashCommand[Any]]) -> None:
+        super().__init__()
+        self._available_commands = available_commands
+
     @override
-    def get_completions(self, document, complete_event):
+    def get_completions(
+        self, document: Document, complete_event: CompleteEvent
+    ) -> Iterable[Completion]:
         text = document.text_before_cursor
 
         # Only autocomplete when the input buffer has no other content.
@@ -73,7 +90,7 @@ class MetaCommandCompleter(Completer):
         typed = token[1:]
         typed_lower = typed.lower()
 
-        for cmd in sorted(get_meta_commands(), key=lambda c: c.name):
+        for cmd in sorted(self._available_commands, key=lambda c: c.name):
             names = [cmd.name] + list(cmd.aliases)
             if typed == "" or any(n.lower().startswith(typed_lower) for n in names):
                 yield Completion(
@@ -84,7 +101,7 @@ class MetaCommandCompleter(Completer):
                 )
 
 
-class FileMentionCompleter(Completer):
+class LocalFileMentionCompleter(Completer):
     """Offer fuzzy `@` path completion by indexing workspace files."""
 
     _FRAGMENT_PATTERN = re.compile(r"[^\s@]+")
@@ -276,7 +293,7 @@ class FileMentionCompleter(Completer):
 
         if index > 0:
             prev = text[index - 1]
-            if prev.isalnum() or prev in FileMentionCompleter._TRIGGER_GUARDS:
+            if prev.isalnum() or prev in LocalFileMentionCompleter._TRIGGER_GUARDS:
                 return None
 
         fragment = text[index + 1 :]
@@ -298,7 +315,9 @@ class FileMentionCompleter(Completer):
             return False
 
     @override
-    def get_completions(self, document, complete_event):
+    def get_completions(
+        self, document: Document, complete_event: CompleteEvent
+    ) -> Iterable[Completion]:
         fragment = self._extract_fragment(document.text_before_cursor)
         if fragment is None:
             return
@@ -308,7 +327,26 @@ class FileMentionCompleter(Completer):
         mention_doc = Document(text=fragment, cursor_position=len(fragment))
         self._fragment_hint = fragment
         try:
-            yield from self._fuzzy.get_completions(mention_doc, complete_event)
+            # 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
+                # preserve original FuzzyCompleter's order in the same category
+                return (cat,)
+
+            candidates.sort(key=_rank)
+            yield from candidates
         finally:
             self._fragment_hint = None
 
@@ -359,7 +397,7 @@ class PromptMode(Enum):
     AGENT = "agent"
     SHELL = "shell"
 
-    def toggle(self) -> "PromptMode":
+    def toggle(self) -> PromptMode:
         return PromptMode.SHELL if self == PromptMode.AGENT else PromptMode.AGENT
 
     def __str__(self) -> str:
@@ -381,12 +419,49 @@ class UserInput(BaseModel):
 
 
 _REFRESH_INTERVAL = 1.0
-_toast_queue: asyncio.Queue[tuple[str, float]] = asyncio.Queue()
 
 
-def toast(message: str, duration: float = 5.0) -> None:
+@dataclass(slots=True)
+class _ToastEntry:
+    topic: str | None
+    """There can be only one toast of each non-None topic in the queue."""
+    message: str
+    duration: float
+
+
+_toast_queues: dict[Literal["left", "right"], deque[_ToastEntry]] = {
+    "left": deque(),
+    "right": deque(),
+}
+"""The queue of toasts to show, including the one currently being shown (the first one)."""
+
+
+def toast(
+    message: str,
+    duration: float = 5.0,
+    topic: str | None = None,
+    immediate: bool = False,
+    position: Literal["left", "right"] = "left",
+) -> None:
+    queue = _toast_queues[position]
     duration = max(duration, _REFRESH_INTERVAL)
-    _toast_queue.put_nowait((message, duration))
+    entry = _ToastEntry(topic=topic, message=message, duration=duration)
+    if topic is not None:
+        # Remove existing toasts with the same topic
+        for existing in list(queue):
+            if existing.topic == topic:
+                queue.remove(existing)
+    if immediate:
+        queue.appendleft(entry)
+    else:
+        queue.append(entry)
+
+
+def _current_toast(position: Literal["left", "right"] = "left") -> _ToastEntry | None:
+    queue = _toast_queues[position]
+    if not queue:
+        return None
+    return queue[0]
 
 
 _ATTACHMENT_PLACEHOLDER_RE = re.compile(
@@ -394,15 +469,36 @@ _ATTACHMENT_PLACEHOLDER_RE = re.compile(
 )
 
 
+def _sanitize_surrogates(text: str) -> str:
+    """Sanitize UTF-16 surrogate characters that cannot be encoded to UTF-8.
+
+    This is particularly common on Windows when copying text from applications
+    that use UTF-16 internally and don't properly convert surrogate pairs.
+    """
+    return text.encode("utf-8", errors="surrogatepass").decode("utf-8", errors="replace")
+
+
 class CustomPromptSession:
-    def __init__(self, status_provider: Callable[[], StatusSnapshot]):
+    def __init__(
+        self,
+        *,
+        status_provider: Callable[[], StatusSnapshot],
+        model_capabilities: set[ModelCapability],
+        model_name: str | None,
+        thinking: bool,
+        agent_mode_slash_commands: Sequence[SlashCommand[Any]],
+        shell_mode_slash_commands: Sequence[SlashCommand[Any]],
+    ) -> None:
         history_dir = get_share_dir() / "user-history"
         history_dir.mkdir(parents=True, exist_ok=True)
-        work_dir_id = md5(str(Path.cwd()).encode()).hexdigest()
+        work_dir_id = md5(str(KaosPath.cwd()).encode(encoding="utf-8")).hexdigest()
         self._history_file = (history_dir / work_dir_id).with_suffix(".jsonl")
         self._status_provider = status_provider
+        self._model_capabilities = model_capabilities
+        self._model_name = model_name
         self._last_history_content: str | None = None
         self._mode: PromptMode = PromptMode.AGENT
+        self._thinking = thinking
         self._attachment_parts: dict[str, ContentPart] = {}
         """Mapping from attachment id to ContentPart."""
 
@@ -418,17 +514,19 @@ class CustomPromptSession:
         # Build completers
         self._agent_mode_completer = merge_completers(
             [
-                MetaCommandCompleter(),
-                FileMentionCompleter(Path.cwd()),
+                SlashCommandCompleter(agent_mode_slash_commands),
+                # TODO(kaos): we need an async KaosFileMentionCompleter
+                LocalFileMentionCompleter(KaosPath.cwd().unsafe_to_local_path()),
             ],
             deduplicate=True,
         )
+        self._shell_mode_completer = SlashCommandCompleter(shell_mode_slash_commands)
 
         # Build key bindings
         _kb = KeyBindings()
 
         @_kb.add("enter", filter=has_completions)
-        def _accept_completion(event: KeyPressEvent) -> None:
+        def _(event: KeyPressEvent) -> None:
             """Accept the first completion when Enter is pressed and completions are shown."""
             buff = event.current_buffer
             if buff.complete_state and buff.complete_state.completions:
@@ -438,45 +536,66 @@ class CustomPromptSession:
                     completion = buff.complete_state.completions[0]
                 buff.apply_completion(completion)
 
-        @_kb.add("escape", "enter", eager=True)
-        @_kb.add("c-j", eager=True)
-        def _insert_newline(event: KeyPressEvent) -> None:
-            """Insert a newline when Alt-Enter or Ctrl-J is pressed."""
-            event.current_buffer.insert_text("\n")
-
         @_kb.add("c-x", eager=True)
-        def _switch_mode(event: KeyPressEvent) -> None:
+        def _(event: KeyPressEvent) -> None:
             self._mode = self._mode.toggle()
             # Apply mode-specific settings
             self._apply_mode(event)
             # Redraw UI
             event.app.invalidate()
 
-        @_kb.add("c-v", eager=True)
-        def _paste(event: KeyPressEvent) -> None:
-            if self._try_paste_image(event):
-                return
-            clipboard_data = event.app.clipboard.get_data()
-            event.current_buffer.paste_clipboard_data(clipboard_data)
+        @_kb.add("escape", "enter", eager=True)
+        @_kb.add("c-j", eager=True)
+        def _(event: KeyPressEvent) -> None:
+            """Insert a newline when Alt-Enter or Ctrl-J is pressed."""
+            event.current_buffer.insert_text("\n")
+
+        if is_clipboard_available():
+
+            @_kb.add("c-v", eager=True)
+            def _(event: KeyPressEvent) -> None:
+                if self._try_paste_image(event):
+                    return
+                clipboard_data = event.app.clipboard.get_data()
+                event.current_buffer.paste_clipboard_data(clipboard_data)
+
+            clipboard = PyperclipClipboard()
+        else:
+            clipboard = None
+
+        @_kb.add("c-_", eager=True)  # Ctrl-/ sends Ctrl-_ in most terminals
+        def _(event: KeyPressEvent) -> None:
+            """Show help by submitting /help command."""
+            buff = event.current_buffer
+            buff.text = "/help"
+            buff.validate_and_handle()
 
-        self._session = PromptSession(
+        self._session = PromptSession[str](
             message=self._render_message,
             # prompt_continuation=FormattedText([("fg:#4d4d4d", "... ")]),
             completer=self._agent_mode_completer,
             complete_while_typing=True,
             key_bindings=_kb,
-            clipboard=PyperclipClipboard(),
+            clipboard=clipboard,
             history=history,
             bottom_toolbar=self._render_bottom_toolbar,
+            style=Style.from_dict({"bottom-toolbar": "noreverse"}),
         )
 
-        self._status_refresh_task: asyncio.Task | None = None
-        self._current_toast: str | None = None
-        self._current_toast_duration: float = 0.0
+        # Allow completion to be triggered when the text is changed,
+        # such as when backspace is used to delete text.
+        @self._session.default_buffer.on_text_changed.add_handler
+        def _(buffer: Buffer) -> None:
+            if buffer.complete_while_typing():
+                buffer.start_completion()
+
+        self._status_refresh_task: asyncio.Task[None] | None = None
 
     def _render_message(self) -> FormattedText:
-        symbol = PROMPT_SYMBOL if self._mode == PromptMode.AGENT else "$"
-        return FormattedText([("bold", f"{getpass.getuser()}{symbol} ")])
+        symbol = PROMPT_SYMBOL if self._mode == PromptMode.AGENT else PROMPT_SYMBOL_SHELL
+        if self._mode == PromptMode.AGENT and self._thinking:
+            symbol = PROMPT_SYMBOL_THINKING
+        return FormattedText([("bold", f"{getpass.getuser()}@{KaosPath.cwd().name}{symbol} ")])
 
     def _apply_mode(self, event: KeyPressEvent | None = None) -> None:
         # Apply mode to the active buffer (not the PromptSession itself)
@@ -486,19 +605,13 @@ class CustomPromptSession:
             buff = None
 
         if self._mode == PromptMode.SHELL:
-            # Cancel any active completion menu
-            with contextlib.suppress(Exception):
-                if buff is not None:
-                    buff.cancel_completion()
             if buff is not None:
-                buff.completer = DummyCompleter()
-                buff.complete_while_typing = Never()
+                buff.completer = self._shell_mode_completer
         else:
             if buff is not None:
                 buff.completer = self._agent_mode_completer
-                buff.complete_while_typing = Always()
 
-    def __enter__(self) -> "CustomPromptSession":
+    def __enter__(self) -> CustomPromptSession:
         if self._status_refresh_task is not None and not self._status_refresh_task.done():
             return self
 
@@ -524,7 +637,7 @@ class CustomPromptSession:
         self._status_refresh_task = asyncio.create_task(_refresh(_REFRESH_INTERVAL))
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback) -> None:
+    def __exit__(self, *_) -> None:
         if self._status_refresh_task is not None and not self._status_refresh_task.done():
             self._status_refresh_task.cancel()
         self._status_refresh_task = None
@@ -534,7 +647,10 @@ class CustomPromptSession:
         """Try to paste an image from the clipboard. Return True if successful."""
         # Try get image from clipboard
         image = ImageGrab.grabclipboard()
-        if isinstance(image, list):
+        if image is None:
+            return False
+
+        if not isinstance(image, Image.Image):
             for item in image:
                 try:
                     with Image.open(item) as img:
@@ -543,9 +659,10 @@ class CustomPromptSession:
                 except Exception:
                     continue
             else:
-                image = None
+                return False
 
-        if image is None:
+        if "image_in" not in self._model_capabilities:
+            console.print("[yellow]Image input is not supported by the selected LLM model[/yellow]")
             return False
 
         attachment_id = f"{random_string(8)}.png"
@@ -570,9 +687,11 @@ class CustomPromptSession:
         return True
 
     async def prompt(self) -> UserInput:
-        with patch_stdout():
+        with patch_stdout(raw=True):
             command = str(await self._session.prompt_async()).strip()
             command = command.replace("\x00", "")  # just in case null bytes are somehow inserted
+            # Sanitize UTF-16 surrogates that may come from Windows clipboard
+            command = _sanitize_surrogates(command)
         self._append_history_entry(command)
 
         # Parse rich content parts
@@ -597,7 +716,11 @@ class CustomPromptSession:
         if remaining_command.strip():
             content.append(TextPart(text=remaining_command.strip()))
 
-        return UserInput(mode=self._mode, content=content, command=command)
+        return UserInput(
+            mode=self._mode,
+            content=content,
+            command=command,
+        )
 
     def _append_history_entry(self, text: str) -> None:
         entry = _HistoryEntry(content=text.strip())
@@ -632,146 +755,49 @@ class CustomPromptSession:
         columns -= len(now_text) + 2
 
         mode = str(self._mode).lower()
+        if self._mode == PromptMode.AGENT:
+            mode_details: list[str] = []
+            if self._model_name:
+                mode_details.append(self._model_name)
+            if self._thinking:
+                mode_details.append("thinking")
+            if mode_details:
+                mode += f" ({', '.join(mode_details)})"
+        status = self._status_provider()
+        if status.yolo_enabled:
+            fragments.extend([("bold fg:#ffff00", "yolo"), ("", " " * 2)])
+            columns -= len("yolo") + 2
         fragments.extend([("", f"{mode}"), ("", " " * 2)])
         columns -= len(mode) + 2
-
-        status = self._status_provider()
-        status_text = self._format_status(status)
-
-        if self._current_toast is not None:
-            fragments.extend([("", self._current_toast), ("", " " * 2)])
-            columns -= len(self._current_toast) + 2
-            self._current_toast_duration -= _REFRESH_INTERVAL
-            if self._current_toast_duration <= 0.0:
-                self._current_toast = None
+        right_text = self._render_right_span(status)
+
+        current_toast_left = _current_toast("left")
+        if current_toast_left is not None:
+            fragments.extend([("", current_toast_left.message), ("", " " * 2)])
+            columns -= len(current_toast_left.message) + 2
+            current_toast_left.duration -= _REFRESH_INTERVAL
+            if current_toast_left.duration <= 0.0:
+                _toast_queues["left"].popleft()
         else:
-            shortcuts = [
-                "ctrl-x: switch mode",
-                "ctrl-j: newline",
-                "ctrl-v: paste",
-                "ctrl-d: exit",
-            ]
-            for shortcut in shortcuts:
-                if columns - len(status_text) > len(shortcut) + 2:
-                    fragments.extend([("", shortcut), ("", " " * 2)])
-                    columns -= len(shortcut) + 2
-                else:
-                    break
+            shortcuts = "ctrl-x: toggle mode  ctrl-/: help"
+            if columns - len(right_text) > len(shortcuts) + 2:
+                fragments.extend([("", shortcuts), ("", " " * 2)])
+                columns -= len(shortcuts) + 2
 
-        if self._current_toast is None and not _toast_queue.empty():
-            self._current_toast, self._current_toast_duration = _toast_queue.get_nowait()
-
-        padding = max(1, columns - len(status_text))
+        padding = max(1, columns - len(right_text))
         fragments.append(("", " " * padding))
-        fragments.append(("", status_text))
+        fragments.append(("", right_text))
 
         return FormattedText(fragments)
 
     @staticmethod
-    def _format_status(status: StatusSnapshot) -> str:
-        bounded = max(0.0, min(status.context_usage, 1.0))
-        return f"context: {bounded:.1%}"
-
-
-def ensure_new_line() -> None:
-    """Ensure the next prompt starts at column 0 regardless of prior command output."""
-
-    if not sys.stdout.isatty() or not sys.stdin.isatty():
-        return
-
-    needs_break = True
-    if sys.platform == "win32":
-        column = _cursor_column_windows()
-        needs_break = column not in (None, 0)
-    else:
-        column = _cursor_column_unix()
-        needs_break = column not in (None, 1)
-
-    if needs_break:
-        _write_newline()
-
-
-def _cursor_column_unix() -> int | None:
-    assert sys.platform != "win32"
-
-    import select
-    import termios
-    import tty
-
-    _CURSOR_QUERY = "\x1b[6n"
-    _CURSOR_POSITION_RE = re.compile(r"\x1b\[(\d+);(\d+)R")
-
-    fd = sys.stdin.fileno()
-    oldterm = termios.tcgetattr(fd)
-
-    try:
-        tty.setcbreak(fd)
-        sys.stdout.write(_CURSOR_QUERY)
-        sys.stdout.flush()
-
-        response = ""
-        deadline = time.monotonic() + 0.2
-        while time.monotonic() < deadline:
-            timeout = max(0.01, deadline - time.monotonic())
-            ready, _, _ = select.select([sys.stdin], [], [], timeout)
-            if not ready:
-                continue
-            try:
-                chunk = os.read(fd, 32)
-            except OSError:
-                break
-            if not chunk:
-                break
-            response += chunk.decode(errors="ignore")
-            match = _CURSOR_POSITION_RE.search(response)
-            if match:
-                return int(match.group(2))
-    finally:
-        termios.tcsetattr(fd, termios.TCSADRAIN, oldterm)
-
-    return None
-
-
-def _cursor_column_windows() -> int | None:
-    assert sys.platform == "win32"
-
-    import ctypes
-    from ctypes import wintypes
-
-    kernel32 = ctypes.windll.kernel32
-    _STD_OUTPUT_HANDLE = -11  # Windows API constant for standard output handle
-    handle = kernel32.GetStdHandle(_STD_OUTPUT_HANDLE)
-    invalid_handle_value = ctypes.c_void_p(-1).value
-    if handle in (0, invalid_handle_value):
-        return None
-
-    class COORD(ctypes.Structure):
-        _fields_ = [("X", wintypes.SHORT), ("Y", wintypes.SHORT)]
-
-    class SMALL_RECT(ctypes.Structure):
-        _fields_ = [
-            ("Left", wintypes.SHORT),
-            ("Top", wintypes.SHORT),
-            ("Right", wintypes.SHORT),
-            ("Bottom", wintypes.SHORT),
-        ]
-
-    class CONSOLE_SCREEN_BUFFER_INFO(ctypes.Structure):
-        _fields_ = [
-            ("dwSize", COORD),
-            ("dwCursorPosition", COORD),
-            ("wAttributes", wintypes.WORD),
-            ("srWindow", SMALL_RECT),
-            ("dwMaximumWindowSize", COORD),
-        ]
-
-    csbi = CONSOLE_SCREEN_BUFFER_INFO()
-    if not kernel32.GetConsoleScreenBufferInfo(handle, ctypes.byref(csbi)):
-        return None
-
-    return int(csbi.dwCursorPosition.X)
-
-
-def _write_newline() -> None:
-    sys.stdout.write("\n")
-    sys.stdout.flush()
+    def _render_right_span(status: StatusSnapshot) -> str:
+        current_toast = _current_toast("right")
+        if current_toast is None:
+            bounded = max(0.0, min(status.context_usage, 1.0))
+            return f"context: {bounded:.1%}"
+
+        current_toast.duration -= _REFRESH_INTERVAL
+        if current_toast.duration <= 0.0:
+            _toast_queues["right"].popleft()
+        return current_toast.message