kimi-cli 0.41__py3-none-any.whl → 0.43__py3-none-any.whl

kimi_cli/ui/shell/prompt.py

@@ -1,19 +1,25 @@
 import asyncio
+import base64
 import contextlib
 import getpass
 import json
 import os
 import re
+import sys
 import time
 from collections.abc import Callable
 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 kosong.base.message import ContentPart, ImageURLPart, TextPart
+from PIL import Image, ImageGrab
 from prompt_toolkit import PromptSession
 from prompt_toolkit.application.current import get_app_or_none
+from prompt_toolkit.clipboard.pyperclip import PyperclipClipboard
 from prompt_toolkit.completion import (
     Completer,
     Completion,
@@ -34,6 +40,9 @@ 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.utils.logging import logger
+from kimi_cli.utils.string import random_string
+
+PROMPT_SYMBOL = "✨"
 
 
 class MetaCommandCompleter(Completer):
@@ -360,6 +369,9 @@ class PromptMode(Enum):
 class UserInput(BaseModel):
     mode: PromptMode
     command: str
+    """The plain text representation of the user input."""
+    content: list[ContentPart]
+    """The rich content parts."""
 
     def __str__(self) -> str:
         return self.command
@@ -377,6 +389,11 @@ def toast(message: str, duration: float = 5.0) -> None:
     _toast_queue.put_nowait((message, duration))
 
 
+_ATTACHMENT_PLACEHOLDER_RE = re.compile(
+    r"\[(?P<type>image):(?P<id>[a-zA-Z0-9_\-\.]+)(?:,(?P<width>\d+)x(?P<height>\d+))?\]"
+)
+
+
 class CustomPromptSession:
     def __init__(self, status_provider: Callable[[], StatusSnapshot]):
         history_dir = get_share_dir() / "user-history"
@@ -386,6 +403,8 @@ class CustomPromptSession:
         self._status_provider = status_provider
         self._last_history_content: str | None = None
         self._mode: PromptMode = PromptMode.AGENT
+        self._attachment_parts: dict[str, ContentPart] = {}
+        """Mapping from attachment id to ContentPart."""
 
         history_entries = _load_history_entries(self._history_file)
         history = InMemoryHistory()
@@ -419,20 +438,34 @@ class CustomPromptSession:
                     completion = buff.complete_state.completions[0]
                 buff.apply_completion(completion)
 
-        @_kb.add("c-k", eager=True)
-        def _toggle_mode(event: KeyPressEvent) -> None:
+        @_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:
             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)
+
         self._session = PromptSession(
             message=self._render_message,
-            prompt_continuation=FormattedText([("fg:#4d4d4d", "... ")]),
+            # prompt_continuation=FormattedText([("fg:#4d4d4d", "... ")]),
             completer=self._agent_mode_completer,
             complete_while_typing=True,
             key_bindings=_kb,
+            clipboard=PyperclipClipboard(),
             history=history,
             bottom_toolbar=self._render_bottom_toolbar,
         )
@@ -442,7 +475,7 @@ class CustomPromptSession:
         self._current_toast_duration: float = 0.0
 
     def _render_message(self) -> FormattedText:
-        symbol = "✨" if self._mode == PromptMode.AGENT else "$"
+        symbol = PROMPT_SYMBOL if self._mode == PromptMode.AGENT else "$"
         return FormattedText([("bold", f"{getpass.getuser()}{symbol} ")])
 
     def _apply_mode(self, event: KeyPressEvent | None = None) -> None:
@@ -495,12 +528,75 @@ class CustomPromptSession:
         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
+        self._attachment_parts.clear()
+
+    def _try_paste_image(self, event: KeyPressEvent) -> bool:
+        """Try to paste an image from the clipboard. Return True if successful."""
+        # Try get image from clipboard
+        image = ImageGrab.grabclipboard()
+        if isinstance(image, list):
+            for item in image:
+                try:
+                    with Image.open(item) as img:
+                        image = img.copy()
+                    break
+                except Exception:
+                    continue
+            else:
+                image = None
+
+        if image is None:
+            return False
+
+        attachment_id = f"{random_string(8)}.png"
+        png_bytes = BytesIO()
+        image.save(png_bytes, format="PNG")
+        png_base64 = base64.b64encode(png_bytes.getvalue()).decode("ascii")
+        image_part = ImageURLPart(
+            image_url=ImageURLPart.ImageURL(
+                url=f"data:image/png;base64,{png_base64}", id=attachment_id
+            )
+        )
+        self._attachment_parts[attachment_id] = image_part
+        logger.debug(
+            "Pasted image from clipboard: {attachment_id}, {image_size}",
+            attachment_id=attachment_id,
+            image_size=image.size,
+        )
+
+        placeholder = f"[image:{attachment_id},{image.width}x{image.height}]"
+        event.current_buffer.insert_text(placeholder)
+        event.app.invalidate()
+        return True
 
     async def prompt(self) -> UserInput:
         with patch_stdout():
             command = str(await self._session.prompt_async()).strip()
         self._append_history_entry(command)
