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

kimi_cli/soul/message.py

@@ -1,69 +1,56 @@
-from kosong.base.message import ContentPart, Message, TextPart
-from kosong.tooling import ToolError, ToolOk, ToolResult
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from kosong.message import Message
 from kosong.tooling.error import ToolRuntimeError
 
+from kimi_cli.llm import ModelCapability
+from kimi_cli.wire.types import (
+    ContentPart,
+    ImageURLPart,
+    TextPart,
+    ThinkPart,
+    ToolResult,
+    VideoURLPart,
+)
+
 
 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):
+def tool_result_to_message(tool_result: ToolResult) -> Message:
+    """Convert a tool result to a message."""
+    if tool_result.return_value.is_error:
+        assert tool_result.return_value.message, "Error return value should have a message"
+        message = tool_result.return_value.message
+        if isinstance(tool_result.return_value, ToolRuntimeError):
             message += "\nThis is an unexpected error and the tool is probably not working."
         content: list[ContentPart] = [system(f"ERROR: {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: list[ContentPart] = []
-    non_text_parts: list[ContentPart] = []
-    for part in content:
-        if isinstance(part, TextPart):
-            text_parts.append(part)
-        else:
-            non_text_parts.append(part)
+        if tool_result.return_value.output:
+            content.extend(_output_to_content_parts(tool_result.return_value.output))
+    else:
+        content: list[ContentPart] = []
+        if tool_result.return_value.message:
+            content.append(system(tool_result.return_value.message))
+        if tool_result.return_value.output:
+            content.extend(_output_to_content_parts(tool_result.return_value.output))
+        if not content:
+            content.append(system("Tool output is empty."))
 
-    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=content,
+        tool_call_id=tool_result.tool_call_id,
     )
-    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."""
+def _output_to_content_parts(
+    output: str | ContentPart | Sequence[ContentPart],
+) -> list[ContentPart]:
     content: list[ContentPart] = []
-    if result.message:
-        content.append(system(result.message))
-    match output := result.output:
+    match output:
         case str(text):
             if text:
                 content.append(TextPart(text=text))
@@ -71,6 +58,19 @@ def tool_ok_to_message_content(result: ToolOk) -> list[ContentPart]:
             content.append(output)
         case _:
             content.extend(output)
-    if not content:
-        content.append(system("Tool output is empty."))
     return content
+
+
+def check_message(
+    message: Message, model_capabilities: set[ModelCapability]
+) -> set[ModelCapability]:
+    """Check the message content, return the missing model capabilities."""
+    capabilities_needed = set[ModelCapability]()
+    for part in message.content:
+        if isinstance(part, ImageURLPart):
+            capabilities_needed.add("image_in")
+        elif isinstance(part, VideoURLPart):
+            capabilities_needed.add("video_in")
+        elif isinstance(part, ThinkPart):
+            capabilities_needed.add("thinking")
+    return capabilities_needed - model_capabilities

kimi_cli/soul/slash.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+import tempfile
+from collections.abc import Awaitable, Callable
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from kosong.message import Message
+from loguru import logger
+
+import kimi_cli.prompts as prompts
+from kimi_cli.soul import wire_send
+from kimi_cli.soul.agent import load_agents_md
+from kimi_cli.soul.context import Context
+from kimi_cli.soul.message import system
+from kimi_cli.utils.slashcmd import SlashCommandRegistry
+from kimi_cli.wire.types import TextPart
+
+if TYPE_CHECKING:
+    from kimi_cli.soul.kimisoul import KimiSoul
+
+type SoulSlashCmdFunc = Callable[[KimiSoul, str], None | Awaitable[None]]
+"""
+A function that runs as a KimiSoul-level slash command.
+
+Raises:
+    Any exception that can be raised by `Soul.run`.
+"""
+
+registry = SlashCommandRegistry[SoulSlashCmdFunc]()
+
+
+@registry.command
+async def init(soul: KimiSoul, args: str):
+    """Analyze the codebase and generate an `AGENTS.md` file"""
+    from kimi_cli.soul.kimisoul import KimiSoul
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        tmp_context = Context(file_backend=Path(temp_dir) / "context.jsonl")
+        tmp_soul = KimiSoul(soul.agent, context=tmp_context)
+        await tmp_soul.run(prompts.INIT)
+
+    agents_md = load_agents_md(soul.runtime.builtin_args.KIMI_WORK_DIR)
+    system_message = system(
+        "The user just ran `/init` slash command. "
+        "The system has analyzed the codebase and generated an `AGENTS.md` file. "
+        f"Latest AGENTS.md file content:\n{agents_md}"
+    )
+    await soul.context.append_message(Message(role="user", content=[system_message]))
+
+
+@registry.command
+async def compact(soul: KimiSoul, args: str):
+    """Compact the context"""
+    if soul.context.n_checkpoints == 0:
+        wire_send(TextPart(text="The context is empty."))
+        return
+
+    logger.info("Running `/compact`")
+    await soul.compact_context()
+    wire_send(TextPart(text="The context has been compacted."))
+
+
+@registry.command
+async def yolo(soul: KimiSoul, args: str):
+    """Toggle YOLO mode (auto-approve all actions)"""
+    if soul.runtime.approval.is_yolo():
+        soul.runtime.approval.set_yolo(False)
+        wire_send(TextPart(text="You only die once! Actions will require approval."))
+    else:
+        soul.runtime.approval.set_yolo(True)
+        wire_send(TextPart(text="You only live once! All actions will be auto-approved."))

kimi_cli/soul/toolset.py

@@ -1,8 +1,46 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import importlib
+import inspect
+import json
 from contextvars import ContextVar
-from typing import override
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import TYPE_CHECKING, Any, Literal, overload
+
+from kosong.tooling import (
+    CallableTool,
+    CallableTool2,
+    HandleResult,
+    Tool,
+    ToolError,
+    ToolOk,
+    Toolset,
+)
+from kosong.tooling.error import (
+    ToolNotFoundError,
+    ToolParseError,
+    ToolRuntimeError,
+)
+from kosong.tooling.mcp import convert_mcp_content
+from kosong.utils.typing import JsonType
+from loguru import logger
+
+from kimi_cli.exception import InvalidToolError, MCPRuntimeError
+from kimi_cli.tools import SkipThisTool
+from kimi_cli.tools.utils import ToolRejectedError
+from kimi_cli.wire.types import ContentPart, ToolCall, ToolResult, ToolReturnValue
+
+if TYPE_CHECKING:
+    import fastmcp
+    import mcp
+    from fastmcp.client.client import CallToolResult
+    from fastmcp.client.transports import ClientTransport
+    from fastmcp.mcp_config import MCPConfig
 
-from kosong.base.message import ToolCall
-from kosong.tooling import HandleResult, SimpleToolset
+    from kimi_cli.soul.agent import Runtime
 
 current_tool_call = ContextVar[ToolCall | None]("current_tool_call", default=None)
 
@@ -15,11 +53,354 @@ def get_current_tool_call_or_none() -> ToolCall | None:
     return current_tool_call.get()
 
 
-class CustomToolset(SimpleToolset):
-    @override
+type ToolType = CallableTool | CallableTool2[Any]
+
+
+if TYPE_CHECKING:
+
+    def type_check(kimi_toolset: KimiToolset):
+        _: Toolset = kimi_toolset
+
+
+class KimiToolset:
+    def __init__(self) -> None:
+        self._tool_dict: dict[str, ToolType] = {}
+        self._mcp_servers: dict[str, MCPServerInfo] = {}
+        self._mcp_loading_task: asyncio.Task[None] | None = None
+
+    def add(self, tool: ToolType) -> None:
+        self._tool_dict[tool.name] = tool
+
+    @overload
+    def find(self, tool_name_or_type: str) -> ToolType | None: ...
+    @overload
+    def find[T: ToolType](self, tool_name_or_type: type[T]) -> T | None: ...
+    def find(self, tool_name_or_type: str | type[ToolType]) -> ToolType | None:
+        if isinstance(tool_name_or_type, str):
+            return self._tool_dict.get(tool_name_or_type)
+        else:
+            for tool in self._tool_dict.values():
+                if isinstance(tool, tool_name_or_type):
+                    return tool
+        return None
+
+    @property
+    def tools(self) -> list[Tool]:
+        return [tool.base for tool in self._tool_dict.values()]
+
     def handle(self, tool_call: ToolCall) -> HandleResult:
         token = current_tool_call.set(tool_call)
         try:
-            return super().handle(tool_call)
+            if tool_call.function.name not in self._tool_dict:
+                return ToolResult(
+                    tool_call_id=tool_call.id,
+                    return_value=ToolNotFoundError(tool_call.function.name),
+                )
+
+            tool = self._tool_dict[tool_call.function.name]
+
+            try:
+                arguments: JsonType = json.loads(tool_call.function.arguments or "{}")
+            except json.JSONDecodeError as e:
+                return ToolResult(tool_call_id=tool_call.id, return_value=ToolParseError(str(e)))
+
+            async def _call():
+                try:
+                    ret = await tool.call(arguments)
+                    return ToolResult(tool_call_id=tool_call.id, return_value=ret)
+                except Exception as e:
+                    return ToolResult(
+                        tool_call_id=tool_call.id, return_value=ToolRuntimeError(str(e))
+                    )
+
+            return asyncio.create_task(_call())
         finally:
             current_tool_call.reset(token)
+
+    @property
+    def mcp_servers(self) -> dict[str, MCPServerInfo]:
+        """Get MCP servers info."""
+        return self._mcp_servers
+
+    def load_tools(self, tool_paths: list[str], dependencies: dict[type[Any], Any]) -> None:
+        """
+        Load tools from paths like `kimi_cli.tools.shell:Shell`.
+
+        Raises:
+            InvalidToolError(KimiCLIException, ValueError): When any tool cannot be loaded.
+        """
+
+        good_tools: list[str] = []
+        bad_tools: list[str] = []
+
+        for tool_path in tool_paths:
+            try:
+                tool = self._load_tool(tool_path, dependencies)
+            except SkipThisTool:
+                logger.info("Skipping tool: {tool_path}", tool_path=tool_path)
+                continue
+            if tool:
+                self.add(tool)
+                good_tools.append(tool_path)
+            else:
+                bad_tools.append(tool_path)
+        logger.info("Loaded tools: {good_tools}", good_tools=good_tools)
+        if bad_tools:
+            raise InvalidToolError(f"Invalid tools: {bad_tools}")
+
+    @staticmethod
+    def _load_tool(tool_path: str, dependencies: dict[type[Any], Any]) -> ToolType | None:
+        logger.debug("Loading tool: {tool_path}", tool_path=tool_path)
+        module_name, class_name = tool_path.rsplit(":", 1)
+        try:
+            module = importlib.import_module(module_name)
+        except ImportError:
+            return None
+        tool_cls = getattr(module, class_name, None)
+        if tool_cls is None:
+            return None
+        args: list[Any] = []
+        if "__init__" in tool_cls.__dict__:
+            # the tool class overrides the `__init__` of base class
+            for param in inspect.signature(tool_cls).parameters.values():
+                if param.kind == inspect.Parameter.KEYWORD_ONLY:
+                    # once we encounter a keyword-only parameter, we stop injecting dependencies
+                    break
+                # all positional parameters should be dependencies to be injected
+                if param.annotation not in dependencies:
+                    raise ValueError(f"Tool dependency not found: {param.annotation}")
+                args.append(dependencies[param.annotation])
+        return tool_cls(*args)
+
+    # TODO(rc): remove `in_background` parameter and always load in background
+    async def load_mcp_tools(
+        self, mcp_configs: list[MCPConfig], runtime: Runtime, in_background: bool = True
+    ) -> None:
+        """
+        Load MCP tools from specified MCP configs.
+
+        Raises:
+            MCPRuntimeError(KimiCLIException, RuntimeError): When any MCP server cannot be
+                connected.
+        """
+        import fastmcp
+        from fastmcp.mcp_config import MCPConfig, RemoteMCPServer
+
+        from kimi_cli.ui.shell.prompt import toast
+
+        async def _check_oauth_tokens(server_url: str) -> bool:
+            """Check if OAuth tokens exist for the server."""
+            try:
+                from fastmcp.client.auth.oauth import FileTokenStorage
+
+                storage = FileTokenStorage(server_url=server_url)
+                tokens = await storage.get_tokens()
+                return tokens is not None
+            except Exception:
+                return False
+
+        def _toast_mcp(message: str) -> None:
+            if in_background:
+                toast(
+                    message,
+                    duration=10.0,
+                    topic="mcp",
+                    immediate=True,
+                    position="right",
+                )
+
+        oauth_servers: dict[str, str] = {}
+
+        async def _connect_server(
+            server_name: str, server_info: MCPServerInfo
+        ) -> tuple[str, Exception | None]:
+            if server_info.status != "pending":
+                return server_name, None
+
+            server_info.status = "connecting"
+            try:
+                async with server_info.client as client:
+                    for tool in await client.list_tools():
+                        server_info.tools.append(
+                            MCPTool(server_name, tool, client, runtime=runtime)
+                        )
+
+                for tool in server_info.tools:
+                    self.add(tool)
+
+                server_info.status = "connected"
+                logger.info("Connected MCP server: {server_name}", server_name=server_name)
+                return server_name, None
+            except Exception as e:
+                logger.error(
+                    "Failed to connect MCP server: {server_name}, error: {error}",
+                    server_name=server_name,
+                    error=e,
+                )
+                server_info.status = "failed"
+                return server_name, e
+
+        async def _connect():
+            _toast_mcp("connecting to mcp servers...")
+            unauthorized_servers: dict[str, str] = {}
+            for server_name, server_info in self._mcp_servers.items():
+                server_url = oauth_servers.get(server_name)
+                if not server_url:
+                    continue
+                if not await _check_oauth_tokens(server_url):
+                    logger.warning(
+                        "Skipping OAuth MCP server '{server_name}': not authorized. "
+                        "Run 'kimi mcp auth {server_name}' first.",
+                        server_name=server_name,
+                    )
+                    server_info.status = "unauthorized"
+                    unauthorized_servers[server_name] = server_url
+
+            tasks = [
+                asyncio.create_task(_connect_server(server_name, server_info))
+                for server_name, server_info in self._mcp_servers.items()
+                if server_info.status == "pending"
+            ]
+            results = await asyncio.gather(*tasks) if tasks else []
+            failed_servers = {name: error for name, error in results if error is not None}
+
+            for mcp_config in mcp_configs:
+                # Skip empty MCP configs (no servers defined)
+                if not mcp_config.mcpServers:
+                    logger.debug("Skipping empty MCP config: {mcp_config}", mcp_config=mcp_config)
+                    continue
+
+            if failed_servers:
+                _toast_mcp("mcp connection failed")
+                raise MCPRuntimeError(f"Failed to connect MCP servers: {failed_servers}")
+            if unauthorized_servers:
+                _toast_mcp("mcp authorization needed")
+            else:
+                _toast_mcp("mcp servers connected")
+
+        for mcp_config in mcp_configs:
+            if not mcp_config.mcpServers:
+                logger.debug("Skipping empty MCP config: {mcp_config}", mcp_config=mcp_config)
+                continue
+
+            for server_name, server_config in mcp_config.mcpServers.items():
+                if isinstance(server_config, RemoteMCPServer) and server_config.auth == "oauth":
+                    oauth_servers[server_name] = server_config.url
+
+                # Add mcp-session-id header for HTTP transports (skip OAuth servers)
+                if (
+                    isinstance(server_config, RemoteMCPServer)
+                    and server_config.auth != "oauth"
+                    and not any(key.lower() == "mcp-session-id" for key in server_config.headers)
+                ):
+                    server_config = server_config.model_copy(deep=True)
+                    server_config.headers["Mcp-Session-Id"] = runtime.session.id
+
+                client = fastmcp.Client(MCPConfig(mcpServers={server_name: server_config}))
+                self._mcp_servers[server_name] = MCPServerInfo(
+                    status="pending", client=client, tools=[]
+                )
+
+        if in_background:
+            self._mcp_loading_task = asyncio.create_task(_connect())
+        else:
+            await _connect()
+
+    async def wait_for_mcp_tools(self) -> None:
+        """Wait for background MCP tool loading to finish."""
+        task = self._mcp_loading_task
+        if not task:
+            return
+        try:
+            await task
+        finally:
+            if self._mcp_loading_task is task and task.done():
+                self._mcp_loading_task = None
+
+    async def cleanup(self) -> None:
+        """Cleanup any resources held by the toolset."""
+        if self._mcp_loading_task:
+            self._mcp_loading_task.cancel()
+            with contextlib.suppress(Exception):
+                await self._mcp_loading_task
+        for server_info in self._mcp_servers.values():
+            await server_info.client.close()
+
+
+@dataclass(slots=True)
+class MCPServerInfo:
+    status: Literal["pending", "connecting", "connected", "failed", "unauthorized"]
+    client: fastmcp.Client[Any]
+    tools: list[MCPTool[Any]]
+
+
+class MCPTool[T: ClientTransport](CallableTool):
+    def __init__(
+        self,
+        server_name: str,
+        mcp_tool: mcp.Tool,
+        client: fastmcp.Client[T],
+        *,
+        runtime: Runtime,
+        **kwargs: Any,
+    ):
+        super().__init__(
+            name=mcp_tool.name,
+            description=(
+                f"This is an MCP (Model Context Protocol) tool from MCP server `{server_name}`.\n\n"
+                f"{mcp_tool.description or 'No description provided.'}"
+            ),
+            parameters=mcp_tool.inputSchema,
+            **kwargs,
+        )
+        self._mcp_tool = mcp_tool
+        self._client = client
+        self._runtime = runtime
+        self._timeout = timedelta(milliseconds=runtime.config.mcp.client.tool_call_timeout_ms)
+        self._action_name = f"mcp:{mcp_tool.name}"
+
+    async def __call__(self, *args: Any, **kwargs: Any) -> ToolReturnValue:
+        description = f"Call MCP tool `{self._mcp_tool.name}`."
+        if not await self._runtime.approval.request(self.name, self._action_name, description):
+            return ToolRejectedError()
+
+        try:
+            async with self._client as client:
+                result = await client.call_tool(
+                    self._mcp_tool.name,
+                    kwargs,
+                    timeout=self._timeout,
+                    raise_on_error=False,
+                )
+                return convert_mcp_tool_result(result)
+        except Exception as e:
+            # fastmcp raises `RuntimeError` on timeout and we cannot tell it from other errors
+            exc_msg = str(e).lower()
+            if "timeout" in exc_msg or "timed out" in exc_msg:
+                return ToolError(
+                    message=(
+                        f"Timeout while calling MCP tool `{self._mcp_tool.name}`. "
+                        "You may explain to the user that the timeout config is set too low."
+                    ),
+                    brief="Timeout",
+                )
+            raise
+
+
+def convert_mcp_tool_result(result: CallToolResult) -> ToolReturnValue:
+    """Convert MCP tool result to kosong tool return value.
+
+    Raises:
+        ValueError: If any content part has unsupported type or mime type.
+    """
+    content: list[ContentPart] = []
+    for part in result.content:
+        content.append(convert_mcp_content(part))
+    if result.is_error:
+        return ToolError(
+            output=content,
+            message="Tool returned an error. The output may be error message or incomplete output",
+            brief="",
+        )
+    else:
+        return ToolOk(output=content)

kimi_cli/toad.py

@@ -0,0 +1,74 @@
+import importlib.util
+import shlex
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+import typer
+
+
+def _default_acp_command() -> list[str]:
+    argv0 = sys.argv[0]
+    if argv0:
+        resolved = shutil.which(argv0)
+        resolved_path = Path(resolved).expanduser() if resolved else Path(argv0).expanduser()
+        if (
+            resolved_path.exists()
+            and resolved_path.suffix != ".py"
+            and not resolved_path.name.startswith(("python", "pypy"))
+        ):
+            return [str(resolved_path), "acp"]
+
+    return [sys.executable, "-m", "kimi_cli.cli", "acp"]
+
+
+def _default_toad_command() -> list[str]:
+    if sys.version_info < (3, 14):
+        typer.echo("`kimi term` requires Python 3.14+ because Toad requires it.", err=True)
+        raise typer.Exit(code=1)
+    if importlib.util.find_spec("toad") is None:
+        typer.echo(
+            "Toad dependency is missing. Run `uv sync --python 3.14` or install kimi-cli with "
+            "Python 3.14.",
+            err=True,
+        )
+        raise typer.Exit(code=1)
+    return [sys.executable, "-m", "toad.cli"]
+
+
+def _extract_project_dir(extra_args: list[str]) -> Path | None:
+    work_dir: str | None = None
+    idx = 0
+    while idx < len(extra_args):
+        arg = extra_args[idx]
+        if arg in ("--work-dir", "-w"):
+            if idx + 1 < len(extra_args):
+                work_dir = extra_args[idx + 1]
+                idx += 2
+                continue
+        elif arg.startswith("--work-dir=") or arg.startswith("-w="):
+            work_dir = arg.split("=", 1)[1]
+        elif arg.startswith("-w") and len(arg) > 2:
+            work_dir = arg[2:]
+        idx += 1
+
+    if not work_dir:
+        return None
+
+    return Path(work_dir).expanduser().resolve()
+
+
+def run_term(ctx: typer.Context) -> None:
+    extra_args = list(ctx.args)
+    acp_args = _default_acp_command()
+    acp_command = shlex.join(acp_args)
+    toad_parts = _default_toad_command()
+    args = [*toad_parts, "acp", acp_command]
+    project_dir = _extract_project_dir(extra_args)
+    if project_dir is not None:
+        args.append(str(project_dir))
+
+    result = subprocess.run(args)
+    if result.returncode != 0:
+        raise typer.Exit(code=result.returncode)

kimi_cli/tools/AGENTS.md

@@ -0,0 +1,5 @@
+# Kimi CLI Tools
+
+## Guidelines
+
+- Except for `Task` tool, tools should not refer to any types in `kimi_cli/wire/`. When importing things like `ToolReturnValue`, `DisplayBlock`, import from `kosong.tooling`.