kimi-cli 0.35__py3-none-any.whl

kimi_cli/soul/message.py

@@ -0,0 +1,76 @@
+from kosong.base.message import ContentPart, Message, TextPart
+from kosong.tooling import ToolError, ToolOk, ToolResult
+from kosong.tooling.error import ToolRuntimeError
+
+
+def system(message: str) -> ContentPart:
+    return TextPart(text=f"<system>{message}</system>")
+
+
+def tool_result_to_messages(tool_result: ToolResult) -> list[Message]:
+    """Convert a tool result to a list of messages."""
+    if isinstance(tool_result.result, ToolError):
+        assert tool_result.result.message, "ToolError should have a message"
+        message = tool_result.result.message
+        if isinstance(tool_result.result, ToolRuntimeError):
+            message += "\nThis is an unexpected error and the tool is probably not working."
+        content = [system(message)]
+        if tool_result.result.output:
+            content.append(TextPart(text=tool_result.result.output))
+        return [
+            Message(
+                role="tool",
+                content=content,
+                tool_call_id=tool_result.tool_call_id,
+            )
+        ]
+
+    content = tool_ok_to_message_content(tool_result.result)
+    text_parts = []
+    non_text_parts = []
+    for part in content:
+        if isinstance(part, TextPart):
+            text_parts.append(part)
+        else:
+            non_text_parts.append(part)
+
+    if not non_text_parts:
+        return [
+            Message(
+                role="tool",
+                content=text_parts,
+                tool_call_id=tool_result.tool_call_id,
+            )
+        ]
+
+    text_parts.append(
+        system(
+            "Tool output contains non-text parts. Non-text parts are sent as a user message below."
+        )
+    )
+    return [
+        Message(
+            role="tool",
+            content=text_parts,
+            tool_call_id=tool_result.tool_call_id,
+        ),
+        Message(role="user", content=non_text_parts),
+    ]
+
+
+def tool_ok_to_message_content(result: ToolOk) -> list[ContentPart]:
+    """Convert a tool return value to a list of message content parts."""
+    content = []
+    if result.message:
+        content.append(system(result.message))
+    match output := result.output:
+        case str(text):
+            if text:
+                content.append(TextPart(text=text))
+        case ContentPart():
+            content.append(output)
+        case _:
+            content.extend(list(output))
+    if not content:
+        content.append(system("Tool output is empty."))
+    return content

kimi_cli/soul/toolset.py

@@ -0,0 +1,25 @@
+from contextvars import ContextVar
+from typing import override
+
+from kosong.base.message import ToolCall
+from kosong.tooling import HandleResult, SimpleToolset
+
+current_tool_call = ContextVar[ToolCall | None]("current_tool_call", default=None)
+
+
+def get_current_tool_call_or_none() -> ToolCall | None:
+    """
+    Get the current tool call or None.
+    Expect to be not None when called from a `__call__` method of a tool.
+    """
+    return current_tool_call.get()
+
+
+class CustomToolset(SimpleToolset):
+    @override
+    def handle(self, tool_call: ToolCall) -> HandleResult:
+        token = current_tool_call.set(tool_call)
+        try:
+            return super().handle(tool_call)
+        finally:
+            current_tool_call.reset(token)

kimi_cli/soul/wire.py