-        return UserInput(mode=self._mode, command=command)
+
+        # Parse rich content parts
+        content: list[ContentPart] = []
+        remaining_command = command
+        while match := _ATTACHMENT_PLACEHOLDER_RE.search(remaining_command):
+            start, end = match.span()
+            if start > 0:
+                content.append(TextPart(text=remaining_command[:start]))
+            attachment_id = match.group("id")
+            part = self._attachment_parts.get(attachment_id)
+            if part is not None:
+                content.append(part)
+            else:
+                logger.warning(
+                    "Attachment placeholder found but no matching attachment part: {placeholder}",
+                    placeholder=match.group(0),
+                )
+                content.append(TextPart(text=match.group(0)))
+            remaining_command = remaining_command[end:]
+
+        if remaining_command.strip():
+            content.append(TextPart(text=remaining_command.strip()))
+
+        return UserInput(mode=self._mode, content=content, command=command)
 
     def _append_history_entry(self, text: str) -> None:
         entry = _HistoryEntry(content=text.strip())
@@ -549,7 +645,9 @@ class CustomPromptSession:
                 self._current_toast = None
         else:
             shortcuts = [
-                "ctrl-k: toggle mode",
+                "ctrl-x: switch mode",
+                "ctrl-j: newline",
+                "ctrl-v: paste",
                 "ctrl-d: exit",
             ]
             for shortcut in shortcuts:
@@ -572,3 +670,107 @@ class CustomPromptSession:
     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()

kimi_cli/ui/shell/replay.py

@@ -0,0 +1,104 @@
+import asyncio
+import contextlib
+import getpass
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+from kosong.base.message import Message, TextPart
+from kosong.tooling import ToolError, ToolOk
+
+from kimi_cli.soul import StatusSnapshot
+from kimi_cli.ui.shell.console import console
+from kimi_cli.ui.shell.prompt import PROMPT_SYMBOL
+from kimi_cli.ui.shell.visualize import visualize
+from kimi_cli.utils.message import message_extract_text, message_stringify
+from kimi_cli.wire import Wire
+from kimi_cli.wire.message import ContentPart, StepBegin, ToolCall, ToolResult
+
+MAX_REPLAY_RUNS = 5
+
+type _ReplayEvent = StepBegin | ToolCall | ContentPart | ToolResult
+
+
+@dataclass(slots=True)
+class _ReplayRun:
+    user_message: Message
+    events: list[_ReplayEvent]
+    n_steps: int = 0
+
+
+async def replay_recent_history(history: Sequence[Message]) -> None:
+    """
+    Replay the most recent user-initiated runs from the provided message history.
+    """
+    start_idx = _find_replay_start(history)
+    if start_idx is None:
+        return
+
+    runs = _build_replay_runs(history[start_idx:])
+    if not runs:
+        return
+
+    for run in runs:
+        wire = Wire()
+        console.print(f"{getpass.getuser()}{PROMPT_SYMBOL} {message_stringify(run.user_message)}")
+        ui_task = asyncio.create_task(
+            visualize(wire.ui_side, initial_status=StatusSnapshot(context_usage=0.0))
+        )
+        for event in run.events:
+            wire.soul_side.send(event)
+            await asyncio.sleep(0)  # yield to UI loop
+        wire.shutdown()
+        with contextlib.suppress(asyncio.QueueShutDown):
+            await ui_task
+
+
+def _is_user_message(message: Message) -> bool:
+    # FIXME: should consider non-text tool call results which are sent as user messages
+    if message.role != "user":
+        return False
+    return not message_extract_text(message).startswith("<system>CHECKPOINT")
+
+
+def _find_replay_start(history: Sequence[Message]) -> int | None:
+    indices = [idx for idx, message in enumerate(history) if _is_user_message(message)]
+    if not indices:
+        return None
+    # only replay last MAX_REPLAY_RUNS messages
+    return indices[max(0, len(indices) - MAX_REPLAY_RUNS)]
+
+
+def _build_replay_runs(history: Sequence[Message]) -> list[_ReplayRun]:
+    runs: list[_ReplayRun] = []
+    current_run: _ReplayRun | None = None
+    for message in history:
+        if _is_user_message(message):
+            # start a new run
+            if current_run is not None:
+                runs.append(current_run)
+            current_run = _ReplayRun(user_message=message, events=[])
+        elif message.role == "assistant":
+            if current_run is None:
+                continue
+            current_run.n_steps += 1
+            current_run.events.append(StepBegin(n=current_run.n_steps))
+            if isinstance(message.content, str):
+                current_run.events.append(TextPart(text=message.content))
+            else:
+                current_run.events.extend(message.content)
+            current_run.events.extend(message.tool_calls or [])
+        elif message.role == "tool":
+            if current_run is None:
+                continue
+            assert message.tool_call_id is not None
+            if isinstance(message.content, list) and any(
+                isinstance(part, TextPart) and part.text.startswith("<system>ERROR")
+                for part in message.content
+            ):
+                result = ToolError(message="", output="", brief="")
+            else:
+                result = ToolOk(output=message.content)
+            current_run.events.append(ToolResult(tool_call_id=message.tool_call_id, result=result))
+    if current_run is not None:
+        runs.append(current_run)
+    return runs

kimi_cli/ui/shell/visualize.py

@@ -8,7 +8,6 @@ from kimi_cli.soul import StatusSnapshot
 from kimi_cli.ui.shell.console import console
 from kimi_cli.ui.shell.keyboard import listen_for_keyboard
 from kimi_cli.ui.shell.liveview import StepLiveView, StepLiveViewWithMarkdown
-from kimi_cli.utils.logging import logger
 from kimi_cli.wire import WireUISide
 from kimi_cli.wire.message import (
     ApprovalRequest,
@@ -46,69 +45,67 @@ async def visualize(
 ):
     """
     A loop to consume agent events and visualize the agent behavior.
-    This loop never raise any exception except asyncio.CancelledError.
 
     Args:
         wire: Communication channel with the agent
         initial_status: Initial status snapshot
         cancel_event: Event that can be set (e.g., by ESC key) to cancel the run
     """
+
     latest_status = initial_status
-    try:
-        # expect a StepBegin
-        assert isinstance(await wire.receive(), StepBegin)
-
-        while True:
-            # TODO: Maybe we can always have a StepLiveView here.
-            #       No need to recreate for each step.
-            with StepLiveViewWithMarkdown(latest_status, cancel_event) as step:
-                async with _keyboard_listener(step):
-                    # spin the moon at the beginning of each step
-                    with console.status("", spinner="moon"):
-                        msg = await wire.receive()
 
-                    if isinstance(msg, CompactionBegin):
-                        with console.status("[cyan]Compacting...[/cyan]"):
-                            msg = await wire.receive()
-                        if isinstance(msg, StepInterrupted):
-                            break
-                        assert isinstance(msg, CompactionEnd)
-                        continue
-
-                    # visualization loop for one step
-                    while True:
-                        match msg:
-                            case TextPart(text=text):
-                                step.append_text(text)
-                            case ContentPart():
-                                # TODO: support more content parts
-                                step.append_text(f"[{msg.__class__.__name__}]")
-                            case ToolCall():
-                                step.append_tool_call(msg)
-                            case ToolCallPart():
-                                step.append_tool_call_part(msg)
-                            case ToolResult():
-                                step.append_tool_result(msg)
-                            case ApprovalRequest():
-                                step.request_approval(msg)
-                            case StatusUpdate(status=status):
-                                latest_status = status
-                                step.update_status(latest_status)
-                            case _:
-                                break  # break the step loop
-                        msg = await wire.receive()
+    # expect a StepBegin
+    assert isinstance(await wire.receive(), StepBegin)
+
+    while True:
+        # TODO: Maybe we can always have a StepLiveView here.
+        #       No need to recreate for each step.
+        with StepLiveViewWithMarkdown(latest_status, cancel_event) as step:
+            async with _keyboard_listener(step):
+                # spin the moon at the beginning of each step
+                with console.status("", spinner="moon"):
+                    msg = await wire.receive()
 
-                    # cleanup the step live view
+                if isinstance(msg, CompactionBegin):
+                    with console.status("[cyan]Compacting...[/cyan]"):
+                        msg = await wire.receive()
                     if isinstance(msg, StepInterrupted):
-                        step.interrupt()
-                    else:
-                        step.finish()
-
-            if isinstance(msg, StepInterrupted):
-                # for StepInterrupted, the visualization loop should end immediately
-                break
-
-            assert isinstance(msg, StepBegin), "expect a StepBegin"
-            # start a new step
-    except asyncio.QueueShutDown:
-        logger.debug("Visualization loop shutting down")
+                        break
+                    assert isinstance(msg, CompactionEnd)
+                    continue
+
+                # visualization loop for one step
+                while True:
+                    match msg:
+                        case TextPart(text=text):
+                            step.append_text(text)
+                        case ContentPart():
+                            # TODO: support more content parts
+                            step.append_text(f"[{msg.__class__.__name__}]")
+                        case ToolCall():
+                            step.append_tool_call(msg)
+                        case ToolCallPart():
+                            step.append_tool_call_part(msg)
+                        case ToolResult():
+                            step.append_tool_result(msg)
+                        case ApprovalRequest():
+                            step.request_approval(msg)
+                        case StatusUpdate(status=status):
+                            latest_status = status
+                            step.update_status(latest_status)
+                        case _:
+                            break  # break the step loop
+                    msg = await wire.receive()
+
+                # cleanup the step live view
+                if isinstance(msg, StepInterrupted):
+                    step.interrupt()
+                else:
+                    step.finish()
+
+        if isinstance(msg, StepInterrupted):
+            # for StepInterrupted, the visualization loop should end immediately
+            break
+
+        assert isinstance(msg, StepBegin), "expect a StepBegin"
+        # start a new step

kimi_cli/utils/message.py

@@ -6,3 +6,17 @@ def message_extract_text(message: Message) -> str:
     if isinstance(message.content, str):
         return message.content
     return "\n".join(part.text for part in message.content if isinstance(part, TextPart))
+
+
+def message_stringify(message: Message) -> str:
+    """Get a string representation of a message."""
+    parts: list[str] = []
+    if isinstance(message.content, str):
+        parts.append(message.content)
+    else:
+        for part in message.content:
+            if isinstance(part, TextPart):
+                parts.append(part.text)
+            else:
+                parts.append(f"[{part.type}]")
+    return "".join(parts)

kimi_cli/utils/signals.py

@@ -0,0 +1,41 @@
+import asyncio
+import contextlib
+import signal
+from collections.abc import Callable
+
+
+def install_sigint_handler(
+    loop: asyncio.AbstractEventLoop, handler: Callable[[], None]
+) -> Callable[[], None]:
+    """
+    Install a SIGINT handler that works on Unix and Windows.
+
+    On Unix event loops, prefer `loop.add_signal_handler`.
+    On Windows (or other platforms) where it is not implemented, fall back to
+    `signal.signal`. The fallback cannot be removed from the loop, but we
+    restore the previous handler on uninstall.
+
+    Returns:
+        A function that removes the installed handler. It is guaranteed that
+        no exceptions are raised when calling the returned function.
+    """
+
+    try:
+        loop.add_signal_handler(signal.SIGINT, handler)
+
+        def remove() -> None:
+            with contextlib.suppress(RuntimeError):
+                loop.remove_signal_handler(signal.SIGINT)
+
+        return remove
+    except RuntimeError:
+        # Windows ProactorEventLoop and some environments do not support
+        # add_signal_handler. Use synchronous signal handling as a fallback.
+        previous = signal.getsignal(signal.SIGINT)
+        signal.signal(signal.SIGINT, lambda signum, frame: handler())
+
+        def remove() -> None:
+            with contextlib.suppress(RuntimeError):
+                signal.signal(signal.SIGINT, previous)
+
+        return remove

kimi_cli/utils/string.py

@@ -1,4 +1,6 @@
+import random
 import re
+import string
 
 _NEWLINE_RE = re.compile(r"[\r\n]+")
 
@@ -10,3 +12,9 @@ def shorten_middle(text: str, width: int, remove_newline: bool = True) -> str:
     if remove_newline:
         text = _NEWLINE_RE.sub(" ", text)
     return text[: width // 2] + "..." + text[-width // 2 :]
+
+
+def random_string(length: int = 8) -> str:
+    """Generate a random string of fixed length."""
+    letters = string.ascii_lowercase
+    return "".join(random.choice(letters) for _ in range(length))

kimi_cli/wire/__init__.py

@@ -25,6 +25,7 @@ class Wire:
         return self._ui_side
 
     def shutdown(self) -> None:
+        logger.debug("Shutting down wire")
         self._queue.shutdown()
 
 
@@ -55,3 +56,15 @@ class WireUISide:
         if not isinstance(msg, ContentPart | ToolCallPart):
             logger.debug("Receiving wire message: {msg}", msg=msg)
         return msg
+
+    def receive_nowait(self) -> WireMessage | None:
+        """
+        Try receive a message without waiting. If no message is available, return None.
+        """
+        try:
+            msg = self._queue.get_nowait()
+        except asyncio.QueueEmpty:
+            return None
+        if not isinstance(msg, ContentPart | ToolCallPart):
+            logger.debug("Receiving wire message: {msg}", msg=msg)
+        return msg