@@ -0,0 +1,101 @@
+import asyncio
+import uuid
+from contextvars import ContextVar
+from enum import Enum
+from typing import NamedTuple
+
+from kosong.base.message import ContentPart, ToolCall, ToolCallPart
+from kosong.tooling import ToolResult
+
+from kimi_cli.soul import StatusSnapshot
+from kimi_cli.utils.logging import logger
+
+
+class StepBegin(NamedTuple):
+    n: int
+
+
+class StepInterrupted(NamedTuple):
+    pass
+
+
+class StatusUpdate(NamedTuple):
+    status: StatusSnapshot
+
+
+type ControlFlowEvent = StepBegin | StepInterrupted | StatusUpdate
+type Event = ControlFlowEvent | ContentPart | ToolCall | ToolCallPart | ToolResult
+
+
+class ApprovalResponse(Enum):
+    APPROVE = "approve"
+    APPROVE_FOR_SESSION = "approve_for_session"
+    REJECT = "reject"
+
+
+class ApprovalRequest:
+    def __init__(self, tool_call_id: str, action: str, description: str):
+        self.id = str(uuid.uuid4())
+        self.tool_call_id = tool_call_id
+        self.action = action
+        self.description = description
+        self._future = asyncio.Future[ApprovalResponse]()
+
+    def __repr__(self) -> str:
+        return (
+            f"ApprovalRequest(id={self.id}, tool_call_id={self.tool_call_id}, "
+            f"action={self.action}, description={self.description})"
+        )
+
+    async def wait(self) -> ApprovalResponse:
+        """
+        Wait for the request to be resolved or cancelled.
+
+        Returns:
+            ApprovalResponse: The response to the approval request.
+        """
+        return await self._future
+
+    def resolve(self, response: ApprovalResponse) -> None:
+        """
+        Resolve the approval request with the given response.
+        This will cause the `wait()` method to return the response.
+        """
+        self._future.set_result(response)
+
+
+type WireMessage = Event | ApprovalRequest
+
+
+class Wire:
+    """
+    A channel for communication between the soul and the UI.
+    """
+
+    def __init__(self):
+        self._queue = asyncio.Queue[WireMessage]()
+
+    def send(self, msg: WireMessage) -> None:
+        if not isinstance(msg, ContentPart | ToolCallPart):
+            logger.debug("Sending wire message: {msg}", msg=msg)
+        self._queue.put_nowait(msg)
+
+    async def receive(self) -> WireMessage:
+        msg = await self._queue.get()
+        if not isinstance(msg, ContentPart | ToolCallPart):
+            logger.debug("Receiving wire message: {msg}", msg=msg)
+        return msg
+
+    def shutdown(self) -> None:
+        self._queue.shutdown()
+
+
+current_wire = ContextVar[Wire | None]("current_wire", default=None)
+
+
+def get_wire_or_none() -> Wire | None:
+    """
+    Get the current wire or None.
+    Expect to be not None when called from anywhere in the agent loop.
+    """
+    return current_wire.get()

kimi_cli/tools/__init__.py

@@ -0,0 +1,85 @@
+import json
+from pathlib import Path
+
+import streamingjson
+from kosong.utils.typing import JsonType
+
+from kimi_cli.utils.string import shorten_middle
+
+
+def extract_subtitle(lexer: streamingjson.Lexer, tool_name: str) -> str | None:
+    try:
+        curr_args: JsonType = json.loads(lexer.complete_json())
+    except json.JSONDecodeError:
+        return None
+    if not curr_args:
+        return None
+    subtitle: str = ""
+    match tool_name:
+        case "Task":
+            if not isinstance(curr_args, dict) or not curr_args.get("description"):
+                return None
+            subtitle = str(curr_args["description"])
+        case "SendDMail":
+            return "El Psy Kongroo"
+        case "Think":
+            if not isinstance(curr_args, dict) or not curr_args.get("thought"):
+                return None
+            subtitle = str(curr_args["thought"])
+        case "SetTodoList":
+            if not isinstance(curr_args, dict) or not curr_args.get("todos"):
+                return None
+            if not isinstance(curr_args["todos"], list):
+                return None
+            for todo in curr_args["todos"]:
+                if not isinstance(todo, dict) or not todo.get("title"):
+                    continue
+                subtitle += f"• {todo['title']}"
+                if todo.get("status"):
+                    subtitle += f" [{todo['status']}]"
+                subtitle += "\n"
+            return "\n" + subtitle.strip()
+        case "Bash":
+            if not isinstance(curr_args, dict) or not curr_args.get("command"):
+                return None
+            subtitle = str(curr_args["command"])
+        case "ReadFile":
+            if not isinstance(curr_args, dict) or not curr_args.get("path"):
+                return None
+            subtitle = _normalize_path(str(curr_args["path"]))
+        case "Glob":
+            if not isinstance(curr_args, dict) or not curr_args.get("pattern"):
+                return None
+            subtitle = str(curr_args["pattern"])
+        case "Grep":
+            if not isinstance(curr_args, dict) or not curr_args.get("pattern"):
+                return None
+            subtitle = str(curr_args["pattern"])
+        case "WriteFile":
+            if not isinstance(curr_args, dict) or not curr_args.get("path"):
+                return None
+            subtitle = _normalize_path(str(curr_args["path"]))
+        case "StrReplaceFile":
+            if not isinstance(curr_args, dict) or not curr_args.get("path"):
+                return None
+            subtitle = _normalize_path(str(curr_args["path"]))
+        case "SearchWeb":
+            if not isinstance(curr_args, dict) or not curr_args.get("query"):
+                return None
+            subtitle = str(curr_args["query"])
+        case "FetchURL":
+            if not isinstance(curr_args, dict) or not curr_args.get("url"):
+                return None
+            subtitle = str(curr_args["url"])
+        case _:
+            subtitle = "".join(lexer.json_content)
+    if tool_name not in ["SetTodoList"]:
+        subtitle = shorten_middle(subtitle, width=50)
+    return subtitle
+
+
+def _normalize_path(path: str) -> str:
+    cwd = str(Path.cwd().absolute())
+    if path.startswith(cwd):
+        path = path[len(cwd) :].lstrip("/\\")
+    return path

kimi_cli/tools/bash/__init__.py

@@ -0,0 +1,97 @@
+import asyncio
+from pathlib import Path
+from typing import override
+
+from kosong.tooling import CallableTool2, ToolReturnType
+from pydantic import BaseModel, Field
+
+from kimi_cli.soul.approval import Approval
+from kimi_cli.tools.utils import ToolRejectedError, ToolResultBuilder, load_desc
+
+MAX_TIMEOUT = 5 * 60
+
+
+class Params(BaseModel):
+    command: str = Field(description="The bash command to execute.")
+    timeout: int = Field(
+        description=(
+            "The timeout in seconds for the command to execute. "
+            "If the command takes longer than this, it will be killed."
+        ),
+        default=60,
+        ge=1,
+        le=MAX_TIMEOUT,
+    )
+
+
+class Bash(CallableTool2[Params]):
+    name: str = "Bash"
+    description: str = load_desc(Path(__file__).parent / "bash.md", {})
+    params: type[Params] = Params
+
+    def __init__(self, approval: Approval, **kwargs):
+        super().__init__(**kwargs)
+        self._approval = approval
+
+    @override
+    async def __call__(self, params: Params) -> ToolReturnType:
+        builder = ToolResultBuilder()
+
+        if not await self._approval.request(
+            f"run command {params.command}", f"Run command `{params.command}`"
+        ):
+            return ToolRejectedError()
+
+        def stdout_cb(line: bytes):
+            line_str = line.decode()
+            builder.write(line_str)
+
+        def stderr_cb(line: bytes):
+            line_str = line.decode()
+            builder.write(line_str)
+
+        try:
+            exitcode = await _stream_subprocess(
+                params.command, stdout_cb, stderr_cb, params.timeout
+            )
+
+            if exitcode == 0:
+                return builder.ok("Command executed successfully.")
+            else:
+                return builder.error(
+                    f"Command failed with exit code: {exitcode}.",
+                    brief=f"Failed with exit code: {exitcode}",
+                )
+        except TimeoutError:
+            return builder.error(
+                f"Command killed by timeout ({params.timeout}s)",
+                brief=f"Killed by timeout ({params.timeout}s)",
+            )
+
+
+async def _stream_subprocess(command: str, stdout_cb, stderr_cb, timeout: int) -> int:
+    async def _read_stream(stream, cb):
+        while True:
+            line = await stream.readline()
+            if line:
+                cb(line)
+            else:
+                break
+
+    # FIXME: if the event loop is cancelled, an exception may be raised when the process finishes
+    process = await asyncio.create_subprocess_shell(
+        command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+    )
+
+    try:
+        await asyncio.wait_for(
+            asyncio.gather(
+                _read_stream(process.stdout, stdout_cb),
+                _read_stream(process.stderr, stderr_cb),
+            ),
+            timeout,
+        )
+        return await process.wait()
+    except TimeoutError:
+        process.kill()
+        raise

kimi_cli/tools/bash/bash.md

@@ -0,0 +1,31 @@
+Execute a shell command. Use this tool to explore the filesystem, edit files, run scripts, get system information, etc.
+
+**Output:**
+The stdout and stderr will be combined and returned as a string. The output may be truncated if it is too long. If the command failed, the exit code will be provided in a system tag.
+
+**Guidelines for safety and security:**
+- Each shell tool call will be executed in a fresh shell environment. The shell variables, current working directory changes, and the shell history is not preserved between calls.
+- The tool call will return after the command is finished. You shall not use this tool to execute an interactive command or a command that may run forever. For possibly long-running commands, you shall set `timeout` argument to a reasonable value.
+- Avoid using `..` to access files or directories outside of the working directory.
+- Avoid modifying files outside of the working directory unless explicitly instructed to do so.
+- Never run commands that require superuser privileges unless explicitly instructed to do so.
+
+**Guidelines for efficiency:**
+- For multiple related commands, use `&&` to chain them in a single call, e.g. `cd /path && ls -la`
+- Use `;` to run commands sequentially regardless of success/failure
+- Use `||` for conditional execution (run second command only if first fails)
+- Use pipe operations (`|`) and redirections (`>`, `>>`) to chain input and output between commands
+- Always quote file paths containing spaces with double quotes (e.g., cd "/path with spaces/")
+- Use `if`, `case`, `for`, `while` control flows to execute complex logic in a single call.
+- Verify directory structure before create/edit/delete files or directories to reduce the risk of failure.
+
+**Commands available:**
+- Shell environment: cd, pwd, export, unset, env
+- File system operations: ls, find, mkdir, rm, cp, mv, touch, chmod, chown
+- File viewing/editing: cat, grep, head, tail, diff, patch
+- Text processing: awk, sed, sort, uniq, wc
+- System information/operations: ps, kill, top, df, free, uname, whoami, id, date
+- Package management: pip, uv, npm, yarn, bun, cargo
+- Network operations: curl, wget, ping, telnet, ssh
+- Archive operations: tar, zip, unzip
+- Other: Other commands available in the shell environment. Check the existence of a command by running `which <command>` before using it.

kimi_cli/tools/dmail/__init__.py

@@ -0,0 +1,38 @@
+from pathlib import Path
+from typing import override
+
+from kosong.tooling import CallableTool2, ToolError, ToolReturnType
+
+from kimi_cli.soul.denwarenji import DenwaRenji, DenwaRenjiError, DMail
+
+NAME = "SendDMail"
+
+
+class SendDMail(CallableTool2):
+    name: str = NAME
+    description: str = (Path(__file__).parent / "dmail.md").read_text()
+    params: type[DMail] = DMail
+
+    def __init__(self, denwa_renji: DenwaRenji, **kwargs):
+        super().__init__(**kwargs)
+        self._denwa_renji = denwa_renji
+
+    @override
+    async def __call__(self, params: DMail) -> ToolReturnType:
+        try:
+            self._denwa_renji.send_dmail(params)
+        except DenwaRenjiError as e:
+            return ToolError(
+                output="",
+                message=f"Failed to send D-Mail. Error: {str(e)}",
+                brief="Failed to send D-Mail",
+            )
+        # always return an error because a successful SendDMail call will never return
+        return ToolError(
+            output="",
+            message=(
+                "If you see this message, the D-Mail was not sent successfully. "
+                "This may be because some other tool that needs approval was rejected."
+            ),
+            brief="D-Mail not sent",
+        )

kimi_cli/tools/dmail/dmail.md

@@ -0,0 +1,15 @@
+Send a message to the past, just like sending a D-Mail in Steins;Gate.
+
+You can see some `user` messages with `CHECKPOINT {checkpoint_id}` wrapped in `<system>` tags in the context. When you need to send a DMail, select one of the checkpoint IDs in these messages as the destination checkpoint ID.
+
+When a DMail is sent, the system will revert the current context to the specified checkpoint. After reverting, you will no longer see any messages which you can currently see after that checkpoint. The message in the DMail will be appended to the end of the context. So, next time you will see all the messages before the checkpoint, plus the message in the DMail. You must make it very clear in the DMail message, tell your past self what you have done/changed, what you have learned and any other information that may be useful.
+
+When sending a DMail, DO NOT do much explanation to the user. The user do not care about this. Just explain to your past self.
+
+Here are some typical scenarios you may want to send a DMail:
+
+- You read a file, found it very large and most of the content is not relevant to the current task. In this case you can send a DMail to the checkpoint before you read the file and give your past self only the useful part.
+- You searched the web, found the result very large.
+  - If you got what you need, you may send a DMail to the checkpoint before you searched the web and give your past self the useful part.
+  - If you did not get what you need, you may send a DMail to tell your past self to try another query.
+- You wrote some code and it did not work as expected. You spent many struggling steps to fix it but the process is not relevant to the ultimate goal. In this case you can send a DMail to the checkpoint before you wrote the code and give your past self the fixed version of the code and tell yourself no need to write it again because you already wrote to the filesystem.

kimi_cli/tools/file/__init__.py

@@ -0,0 +1,21 @@
+class FileOpsWindow:
+    """Maintains a window of file operations."""
+
+    pass
+
+
+from .glob import Glob  # noqa: E402
+from .grep import Grep  # noqa: E402
+from .patch import PatchFile  # noqa: E402
+from .read import ReadFile  # noqa: E402
+from .replace import StrReplaceFile  # noqa: E402
+from .write import WriteFile  # noqa: E402
+
+__all__ = (
+    "ReadFile",
+    "Glob",
+    "Grep",
+    "WriteFile",
+    "StrReplaceFile",
+    "PatchFile",
+)

kimi_cli/tools/file/glob.md

@@ -0,0 +1,17 @@
+Find files and directories using glob patterns. This tool supports standard glob syntax like `*`, `?`, and `**` for recursive searches.
+
+**When to use:**
+- Find files matching specific patterns (e.g., all Python files: `*.py`)
+- Search for files recursively in subdirectories (e.g., `src/**/*.js`)
+- Locate configuration files (e.g., `*.config.*`, `*.json`)
+- Find test files (e.g., `test_*.py`, `*_test.go`)
+
+**Example patterns:**
+- `*.py` - All Python files in current directory
+- `src/**/*.js` - All JavaScript files in src directory recursively
+- `test_*.py` - Python test files starting with "test_"
+- `*.config.{js,ts}` - Config files with .js or .ts extension
+
+**Bad example patterns:**
+- `**`, `**/*.py` - Any pattern starting with '**' will be rejected. Because it would recursively search all directories and subdirectories, which is very likely to yield large result that exceeds your context size. Always use more specific patterns like `src/**/*.py` instead.
+- `node_modules/**/*.js` - Although this does not start with '**', it would still highly possible to yield large result because `node_modules` is well-known to contain too many directories and files. Avoid recursivelly searching in such directories, other examples include `venv`, `.venv`, `__pycache__`, `target`. If you really need to search in a dependency, use more specific patterns like `node_modules/react/src/*` instead.

kimi_cli/tools/file/glob.py

@@ -0,0 +1,149 @@
+"""Glob tool implementation."""
+
+import asyncio
+from pathlib import Path
+from typing import override
+
+import aiofiles.os
+from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
+from pydantic import BaseModel, Field
+
+from kimi_cli.agent import BuiltinSystemPromptArgs
+from kimi_cli.tools.utils import load_desc
+
+MAX_MATCHES = 1000
+
+
+class Params(BaseModel):
+    pattern: str = Field(description=("Glob pattern to match files/directories."))
+    directory: str | None = Field(
+        description=(
+            "Absolute path to the directory to search in (defaults to working directory)."
+        ),
+        default=None,
+    )
+    include_dirs: bool = Field(
+        description="Whether to include directories in results.",
+        default=True,
+    )
+
+
+class Glob(CallableTool2[Params]):
+    name: str = "Glob"
+    description: str = load_desc(
+        Path(__file__).parent / "glob.md",
+        {
+            "MAX_MATCHES": str(MAX_MATCHES),
+        },
+    )
+    params: type[Params] = Params
+
+    def __init__(self, builtin_args: BuiltinSystemPromptArgs, **kwargs):
+        super().__init__(**kwargs)
+        self._work_dir = builtin_args.KIMI_WORK_DIR
+
+    async def _validate_pattern(self, pattern: str) -> ToolError | None:
+        """Validate that the pattern is safe to use."""
+        if pattern.startswith("**"):
+            # TODO: give a `ls -la` result as the output
+            ls_result = await aiofiles.os.listdir(self._work_dir)
+            return ToolError(
+                output="\n".join(ls_result),
+                message=(
+                    f"Pattern `{pattern}` starts with '**' which is not allowed. "
+                    "This would recursively search all directories and may include large "
+                    "directories like `node_modules`. Use more specific patterns instead. "
+                    "For your convenience, a list of all files and directories in the "
+                    "top level of the working directory is provided below."
+                ),
+                brief="Unsafe pattern",
+            )
+        return None
+
+    def _validate_directory(self, directory: Path) -> ToolError | None:
+        """Validate that the directory is safe to search."""
+        resolved_dir = directory.resolve()
+        resolved_work_dir = self._work_dir.resolve()
+
+        # Ensure the directory is within work directory
+        if not str(resolved_dir).startswith(str(resolved_work_dir)):
+            return ToolError(
+                message=(
+                    f"`{directory}` is outside the working directory. "
+                    "You can only search within the working directory."
+                ),
+                brief="Directory outside working directory",
+            )
+        return None
+
+    @override
+    async def __call__(self, params: Params) -> ToolReturnType:
+        try:
+            # Validate pattern safety
+            pattern_error = await self._validate_pattern(params.pattern)
+            if pattern_error:
+                return pattern_error
+
+            dir_path = Path(params.directory) if params.directory else self._work_dir
+
+            if not dir_path.is_absolute():
+                return ToolError(
+                    message=(
+                        f"`{params.directory}` is not an absolute path. "
+                        "You must provide an absolute path to search."
+                    ),
+                    brief="Invalid directory",
+                )
+
+            # Validate directory safety
+            dir_error = self._validate_directory(dir_path)
+            if dir_error:
+                return dir_error
+
+            if not dir_path.exists():
+                return ToolError(
+                    message=f"`{params.directory}` does not exist.",
+                    brief="Directory not found",
+                )
+            if not dir_path.is_dir():
+                return ToolError(
+                    message=f"`{params.directory}` is not a directory.",
+                    brief="Invalid directory",
+                )
+
+            def _glob(pattern: str) -> list[Path]:
+                return list(dir_path.glob(pattern))
+
+            # Perform the glob search - users can use ** directly in pattern
+            matches = await asyncio.to_thread(_glob, params.pattern)
+
+            # Filter out directories if not requested
+            if not params.include_dirs:
+                matches = [p for p in matches if p.is_file()]
+
+            # Sort for consistent output
+            matches.sort()
+
+            # Limit matches
+            message = (
+                f"Found {len(matches)} matches for pattern `{params.pattern}`."
+                if len(matches) > 0
+                else "No matches found for pattern `{params.pattern}`."
+            )
+            if len(matches) > MAX_MATCHES:
+                matches = matches[:MAX_MATCHES]
+                message += (
+                    f" Only the first {MAX_MATCHES} matches are returned. "
+                    "You may want to use a more specific pattern."
+                )
+
+            return ToolOk(
+                output="\n".join(str(p.relative_to(dir_path)) for p in matches),
+                message=message,
+            )
+
+        except Exception as e:
+            return ToolError(
+                message=f"Failed to search for pattern {params.pattern}. Error: {e}",
+                brief="Glob failed",
+            )

kimi_cli/tools/file/grep.md

@@ -0,0 +1,5 @@
+A powerful search tool based-on ripgrep.
+
+**Tips:**
+- ALWAYS use Grep tool instead of running `grep` or `rg` command with Bash tool.
+- Use the ripgrep pattern syntax, not grep syntax. E.g. you need to escape braces like `\\{` to search for `{`.